diff --git a/docs/detection/tools/save_detections.md b/docs/detection/tools/save_detections.md index e49bc129..bdd7c9dc 100644 --- a/docs/detection/tools/save_detections.md +++ b/docs/detection/tools/save_detections.md @@ -9,4 +9,4 @@ status: new

CSV Sink

-:::supervision.detection.tools.csv_sink.CSVSink \ No newline at end of file +:::supervision.detection.tools.csv_sink.CSVSink diff --git a/supervision/detection/tools/csv_sink.py b/supervision/detection/tools/csv_sink.py index d80aeea1..c1226c48 100644 --- a/supervision/detection/tools/csv_sink.py +++ b/supervision/detection/tools/csv_sink.py @@ -4,7 +4,6 @@ import csv import os from typing import Any, Dict, List, Optional -import numpy as np from supervision.detection.core import Detections @@ -18,6 +17,7 @@ BASE_HEADER = [ "tracker_id", ] + class CSVSink: """ A utility class for saving detection data to a CSV file. This class is designed to @@ -63,7 +63,7 @@ class CSVSink: self.file: Optional[open] = None self.writer: Optional[csv.writer] = None self.header_written = False - self.field_names = [] + self.field_names = [] def __enter__(self) -> CSVSink: """ @@ -104,7 +104,7 @@ class CSVSink: parent_directory = os.path.dirname(self.file_name) if parent_directory and not os.path.exists(parent_directory): os.makedirs(parent_directory) - + self.file = open(self.file_name, "w", newline="") self.writer = csv.writer(self.file) @@ -139,18 +139,24 @@ class CSVSink: "y_min": detections.xyxy[i][1], "x_max": detections.xyxy[i][2], "y_max": detections.xyxy[i][3], - "class_id": "" if detections.class_id is None else str(detections.class_id[i]), - "confidence": "" if detections.confidence is None else str(detections.confidence[i]), - "tracker_id": "" if detections.tracker_id is None else str(detections.tracker_id[i]), + "class_id": "" + if detections.class_id is None + else str(detections.class_id[i]), + "confidence": "" + if detections.confidence is None + else str(detections.confidence[i]), + "tracker_id": "" + if detections.tracker_id is None + else str(detections.tracker_id[i]), } if hasattr(detections, "data"): for key, value in detections.data.items(): - if value.ndim == 0: + if value.ndim == 0: row[key] = value else: row[key] = value[i] - + if custom_data: row.update(custom_data) parsed_rows.append(row) diff --git a/test/detection/test_csv.py b/test/detection/test_csv.py index 7b9f13f6..f87f5931 100644 --- a/test/detection/test_csv.py +++ b/test/detection/test_csv.py @@ -1,6 +1,4 @@ import csv -import os -from contextlib import ExitStack as DoesNotRaise from test.test_utils import mock_detections from typing import Any, Dict, List @@ -18,7 +16,7 @@ import supervision as sv confidence=[0.7, 0.8], class_id=[0, 0], tracker_id=[0, 1], - data={"class_name": ["person", "person"]} + data={"class_name": ["person", "person"]}, ), {"frame_number": 42}, mock_detections( @@ -26,7 +24,7 @@ import supervision as sv confidence=[0.6, 0.9], class_id=[1, 1], tracker_id=[2, 3], - data={"class_name": ["car", "car"]} + data={"class_name": ["car", "car"]}, ), {"frame_number": 43}, "test_detections.csv", @@ -46,19 +44,19 @@ import supervision as sv ["50.0", "60.0", "70.0", "80.0", "0", "0.8", "1", "person", "42"], ["15.0", "25.0", "35.0", "45.0", "1", "0.6", "2", "car", "43"], ["55.0", "65.0", "75.0", "85.0", "1", "0.9", "3", "car", "43"], - ] + ], ), # multiple detections ( mock_detections( xyxy=[[60, 70, 80, 90], [100, 110, 120, 130]], tracker_id=[4, 5], - data={"class_name": ["bike", "dog"]} + data={"class_name": ["bike", "dog"]}, ), {"frame_number": 44}, mock_detections( xyxy=[[65, 75, 85, 95], [105, 115, 125, 135]], confidence=[0.5, 0.4], - data={"class_name": ["tree", "cat"]} + data={"class_name": ["tree", "cat"]}, ), {"frame_number": 45}, "test_detections_missing_fields.csv", @@ -74,22 +72,22 @@ import supervision as sv "class_name", "frame_number", ], - ["60.0", "70.0", "80.0", "90.0", "", "", "4", "bike", "44"], + ["60.0", "70.0", "80.0", "90.0", "", "", "4", "bike", "44"], ["100.0", "110.0", "120.0", "130.0", "", "", "5", "dog", "44"], ["65.0", "75.0", "85.0", "95.0", "", "0.5", "", "tree", "45"], ["105.0", "115.0", "125.0", "135.0", "", "0.4", "", "cat", "45"], - ] + ], ), # missing fields ( mock_detections( xyxy=[[10, 11, 12, 13]], confidence=[0.95], - data={"class_name": "unknown", "is_detected": True, "score": 1} + data={"class_name": "unknown", "is_detected": True, "score": 1}, ), {"frame_number": 46}, mock_detections( xyxy=[[14, 15, 16, 17]], - data={"class_name": "artifact", "is_detected": False, "score": 0.85} + data={"class_name": "artifact", "is_detected": False, "score": 0.85}, ), {"frame_number": 47}, "test_detections_varied_data.csv", @@ -107,19 +105,49 @@ import supervision as sv "is_detected", "score", ], - ["10.0", "11.0", "12.0", "13.0", "", "0.95", "", "unknown", "46", "True", "1"], - ["14.0", "15.0", "16.0", "17.0", "", "", "", "artifact", "47", "False", "0.85"], - ] - ), # Inconsistent Data Types + [ + "10.0", + "11.0", + "12.0", + "13.0", + "", + "0.95", + "", + "unknown", + "46", + "True", + "1", + ], + [ + "14.0", + "15.0", + "16.0", + "17.0", + "", + "", + "", + "artifact", + "47", + "False", + "0.85", + ], + ], + ), # Inconsistent Data Types ( mock_detections( xyxy=[[20, 21, 22, 23]], ), - {"metadata": {"sensor_id": 101, "location": "north"}, "tags": ["urgent", "review"]}, + { + "metadata": {"sensor_id": 101, "location": "north"}, + "tags": ["urgent", "review"], + }, mock_detections( xyxy=[[14, 15, 16, 17]], ), - {"metadata": {"sensor_id": 104, "location": "west"}, "tags": ["not-urgent", "done"]}, + { + "metadata": {"sensor_id": 104, "location": "west"}, + "tags": ["not-urgent", "done"], + }, "test_detections_complex_data.csv", [ [ @@ -133,28 +161,47 @@ import supervision as sv "metadata", "tags", ], - ["20.0", "21.0", "22.0", "23.0", "", "", "", "{'sensor_id': 101, 'location': 'north'}", "['urgent', 'review']"], - ["14.0", "15.0", "16.0", "17.0", "", "", "", "{'sensor_id': 104, 'location': 'west'}", "['not-urgent', 'done']"], - ] - ), # Complex Data + [ + "20.0", + "21.0", + "22.0", + "23.0", + "", + "", + "", + "{'sensor_id': 101, 'location': 'north'}", + "['urgent', 'review']", + ], + [ + "14.0", + "15.0", + "16.0", + "17.0", + "", + "", + "", + "{'sensor_id': 104, 'location': 'west'}", + "['not-urgent', 'done']", + ], + ], + ), # Complex Data ], ) - def test_csv_sink( detections: mock_detections, custom_data: Dict[str, Any], second_detections: mock_detections, second_custom_data: Dict[str, Any], file_name: str, - expected_result: List[List[Any]] + expected_result: List[List[Any]], ) -> None: - with sv.CSVSink(file_name) as sink: sink.append(detections, custom_data) sink.append(second_detections, second_custom_data) assert_csv_equal(file_name, expected_result) + @pytest.mark.parametrize( "detections, custom_data, second_detections, second_custom_data, file_name, expected_result", [ @@ -164,7 +211,7 @@ def test_csv_sink( confidence=[0.7, 0.8], class_id=[0, 0], tracker_id=[0, 1], - data={"class_name": ["person", "person"]} + data={"class_name": ["person", "person"]}, ), {"frame_number": 42}, mock_detections( @@ -172,7 +219,7 @@ def test_csv_sink( confidence=[0.6, 0.9], class_id=[1, 1], tracker_id=[2, 3], - data={"class_name": ["car", "car"]} + data={"class_name": ["car", "car"]}, ), {"frame_number": 43}, "test_detections.csv", @@ -192,19 +239,19 @@ def test_csv_sink( ["50.0", "60.0", "70.0", "80.0", "0", "0.8", "1", "person", "42"], ["15.0", "25.0", "35.0", "45.0", "1", "0.6", "2", "car", "43"], ["55.0", "65.0", "75.0", "85.0", "1", "0.9", "3", "car", "43"], - ] + ], ), # multiple detections ( mock_detections( xyxy=[[60, 70, 80, 90], [100, 110, 120, 130]], tracker_id=[4, 5], - data={"class_name": ["bike", "dog"]} + data={"class_name": ["bike", "dog"]}, ), {"frame_number": 44}, mock_detections( xyxy=[[65, 75, 85, 95], [105, 115, 125, 135]], confidence=[0.5, 0.4], - data={"class_name": ["tree", "cat"]} + data={"class_name": ["tree", "cat"]}, ), {"frame_number": 45}, "test_detections_missing_fields.csv", @@ -220,22 +267,22 @@ def test_csv_sink( "class_name", "frame_number", ], - ["60.0", "70.0", "80.0", "90.0", "", "", "4", "bike", "44"], + ["60.0", "70.0", "80.0", "90.0", "", "", "4", "bike", "44"], ["100.0", "110.0", "120.0", "130.0", "", "", "5", "dog", "44"], ["65.0", "75.0", "85.0", "95.0", "", "0.5", "", "tree", "45"], ["105.0", "115.0", "125.0", "135.0", "", "0.4", "", "cat", "45"], - ] + ], ), # missing fields ( mock_detections( xyxy=[[10, 11, 12, 13]], confidence=[0.95], - data={"class_name": "unknown", "is_detected": True, "score": 1} + data={"class_name": "unknown", "is_detected": True, "score": 1}, ), {"frame_number": 46}, mock_detections( xyxy=[[14, 15, 16, 17]], - data={"class_name": "artifact", "is_detected": False, "score": 0.85} + data={"class_name": "artifact", "is_detected": False, "score": 0.85}, ), {"frame_number": 47}, "test_detections_varied_data.csv", @@ -253,19 +300,49 @@ def test_csv_sink( "is_detected", "score", ], - ["10.0", "11.0", "12.0", "13.0", "", "0.95", "", "unknown", "46", "True", "1"], - ["14.0", "15.0", "16.0", "17.0", "", "", "", "artifact", "47", "False", "0.85"], - ] - ), # Inconsistent Data Types + [ + "10.0", + "11.0", + "12.0", + "13.0", + "", + "0.95", + "", + "unknown", + "46", + "True", + "1", + ], + [ + "14.0", + "15.0", + "16.0", + "17.0", + "", + "", + "", + "artifact", + "47", + "False", + "0.85", + ], + ], + ), # Inconsistent Data Types ( mock_detections( xyxy=[[20, 21, 22, 23]], ), - {"metadata": {"sensor_id": 101, "location": "north"}, "tags": ["urgent", "review"]}, + { + "metadata": {"sensor_id": 101, "location": "north"}, + "tags": ["urgent", "review"], + }, mock_detections( xyxy=[[14, 15, 16, 17]], ), - {"metadata": {"sensor_id": 104, "location": "west"}, "tags": ["not-urgent", "done"]}, + { + "metadata": {"sensor_id": 104, "location": "west"}, + "tags": ["not-urgent", "done"], + }, "test_detections_complex_data.csv", [ [ @@ -279,13 +356,32 @@ def test_csv_sink( "metadata", "tags", ], - ["20.0", "21.0", "22.0", "23.0", "", "", "", "{'sensor_id': 101, 'location': 'north'}", "['urgent', 'review']"], - ["14.0", "15.0", "16.0", "17.0", "", "", "", "{'sensor_id': 104, 'location': 'west'}", "['not-urgent', 'done']"], - ] - ), # Complex Data + [ + "20.0", + "21.0", + "22.0", + "23.0", + "", + "", + "", + "{'sensor_id': 101, 'location': 'north'}", + "['urgent', 'review']", + ], + [ + "14.0", + "15.0", + "16.0", + "17.0", + "", + "", + "", + "{'sensor_id': 104, 'location': 'west'}", + "['not-urgent', 'done']", + ], + ], + ), # Complex Data ], ) - def test_csv_sink_manual( detections: mock_detections, custom_data: Dict[str, Any], @@ -302,12 +398,13 @@ def test_csv_sink_manual( assert_csv_equal(file_name, expected_result) -def assert_csv_equal(file_name, expected_rows): + +def assert_csv_equal(file_name, expected_rows): with open(file_name, 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]}" - - #os.remove(file_name) \ No newline at end of file + + # os.remove(file_name)