112 lines
3.5 KiB
Python
112 lines
3.5 KiB
Python
import csv
|
|
import os
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import supervision as sv
|
|
from supervision.detection.core import Detections
|
|
|
|
|
|
@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 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 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",
|
|
],
|
|
[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 = sv.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 didn't match expected output: {row} != {expected_rows[i]}"
|
|
|
|
# Clean up by removing the test CSV file
|
|
os.remove(csv_filename)
|