From 5e339f99fc14067323bb69f227f2dd5f0442374b Mon Sep 17 00:00:00 2001 From: Adonai Vera <45982251+AdonaiVera@users.noreply.github.com> Date: Wed, 31 Jan 2024 01:32:49 -0500 Subject: [PATCH] [CSVSink] - allowing to serialise Detections to a CSV file ready for prod --- supervision/utils/file.py | 86 ++++++++++++++++++++++++++++++++++++++- test/utils/test_csv.py | 83 +++++++++++++++++++++++++++++++++++++ 2 files changed, 167 insertions(+), 2 deletions(-) create mode 100644 test/utils/test_csv.py diff --git a/supervision/utils/file.py b/supervision/utils/file.py index dfce7eb0..9c32d775 100644 --- a/supervision/utils/file.py +++ b/supervision/utils/file.py @@ -1,10 +1,93 @@ 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 CSVSink: + """ + A utility class for saving detection data to a CSV file. This class is designed to + efficiently serialize detection objects into a CSV 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. + + Args: + filename (str): The name of the CSV file where the detections will be stored. + Defaults to 'output.csv'. + + Usage: + ```python + from supervision.utils.detections import Detections + # Initialize CSVSink with a filename + csv_sink = CSVSink('my_detections.csv') + + # Assuming detections is an instance of Detections containing detection data + detections = Detections(...) + + # Open the CSVSink context, append detection data, and close the file automatically + with csv_sink as sink: + sink.append(detections, custom_data={'frame': 1}) + ``` + """ + def __init__(self, filename: str = 'output.csv'): + self.filename = filename + self.file: Optional[open] = None + self.writer: Optional[csv.writer] = None + self.header_written = False + self.fieldnames = [] # To keep track of header names + + def __enter__(self) -> 'CSVSink': + self.open() + return self + + def __exit__(self, exc_type: Optional[type], exc_val: Optional[Exception], exc_tb: Optional[Any]) -> None: + self.close() + + def open(self) -> None: + self.file = open(self.filename, 'w', newline='') + self.writer = csv.writer(self.file) + + def close(self) -> None: + if self.file: + self.file.close() + + def append(self, detections: Detections, custom_data: Dict[str, Any] = None) -> None: + if not self.writer: + raise Exception(f"Cannot append to CSV: The file '{self.filename}' is not open. Ensure that the 'open' method is called before appending data.") + if not self.header_written: + self.write_header(detections, custom_data) + for i in range(len(detections.xyxy)): + self.write_detection_row(detections, i, custom_data) + + def write_header(self, detections: Detections, custom_data: Dict[str, Any]) -> None: + base_header = ['x_min', 'y_min', 'x_max', 'y_max', 'class_id', 'confidence', 'tracker_id'] + dynamic_header = sorted(set(custom_data.keys()) | set(getattr(detections, 'data', {}).keys())) + self.fieldnames = base_header + dynamic_header + self.dynamic_fields = dynamic_header # Store only the dynamic part + self.writer.writerow(self.fieldnames) + self.header_written = True + + def write_detection_row(self, detections: Detections, index: int, custom_data: Dict[str, Any]) -> None: + row_base = [ + detections.xyxy[index][0], detections.xyxy[index][1], + detections.xyxy[index][2], detections.xyxy[index][3], + detections.class_id[index], detections.confidence[index], + detections.tracker_id[index] + ] + dynamic_data = {} + if hasattr(detections, 'data'): + for key, value in detections.data.items(): + dynamic_data[key] = value[index] + if custom_data: + dynamic_data.update(custom_data) + + row_dynamic = [dynamic_data.get(key) for key in self.fieldnames[7:]] + self.writer.writerow(row_base + row_dynamic) class NumpyJsonEncoder(json.JSONEncoder): def default(self, obj): @@ -16,7 +99,6 @@ class NumpyJsonEncoder(json.JSONEncoder): return obj.tolist() return super(NumpyJsonEncoder, self).default(obj) - def list_files_with_extensions( directory: Union[str, Path], extensions: Optional[List[str]] = None ) -> List[Path]: diff --git a/test/utils/test_csv.py b/test/utils/test_csv.py new file mode 100644 index 00000000..c4dd5045 --- /dev/null +++ b/test/utils/test_csv.py @@ -0,0 +1,83 @@ +import os +import csv +import pytest +import numpy as np +from supervision.utils.file import CSVSink +from supervision.detection.core import Detections + +#pytest test/utils/test_csv.py +@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': np.array(['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': np.array(['car', 'car'])} + ) + + custom_data = {'frame_number': 42} + second_custom_data = {'frame_number': 43} + + return detections, custom_data, second_detections, second_custom_data + +def test_csv_sink(detection_instances): + detections, custom_data, second_detections, second_custom_data = detection_instances + csv_filename = "test_detections.csv" + expected_rows = [ + ["x_min", "y_min", "x_max", "y_max", "class_id", "confidence", "tracker_id", "class_name", "frame_number"], + [10, 20, 30, 40, 0, 0.7, 0, "person", 42], + [50, 60, 70, 80, 0, 0.8, 1, "person", 42], + [15, 25, 35, 45, 1, 0.6, 2, "car", 43], + [55, 65, 75, 85, 1, 0.9, 3, "car", 43] + ] + + # Using the CSVSink class to write the detection data to a CSV file + with CSVSink(filename=csv_filename) as sink: + sink.append(detections, custom_data) + sink.append(second_detections, second_custom_data) + + # Read back the CSV file and verify its contents + with open(csv_filename, mode='r', newline='') as file: + reader = csv.reader(file) + for i, row in enumerate(reader): + assert [str(item) for item in expected_rows[i]] == row, f"Row in CSV file did not match expected output: {row} != {expected_rows[i]}" + + # Clean up by removing the test CSV file + os.remove(csv_filename) + +def test_csv_sink_manual(detection_instances): + detections, custom_data, second_detections, second_custom_data = detection_instances + csv_filename = "test_detections.csv" + expected_rows = [ + ["x_min", "y_min", "x_max", "y_max", "class_id", "confidence", "tracker_id", "class_name", "frame_number"], + [10, 20, 30, 40, 0, 0.7, 0, "person", 42], + [50, 60, 70, 80, 0, 0.8, 1, "person", 42], + [15, 25, 35, 45, 1, 0.6, 2, "car", 43], + [55, 65, 75, 85, 1, 0.9, 3, "car", 43] + ] + + # Using the CSVSink class to write the detection data to a CSV file + sink = CSVSink(filename=csv_filename) + sink.open() + sink.append(detections, custom_data) + sink.append(second_detections, second_custom_data) + sink.close() + + # Read back the CSV file and verify its contents + with open(csv_filename, mode='r', newline='') as file: + reader = csv.reader(file) + for i, row in enumerate(reader): + assert [str(item) for item in expected_rows[i]] == row, f"Row in CSV file did not match expected output: {row} != {expected_rows[i]}" + + # Clean up by removing the test CSV file + os.remove(csv_filename)