import json import os from typing import Any import numpy as np import pytest import supervision as sv from tests.helpers import _create_detections @pytest.mark.parametrize( ( "detections", "custom_data", "second_detections", "second_custom_data", "file_name", "expected_result", ), [ ( _create_detections( xyxy=[[10, 20, 30, 40], [50, 60, 70, 80]], confidence=[0.7, 0.8], class_id=[0, 0], tracker_id=[0, 1], data={"class_name": ["person", "person"]}, ), {"frame_number": 42}, _create_detections( xyxy=[[15, 25, 35, 45], [55, 65, 75, 85]], confidence=[0.6, 0.9], class_id=[1, 1], tracker_id=[2, 3], data={"class_name": ["car", "car"]}, ), {"frame_number": 43}, "test_detections.json", [ { "x_min": 10, "y_min": 20, "x_max": 30, "y_max": 40, "class_id": 0, "confidence": 0.699999988079071, "tracker_id": 0, "class_name": "person", "frame_number": 42, }, { "x_min": 50, "y_min": 60, "x_max": 70, "y_max": 80, "class_id": 0, "confidence": 0.800000011920929, "tracker_id": 1, "class_name": "person", "frame_number": 42, }, { "x_min": 15, "y_min": 25, "x_max": 35, "y_max": 45, "class_id": 1, "confidence": 0.6000000238418579, "tracker_id": 2, "class_name": "car", "frame_number": 43, }, { "x_min": 55, "y_min": 65, "x_max": 75, "y_max": 85, "class_id": 1, "confidence": 0.8999999761581421, "tracker_id": 3, "class_name": "car", "frame_number": 43, }, ], ), # Multiple detections ( _create_detections( xyxy=[[60, 70, 80, 90], [100, 110, 120, 130]], tracker_id=[4, 5], data={"class_name": ["bike", "dog"]}, ), {"frame_number": 44}, _create_detections( xyxy=[[65, 75, 85, 95], [105, 115, 125, 135]], confidence=[0.5, 0.4], data={"class_name": ["tree", "cat"]}, ), {"frame_number": 45}, "test_detections_missing_fields.json", [ { "x_min": 60, "y_min": 70, "x_max": 80, "y_max": 90, "class_id": "", "confidence": "", "tracker_id": 4, "class_name": "bike", "frame_number": 44, }, { "x_min": 100, "y_min": 110, "x_max": 120, "y_max": 130, "class_id": "", "confidence": "", "tracker_id": 5, "class_name": "dog", "frame_number": 44, }, { "x_min": 65, "y_min": 75, "x_max": 85, "y_max": 95, "class_id": "", "confidence": 0.5, "tracker_id": "", "class_name": "tree", "frame_number": 45, }, { "x_min": 105, "y_min": 115, "x_max": 125, "y_max": 135, "class_id": "", "confidence": 0.4000000059604645, "tracker_id": "", "class_name": "cat", "frame_number": 45, }, ], ), # Missing fields ( _create_detections( xyxy=[[10, 11, 12, 13]], confidence=[0.95], data={"class_name": "unknown", "is_detected": True, "score": 1}, ), {"frame_number": 46}, _create_detections( xyxy=[[14, 15, 16, 17]], data={"class_name": "artifact", "is_detected": False, "score": 0.85}, ), {"frame_number": 47}, "test_detections_varied_data.json", [ { "x_min": 10, "y_min": 11, "x_max": 12, "y_max": 13, "class_id": "", "confidence": 0.949999988079071, "tracker_id": "", "class_name": "unknown", "is_detected": True, "score": 1, "frame_number": 46, }, { "x_min": 14, "y_min": 15, "x_max": 16, "y_max": 17, "class_id": "", "confidence": "", "tracker_id": "", "class_name": "artifact", "is_detected": False, "score": 0.85, "frame_number": 47, }, ], ), # Inconsistent Data Types ( _create_detections( xyxy=[[20, 21, 22, 23]], ), { "metadata": {"sensor_id": 101, "location": "north"}, "tags": ["urgent", "review"], }, _create_detections( xyxy=[[14, 15, 16, 17]], ), { "metadata": {"sensor_id": 104, "location": "west"}, "tags": ["not-urgent", "done"], }, "test_detections_complex_data.json", [ { "x_min": 20, "y_min": 21, "x_max": 22, "y_max": 23, "class_id": "", "confidence": "", "tracker_id": "", "metadata": {"sensor_id": 101, "location": "north"}, "tags": ["urgent", "review"], }, { "x_min": 14, "y_min": 15, "x_max": 16, "y_max": 17, "class_id": "", "confidence": "", "tracker_id": "", "metadata": {"sensor_id": 104, "location": "west"}, "tags": ["not-urgent", "done"], }, ], ), # Complex Data ( _create_detections( xyxy=[[10, 20, 30, 40], [50, 60, 70, 80]], confidence=[0.9, 0.8], class_id=[0, 1], ), {"area": np.array([400.0, 400.0])}, _create_detections( xyxy=[[15, 25, 35, 45]], confidence=[0.7], class_id=[2], ), {"area": np.array([400.0])}, "test_detections_array_custom_data.json", [ { "x_min": 10, "y_min": 20, "x_max": 30, "y_max": 40, "class_id": 0, "confidence": 0.8999999761581421, "tracker_id": "", "area": 400.0, }, { "x_min": 50, "y_min": 60, "x_max": 70, "y_max": 80, "class_id": 1, "confidence": 0.800000011920929, "tracker_id": "", "area": 400.0, }, { "x_min": 15, "y_min": 25, "x_max": 35, "y_max": 45, "class_id": 2, "confidence": 0.699999988079071, "tracker_id": "", "area": 400.0, }, ], ), # numpy array in custom_data sliced per detection row ( _create_detections( xyxy=[[10, 20, 30, 40], [50, 60, 70, 80]], confidence=[0.9, 0.8], class_id=[0, 1], ), {"ids": ["a", "b"], "tags": ("x", "y")}, _create_detections( xyxy=[[15, 25, 35, 45]], confidence=[0.7], class_id=[2], ), {"ids": ["c"], "tags": ("z",)}, "test_detections_list_custom_data.json", [ { "x_min": 10, "y_min": 20, "x_max": 30, "y_max": 40, "class_id": 0, "confidence": 0.8999999761581421, "tracker_id": "", "ids": "a", "tags": "x", }, { "x_min": 50, "y_min": 60, "x_max": 70, "y_max": 80, "class_id": 1, "confidence": 0.800000011920929, "tracker_id": "", "ids": "b", "tags": "y", }, { "x_min": 15, "y_min": 25, "x_max": 35, "y_max": 45, "class_id": 2, "confidence": 0.699999988079071, "tracker_id": "", "ids": "c", "tags": "z", }, ], ), # list/tuple custom_data matching detection count is sliced per row ( sv.Detections( xyxy=np.array([[10, 20, 30, 40], [50, 60, 70, 80]]), data={"labels": ["person", "car"]}, ), None, sv.Detections( xyxy=np.array([[15, 25, 35, 45]]), data={"labels": ["bus"]}, ), None, "test_detections_plain_list_data.json", [ { "x_min": 10.0, "y_min": 20.0, "x_max": 30.0, "y_max": 40.0, "class_id": "", "confidence": "", "tracker_id": "", "labels": "person", }, { "x_min": 50.0, "y_min": 60.0, "x_max": 70.0, "y_max": 80.0, "class_id": "", "confidence": "", "tracker_id": "", "labels": "car", }, { "x_min": 15.0, "y_min": 25.0, "x_max": 35.0, "y_max": 45.0, "class_id": "", "confidence": "", "tracker_id": "", "labels": "bus", }, ], ), # plain Python list in detections.data is sliced per row without custom_data ], ) def test_json_sink( detections: sv.Detections, custom_data: dict[str, Any] | None, second_detections: sv.Detections, second_custom_data: dict[str, Any] | None, file_name: str, expected_result: list[list[Any]], ) -> None: with sv.JSONSink(file_name) as sink: if custom_data is None: sink.append(detections) else: sink.append(detections, custom_data) if second_custom_data is None: sink.append(second_detections) else: sink.append(second_detections, second_custom_data) assert_json_equal(file_name, expected_result) @pytest.mark.parametrize( ("scalar", "expected"), [ pytest.param(np.int64(7), 7, id="np_int64"), pytest.param(np.float32(0.5), pytest.approx(0.5), id="np_float32"), pytest.param(np.float64(1.5), pytest.approx(1.5), id="np_float64"), pytest.param(np.int32(3), 3, id="np_int32"), pytest.param(np.bool_(True), True, id="np_bool"), ], ) def test_json_sink_serializes_numpy_scalar_custom_data( tmp_path: Any, scalar: Any, expected: Any ) -> None: """NumPy scalar in custom_data serializes as a JSON number or boolean.""" file_name = str(tmp_path / "test_numpy_scalar.json") detections = sv.Detections( xyxy=np.array([[0, 0, 10, 10]], dtype=np.float32), class_id=np.array([0]), confidence=np.array([0.9]), ) with sv.JSONSink(file_name) as sink: sink.append(detections, custom_data={"value": scalar}) with open(file_name) as f: data = json.load(f) assert data[0]["value"] == expected def test_json_sink_serializes_nested_numpy_array_custom_data(tmp_path: Any) -> None: """NumPy array nested inside a custom_data dict value serializes as a JSON array.""" file_name = str(tmp_path / "test_nested_array.json") detections = sv.Detections( xyxy=np.array([[0, 0, 10, 10]], dtype=np.float32), class_id=np.array([0]), ) with sv.JSONSink(file_name) as sink: sink.append(detections, custom_data={"meta": {"arr": np.array([1, 2, 3])}}) with open(file_name) as f: data = json.load(f) assert data[0]["meta"]["arr"] == [1, 2, 3] @pytest.mark.parametrize( ("custom_data", "expected_value"), [ pytest.param( {"embedding": np.array([1, 2, 3])}, [1, 2, 3], id="custom_data_ndarray_length_mismatch", ) ], ) def test_json_sink_broadcasts_ndarray_when_length_mismatches_detection_count( tmp_path: Any, custom_data: dict[str, Any] | None, expected_value: list[float] ) -> None: """Mismatched ndarray data is broadcast and serialized as a JSON array.""" file_name = str(tmp_path / "test_mismatched_array.json") detections = sv.Detections( xyxy=np.array([[0, 0, 10, 10], [20, 20, 30, 30]]), ) with sv.JSONSink(file_name) as sink: sink.append(detections, custom_data=custom_data) with open(file_name) as f: data = json.load(f) assert data[0]["embedding"] == expected_value assert data[1]["embedding"] == expected_value def test_json_sink_serializes_matching_ndarray_rows_as_json_arrays( tmp_path: Any, ) -> None: """Matching 2D ndarray row data serializes as JSON arrays, not strings.""" file_name = str(tmp_path / "test_matching_array_rows.json") detections = sv.Detections( xyxy=np.array([[0, 0, 10, 10], [20, 20, 30, 30]]), data={"embedding": np.array([[1, 2], [3, 4]])}, ) with sv.JSONSink(file_name) as sink: sink.append(detections, custom_data={"score": np.array([0.5, 0.75])}) with open(file_name) as f: data = json.load(f) assert data[0]["embedding"] == [1, 2] assert data[1]["embedding"] == [3, 4] assert data[0]["score"] == 0.5 assert data[1]["score"] == 0.75 def test_json_default_raises_for_unserializable_type() -> None: """_json_default raises TypeError for non-numpy objects.""" with pytest.raises(TypeError, match="is not JSON serializable"): sv.JSONSink._json_default(object()) def assert_json_equal(file_name, expected_rows): with open(file_name) as file: data = json.load(file) assert data == expected_rows, ( f"Data in JSON file didn't match expected output: {data} != {expected_rows}" ) os.remove(file_name)