[CSVSink] - allowing to serialise Detections to a CSV file ready for prod
This commit is contained in:
Adonai Vera
parent
e801879e70
commit
5e339f99fc
|
|
@ -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]:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
Loading…
Reference in New Issue