84 lines
3.3 KiB
Python
84 lines
3.3 KiB
Python
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)
|