Applied feedback from SkalskiP in PR #818

Author: AdonaiVera

docs/detection/tools/csv_sink.md

@@ -0,0 +1,8 @@
+---
+comments: true
+status: new
+---
+
+## Save CSV Detection
+
+:::supervision.detection.tools.csv_sink.CSVSink

mkdocs.yml

@@ -41,6 +41,7 @@ nav:
         - Polygon Zone: detection/tools/polygon_zone.md
         - Inference Slicer: detection/tools/inference_slicer.md
         - Detection Smoother: detection/tools/smoother.md
+        - Save CSV Detection: detection/tools/csv_sink.md
     - Annotators: annotators.md
     - Trackers: trackers.md
     - Datasets: datasets.md

supervision/__init__.py

@@ -39,6 +39,7 @@
 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,

supervision/detection/tools/csv_sink.py

@@ -0,0 +1,141 @@
+from __future__ import annotations
+
+import csv
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Union
+
+import numpy as np
+
+from supervision.detection.core import Detections
+
+BASE_HEADER = [
+    "x_min",
+    "y_min",
+    "x_max",
+    "y_max",
+    "class_id",
+    "confidence",
+    "tracker_id",
+]
+
+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'.
+
+    Example:
+        ```python
+        import numpy as np
+        import supervision as sv
+        from ultralytics import YOLO
+        import time
+
+        model = YOLO("yolov8n.pt")
+        tracker = sv.ByteTrack()
+        box_annotator = sv.BoundingBoxAnnotator()
+        label_annotator = sv.LabelAnnotator()
+        csv_sink = sv.CSVSink(...)
+
+        def callback(frame: np.ndarray, _: int) -> np.ndarray:
+            start_time = time.time()
+            results = model(frame)[0]
+            detections = sv.Detections.from_ultralytics(results)
+            detections = tracker.update_with_detections(detections)
+
+            labels = [
+                f"#{tracker_id} {results.names[class_id]}"
+                for class_id, tracker_id
+                in zip(detections.class_id, detections.tracker_id)
+            ]
+            time_frame = (time.time() - start_time)
+
+            csv_sink.append(detections, custom_data={"processing_time": time_frame})
+
+            annotated_frame = box_annotator.annotate(
+                frame.copy(), detections=detections)
+            return label_annotator.annotate(
+                annotated_frame, detections=detections, labels=labels)
+
+        csv_sink.open()
+        sv.process_video(
+            source_path="people-walking.mp4",
+            target_path="result.mp4",
+            callback=callback
+        )
+        csv_sink.close()
+        ``` 
+    """ # noqa: E501 // docs
+
+    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()
+
+    @staticmethod
+    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 = {
+                "x_min": detections.xyxy[i][0],
+                "y_min": detections.xyxy[i][1],
+                "x_max": detections.xyxy[i][2],
+                "y_max": detections.xyxy[i][3],
+                "class_id": detections.class_id[i],
+                "confidence": detections.confidence[i],
+                "tracker_id": detections.tracker_id[i],
+            }
+            if hasattr(detections, "data"):
+                for key, value in detections.data.items():
+                    row[key] = value[i]
+            if custom_data:
+                row.update(custom_data)
+            parsed_rows.append(row)
+        return parsed_rows
+
+    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.")
+        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])
+
+    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

supervision/utils/file.py

@@ -1,93 +1,10 @@
 import json
-import csv
 from pathlib import Path
-from typing import List, Optional, Union, Dict, Any
-from supervision.detection.core import Detections
+from typing import List, Optional, Union
 
 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):
@@ -99,6 +16,7 @@ def default(self, obj):
             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]:
@@ -228,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)
+        yaml.dump(data, outfile, sort_keys=False, default_flow_style=None)

test/utils/test_csv.py renamed to test/detection/test_csv.py

@@ -2,10 +2,9 @@
 import csv
 import pytest
 import numpy as np
-from supervision.utils.file import CSVSink
 from supervision.detection.core import Detections
+import supervision as sv
 
-#pytest test/utils/test_csv.py
 @pytest.fixture(scope="module")
 def detection_instances():
     # Setup detection instances as per the provided example
@@ -42,15 +41,15 @@ def test_csv_sink(detection_instances):
     ]
 
     # Using the CSVSink class to write the detection data to a CSV file
-    with CSVSink(filename=csv_filename) as sink:
+    with sv.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]}"
+            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)
@@ -67,7 +66,7 @@ def test_csv_sink_manual(detection_instances):
     ]
 
     # Using the CSVSink class to write the detection data to a CSV file
-    sink = CSVSink(filename=csv_filename)
+    sink = sv.CSVSink(filename=csv_filename)
     sink.open()
     sink.append(detections, custom_data)
     sink.append(second_detections, second_custom_data)
@@ -77,7 +76,7 @@ def test_csv_sink_manual(detection_instances):
     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]}"
+            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)