diff --git a/supervision/utils/file.py b/supervision/utils/file.py index dfce7eb0..a827ad18 100644 --- a/supervision/utils/file.py +++ b/supervision/utils/file.py @@ -1,10 +1,77 @@ import json +import csv from pathlib import Path -from typing import List, Optional, Union +from typing import List, Optional, Union, Dict, Any +from supervision.detection.core import Detections import numpy as np import yaml +class JSONSink: + """ + A utility class for saving detection data to a JSON file. This class is designed to + efficiently serialize detection objects into a JSON format, allowing for the inclusion of + bounding box coordinates and additional attributes like confidence, class ID, and tracker ID. + + The class supports the capability to include custom data alongside the detection fields, + providing flexibility for logging various types of information in a structured JSON format. + + Args: + filename (str): The name of the JSON file where the detections will be stored. + Defaults to 'output.json'. + + Usage: + ```python + from supervision.utils.detections import Detections + # Initialize JSONSink with a filename + json_sink = JSONSink('my_detections.json') + + # Assuming detections is an instance of Detections containing detection data + detections = Detections(...) + + # Open the JSONSink context, append detection data, and close the file automatically + with json_sink as sink: + sink.append(detections, custom_data={'frame': 1}) + ``` + """ + def __init__(self, filename: str = 'output.json'): + self.filename: str = filename + self.file: Optional[open] = None + self.data: List[Dict[str, Any]] = [] + + def __enter__(self) -> 'JSONSink': + self.open() + return self + + def __exit__(self, exc_type: Optional[type], exc_val: Optional[Exception], exc_tb: Optional[Any]) -> None: + self.write_and_close() + + def open(self) -> None: + self.file = open(self.filename, 'w') + + def write_and_close(self) -> None: + if self.file: + json.dump(self.data, self.file, indent=4) + self.file.close() + + def append(self, detections: Detections, custom_data: Dict[str, Any] = None) -> None: + for i in range(len(detections.xyxy)): + detection_data = { + 'x_min': int(detections.xyxy[i][0]), + 'y_min': int(detections.xyxy[i][1]), + 'x_max': int(detections.xyxy[i][2]), + 'y_max': int(detections.xyxy[i][3]), + 'class_id': int(detections.class_id[i]), + 'confidence': float(detections.confidence[i]), + 'tracker_id': int(detections.tracker_id[i]) + } + + for key, value in detections.data.items(): + detection_data[key] = value[i] if hasattr(value, '__getitem__') else value + if custom_data: + detection_data.update(custom_data) + + self.data.append(detection_data) class NumpyJsonEncoder(json.JSONEncoder): def default(self, obj): diff --git a/test/utils/test_json.py b/test/utils/test_json.py new file mode 100644 index 00000000..09e128e9 --- /dev/null +++ b/test/utils/test_json.py @@ -0,0 +1,110 @@ +import os +import csv +import pytest +import numpy as np +import json +from supervision.utils.file import JSONSink +from supervision.detection.core import Detections + +@pytest.fixture(scope="module") +def detection_instances(): + # Setup detection instances as per the provided example + detections = Detections( + xyxy=np.array([[10, 20, 30, 40], [50, 60, 70, 80]]), + confidence=np.array([0.7, 0.8]), + class_id=np.array([0, 0]), + tracker_id=np.array([0, 1]), + data={'class_name': ['person', 'person']} + ) + + second_detections = Detections( + xyxy=np.array([[15, 25, 35, 45], [55, 65, 75, 85]]), + confidence=np.array([0.6, 0.9]), + class_id=np.array([1, 1]), + tracker_id=np.array([2, 3]), + data={'class_name': ['car', 'car']} + ) + + custom_data = {'frame_number': 42} + second_custom_data = {'frame_number': 43} + + return detections, custom_data, second_detections, second_custom_data + +def test_json_sink(detection_instances): + detections, custom_data, second_detections, second_custom_data = detection_instances + json_filename = "test_detections.json" + expected_data = [ + { + "x_min": 10, "y_min": 20, "x_max": 30, "y_max": 40, + "class_id": 0, "confidence": 0.7, "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.8, "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.6, "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.9, "tracker_id": 3, "class_name": "car", + "frame_number": 43 + } + ] + + # Using the JSONSink class to write the detection data to a JSON file + with JSONSink(filename=json_filename) as sink: + sink.append(detections, custom_data) + sink.append(second_detections, second_custom_data) + + # Read back the JSON file and verify its contents + with open(json_filename, 'r') as file: + data = json.load(file) + assert data == expected_data, f"Data in JSON file did not match expected output: {data} != {expected_data}" + + # Clean up by removing the test JSON file + os.remove(json_filename) + +def test_csv_sink_manual(detection_instances): + detections, custom_data, second_detections, second_custom_data = detection_instances + json_filename = "test_detections.json" + expected_data = [ + { + "x_min": 10, "y_min": 20, "x_max": 30, "y_max": 40, + "class_id": 0, "confidence": 0.7, "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.8, "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.6, "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.9, "tracker_id": 3, "class_name": "car", + "frame_number": 43 + } + ] + + sink = JSONSink(filename=json_filename) + sink.open() + sink.append(detections, custom_data) + sink.append(second_detections, second_custom_data) + sink.write_and_close() + + # Read back the JSON file and verify its contents + with open(json_filename, 'r') as file: + data = json.load(file) + assert data == expected_data, f"Data in JSON file did not match expected output: {data} != {expected_data}" + + # Clean up by removing the test JSON file + os.remove(json_filename) \ No newline at end of file