diff --git a/supervision/__init__.py b/supervision/__init__.py index c5e6d641..76abe599 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -36,10 +36,10 @@ from supervision.dataset.core import ( from supervision.detection.annotate import BoxAnnotator from supervision.detection.core import Detections from supervision.detection.line_counter import LineZone, LineZoneAnnotator +from supervision.detection.tools.csv_sink import CSVSink from supervision.detection.tools.inference_slicer import InferenceSlicer from supervision.detection.tools.polygon_zone import PolygonZone, PolygonZoneAnnotator from supervision.detection.tools.smoother import DetectionsSmoother -from supervision.detection.tools.csv_sink import CSVSink from supervision.detection.utils import ( box_iou_batch, calculate_masks_centroids, diff --git a/supervision/detection/tools/csv_sink.py b/supervision/detection/tools/csv_sink.py index f4e0c8db..c520247c 100644 --- a/supervision/detection/tools/csv_sink.py +++ b/supervision/detection/tools/csv_sink.py @@ -1,10 +1,8 @@ from __future__ import annotations import csv -from pathlib import Path -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, List, Optional -import numpy as np from supervision.detection.core import Detections @@ -18,6 +16,7 @@ BASE_HEADER = [ "tracker_id", ] + class CSVSink: """ A utility class for saving detection data to a CSV file. This class is designed to @@ -71,8 +70,8 @@ class CSVSink: callback=callback ) csv_sink.close() - ``` - """ # noqa: E501 // docs + ``` + """ # noqa: E501 // docs def __init__(self, filename: str = "output.csv"): self.filename = filename @@ -102,7 +101,9 @@ class CSVSink: self.file.close() @staticmethod - def parse_detection_data(detections: Detections, custom_data: Dict[str, Any] = None) -> List[Dict[str, Any]]: + def parse_detection_data( + detections: Detections, custom_data: Dict[str, Any] = None + ) -> List[Dict[str, Any]]: parsed_rows = [] for i in range(len(detections.xyxy)): row = { @@ -122,20 +123,26 @@ class CSVSink: parsed_rows.append(row) return parsed_rows - def append(self, detections: Detections, custom_data: Dict[str, Any] = None) -> None: + 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.") + raise Exception( + f"Cannot append to CSV: The file '{self.filename}' is not open." + ) if not self.header_written: self.write_header(detections, custom_data) parsed_rows = CSVSink.parse_detection_data(detections, custom_data) for row in parsed_rows: - self.writer.writerow([row.get(fieldname, "") for fieldname in self.fieldnames]) + self.writer.writerow( + [row.get(fieldname, "") for fieldname in self.fieldnames] + ) - def write_header( - self, detections: Detections, custom_data: Dict[str, Any] - ) -> None: - dynamic_header = sorted(set(custom_data.keys()) | set(getattr(detections, "data", {}).keys())) + def write_header(self, detections: Detections, custom_data: Dict[str, Any]) -> None: + dynamic_header = sorted( + set(custom_data.keys()) | set(getattr(detections, "data", {}).keys()) + ) self.fieldnames = BASE_HEADER + dynamic_header self.writer.writerow(self.fieldnames) self.header_written = True diff --git a/supervision/utils/file.py b/supervision/utils/file.py index 84eb637c..dfce7eb0 100644 --- a/supervision/utils/file.py +++ b/supervision/utils/file.py @@ -146,4 +146,4 @@ def save_yaml_file(data: dict, file_path: str) -> None: """ with open(file_path, "w") as outfile: - yaml.dump(data, outfile, sort_keys=False, default_flow_style=None) \ No newline at end of file + yaml.dump(data, outfile, sort_keys=False, default_flow_style=None) diff --git a/test/detection/test_csv.py b/test/detection/test_csv.py index b4412c3c..9af60b20 100644 --- a/test/detection/test_csv.py +++ b/test/detection/test_csv.py @@ -1,9 +1,12 @@ -import os import csv -import pytest +import os + import numpy as np -from supervision.detection.core import Detections +import pytest + import supervision as sv +from supervision.detection.core import Detections + @pytest.fixture(scope="module") def detection_instances(): @@ -13,31 +16,42 @@ def detection_instances(): 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'])} + 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'])} + data={"class_name": np.array(["car", "car"])}, ) - - custom_data = {'frame_number': 42} - second_custom_data = {'frame_number': 43} - + + 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"], + [ + "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] + [55, 65, 75, 85, 1, 0.9, 3, "car", 43], ] # Using the CSVSink class to write the detection data to a CSV file @@ -46,23 +60,36 @@ def test_csv_sink(detection_instances): 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: + 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 didn't match expected output: {row} != {expected_rows[i]}" + assert ( + [str(item) for item in expected_rows[i]] == row + ), f"Row in CSV didn't 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"], + [ + "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] + [55, 65, 75, 85, 1, 0.9, 3, "car", 43], ] # Using the CSVSink class to write the detection data to a CSV file @@ -73,10 +100,12 @@ def test_csv_sink_manual(detection_instances): sink.close() # Read back the CSV file and verify its contents - with open(csv_filename, mode='r', newline='') as file: + 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 didn't match expected output: {row} != {expected_rows[i]}" + assert ( + [str(item) for item in expected_rows[i]] == row + ), f"Row in CSV didn't match expected output: {row} != {expected_rows[i]}" # Clean up by removing the test CSV file os.remove(csv_filename)