9194ef5abd
Docs/Test Workflow / Test docs build (push) Failing after 0s
Check links & references / links-check (push) Failing after 1s
Pytest/Test Workflow / Import Test and Pytest Run (ubuntu-latest, 3.10) (push) Failing after 0s
Pytest/Test Workflow / Import Test and Pytest Run (ubuntu-latest, 3.11) (push) Failing after 0s
PR Conflict Labeler / main (push) Failing after 2s
Pytest/Test Workflow / Import Test and Pytest Run (ubuntu-latest, 3.12) (push) Failing after 2s
Pytest/Test Workflow / Import Test and Pytest Run (ubuntu-latest, 3.13) (push) Failing after 0s
Pytest/Test Workflow / Build this Package (push) Failing after 5s
Pytest/Test Workflow / Import Test and Pytest Run (macos-latest, 3.10) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (macos-latest, 3.11) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (macos-latest, 3.12) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (macos-latest, 3.13) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (windows-latest, 3.10) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (windows-latest, 3.11) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (windows-latest, 3.12) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (windows-latest, 3.13) (push) Has been cancelled
Pytest/Test Workflow / testing-guardian (push) Has been cancelled
565 lines
18 KiB
Python
565 lines
18 KiB
Python
import csv
|
|
import os
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import supervision as sv
|
|
from supervision.detection.tools.csv_sink import CSVSink
|
|
from tests.helpers import _create_detections
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
(
|
|
"detections",
|
|
"custom_data",
|
|
"second_detections",
|
|
"second_custom_data",
|
|
"file_name",
|
|
"expected_result",
|
|
),
|
|
[
|
|
(
|
|
_create_detections(
|
|
xyxy=[[10, 20, 30, 40], [50, 60, 70, 80]],
|
|
confidence=[0.7, 0.8],
|
|
class_id=[0, 0],
|
|
tracker_id=[0, 1],
|
|
data={"class_name": ["person", "person"]},
|
|
),
|
|
{"frame_number": 42},
|
|
_create_detections(
|
|
xyxy=[[15, 25, 35, 45], [55, 65, 75, 85]],
|
|
confidence=[0.6, 0.9],
|
|
class_id=[1, 1],
|
|
tracker_id=[2, 3],
|
|
data={"class_name": ["car", "car"]},
|
|
),
|
|
{"frame_number": 43},
|
|
"test_detections.csv",
|
|
[
|
|
[
|
|
"x_min",
|
|
"y_min",
|
|
"x_max",
|
|
"y_max",
|
|
"class_id",
|
|
"confidence",
|
|
"tracker_id",
|
|
"class_name",
|
|
"frame_number",
|
|
],
|
|
["10.0", "20.0", "30.0", "40.0", "0", "0.7", "0", "person", "42"],
|
|
["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
|
|
(
|
|
_create_detections(
|
|
xyxy=[[60, 70, 80, 90], [100, 110, 120, 130]],
|
|
tracker_id=[4, 5],
|
|
data={"class_name": ["bike", "dog"]},
|
|
),
|
|
{"frame_number": 44},
|
|
_create_detections(
|
|
xyxy=[[65, 75, 85, 95], [105, 115, 125, 135]],
|
|
confidence=[0.5, 0.4],
|
|
data={"class_name": ["tree", "cat"]},
|
|
),
|
|
{"frame_number": 45},
|
|
"test_detections_missing_fields.csv",
|
|
[
|
|
[
|
|
"x_min",
|
|
"y_min",
|
|
"x_max",
|
|
"y_max",
|
|
"class_id",
|
|
"confidence",
|
|
"tracker_id",
|
|
"class_name",
|
|
"frame_number",
|
|
],
|
|
["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
|
|
(
|
|
_create_detections(
|
|
xyxy=[[10, 11, 12, 13]],
|
|
confidence=[0.95],
|
|
data={"class_name": "unknown", "is_detected": True, "score": 1},
|
|
),
|
|
{"frame_number": 46},
|
|
_create_detections(
|
|
xyxy=[[14, 15, 16, 17]],
|
|
data={"class_name": "artifact", "is_detected": False, "score": 0.85},
|
|
),
|
|
{"frame_number": 47},
|
|
"test_detections_varied_data.csv",
|
|
[
|
|
[
|
|
"x_min",
|
|
"y_min",
|
|
"x_max",
|
|
"y_max",
|
|
"class_id",
|
|
"confidence",
|
|
"tracker_id",
|
|
"class_name",
|
|
"frame_number",
|
|
"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
|
|
(
|
|
_create_detections(
|
|
xyxy=[[20, 21, 22, 23]],
|
|
),
|
|
{
|
|
"metadata": {"sensor_id": 101, "location": "north"},
|
|
"tags": ["urgent", "review"],
|
|
},
|
|
_create_detections(
|
|
xyxy=[[14, 15, 16, 17]],
|
|
),
|
|
{
|
|
"metadata": {"sensor_id": 104, "location": "west"},
|
|
"tags": ["not-urgent", "done"],
|
|
},
|
|
"test_detections_complex_data.csv",
|
|
[
|
|
[
|
|
"x_min",
|
|
"y_min",
|
|
"x_max",
|
|
"y_max",
|
|
"class_id",
|
|
"confidence",
|
|
"tracker_id",
|
|
"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
|
|
(
|
|
_create_detections(
|
|
xyxy=[[10, 20, 30, 40], [50, 60, 70, 80]],
|
|
confidence=[0.9, 0.8],
|
|
class_id=[0, 1],
|
|
),
|
|
{"area": np.array([400.0, 400.0]), "frame": 5},
|
|
_create_detections(
|
|
xyxy=[[15, 25, 35, 45]],
|
|
confidence=[0.7],
|
|
class_id=[2],
|
|
),
|
|
{"area": np.array([100.0]), "frame": 6},
|
|
"test_detections_mixed_custom_data.csv",
|
|
[
|
|
[
|
|
"x_min",
|
|
"y_min",
|
|
"x_max",
|
|
"y_max",
|
|
"class_id",
|
|
"confidence",
|
|
"tracker_id",
|
|
"area",
|
|
"frame",
|
|
],
|
|
["10.0", "20.0", "30.0", "40.0", "0", "0.9", "", "400.0", "5"],
|
|
["50.0", "60.0", "70.0", "80.0", "1", "0.8", "", "400.0", "5"],
|
|
["15.0", "25.0", "35.0", "45.0", "2", "0.7", "", "100.0", "6"],
|
|
],
|
|
), # mixed custom_data: ndarray sliced per row, scalar broadcast to all rows
|
|
(
|
|
_create_detections(
|
|
xyxy=[[10, 20, 30, 40], [50, 60, 70, 80]],
|
|
confidence=[0.9, 0.8],
|
|
class_id=[0, 1],
|
|
),
|
|
{"ids": ["a", "b"], "tags": ("x", "y"), "frame": 7},
|
|
_create_detections(
|
|
xyxy=[[15, 25, 35, 45]],
|
|
confidence=[0.7],
|
|
class_id=[2],
|
|
),
|
|
{"ids": ["c"], "tags": ("z",), "frame": 8},
|
|
"test_detections_list_custom_data.csv",
|
|
[
|
|
[
|
|
"x_min",
|
|
"y_min",
|
|
"x_max",
|
|
"y_max",
|
|
"class_id",
|
|
"confidence",
|
|
"tracker_id",
|
|
"frame",
|
|
"ids",
|
|
"tags",
|
|
],
|
|
["10.0", "20.0", "30.0", "40.0", "0", "0.9", "", "7", "a", "x"],
|
|
["50.0", "60.0", "70.0", "80.0", "1", "0.8", "", "7", "b", "y"],
|
|
["15.0", "25.0", "35.0", "45.0", "2", "0.7", "", "8", "c", "z"],
|
|
],
|
|
), # list/tuple custom_data matching detection count is sliced per row
|
|
(
|
|
sv.Detections(
|
|
xyxy=np.array([[10, 20, 30, 40], [50, 60, 70, 80]]),
|
|
data={"labels": ["person", "car"]},
|
|
),
|
|
None,
|
|
sv.Detections(
|
|
xyxy=np.array([[15, 25, 35, 45]]),
|
|
data={"labels": ["bus"]},
|
|
),
|
|
None,
|
|
"test_detections_plain_list_data.csv",
|
|
[
|
|
[
|
|
"x_min",
|
|
"y_min",
|
|
"x_max",
|
|
"y_max",
|
|
"class_id",
|
|
"confidence",
|
|
"tracker_id",
|
|
"labels",
|
|
],
|
|
["10", "20", "30", "40", "", "", "", "person"],
|
|
["50", "60", "70", "80", "", "", "", "car"],
|
|
["15", "25", "35", "45", "", "", "", "bus"],
|
|
],
|
|
), # plain Python list in detections.data is sliced per row without custom_data
|
|
],
|
|
)
|
|
def test_csv_sink(
|
|
detections: sv.Detections,
|
|
custom_data: dict[str, Any] | None,
|
|
second_detections: sv.Detections,
|
|
second_custom_data: dict[str, Any] | None,
|
|
file_name: str,
|
|
expected_result: list[list[Any]],
|
|
) -> None:
|
|
with sv.CSVSink(file_name) as sink:
|
|
if custom_data is None:
|
|
sink.append(detections)
|
|
else:
|
|
sink.append(detections, custom_data)
|
|
if second_custom_data is None:
|
|
sink.append(second_detections)
|
|
else:
|
|
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",
|
|
),
|
|
[
|
|
(
|
|
_create_detections(
|
|
xyxy=[[10, 20, 30, 40], [50, 60, 70, 80]],
|
|
confidence=[0.7, 0.8],
|
|
class_id=[0, 0],
|
|
tracker_id=[0, 1],
|
|
data={"class_name": ["person", "person"]},
|
|
),
|
|
{"frame_number": 42},
|
|
_create_detections(
|
|
xyxy=[[15, 25, 35, 45], [55, 65, 75, 85]],
|
|
confidence=[0.6, 0.9],
|
|
class_id=[1, 1],
|
|
tracker_id=[2, 3],
|
|
data={"class_name": ["car", "car"]},
|
|
),
|
|
{"frame_number": 43},
|
|
"test_detections.csv",
|
|
[
|
|
[
|
|
"x_min",
|
|
"y_min",
|
|
"x_max",
|
|
"y_max",
|
|
"class_id",
|
|
"confidence",
|
|
"tracker_id",
|
|
"class_name",
|
|
"frame_number",
|
|
],
|
|
["10.0", "20.0", "30.0", "40.0", "0", "0.7", "0", "person", "42"],
|
|
["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
|
|
(
|
|
_create_detections(
|
|
xyxy=[[60, 70, 80, 90], [100, 110, 120, 130]],
|
|
tracker_id=[4, 5],
|
|
data={"class_name": ["bike", "dog"]},
|
|
),
|
|
{"frame_number": 44},
|
|
_create_detections(
|
|
xyxy=[[65, 75, 85, 95], [105, 115, 125, 135]],
|
|
confidence=[0.5, 0.4],
|
|
data={"class_name": ["tree", "cat"]},
|
|
),
|
|
{"frame_number": 45},
|
|
"test_detections_missing_fields.csv",
|
|
[
|
|
[
|
|
"x_min",
|
|
"y_min",
|
|
"x_max",
|
|
"y_max",
|
|
"class_id",
|
|
"confidence",
|
|
"tracker_id",
|
|
"class_name",
|
|
"frame_number",
|
|
],
|
|
["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
|
|
(
|
|
_create_detections(
|
|
xyxy=[[10, 11, 12, 13]],
|
|
confidence=[0.95],
|
|
data={"class_name": "unknown", "is_detected": True, "score": 1},
|
|
),
|
|
{"frame_number": 46},
|
|
_create_detections(
|
|
xyxy=[[14, 15, 16, 17]],
|
|
data={"class_name": "artifact", "is_detected": False, "score": 0.85},
|
|
),
|
|
{"frame_number": 47},
|
|
"test_detections_varied_data.csv",
|
|
[
|
|
[
|
|
"x_min",
|
|
"y_min",
|
|
"x_max",
|
|
"y_max",
|
|
"class_id",
|
|
"confidence",
|
|
"tracker_id",
|
|
"class_name",
|
|
"frame_number",
|
|
"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
|
|
(
|
|
_create_detections(
|
|
xyxy=[[20, 21, 22, 23]],
|
|
),
|
|
{
|
|
"metadata": {"sensor_id": 101, "location": "north"},
|
|
"tags": ["urgent", "review"],
|
|
},
|
|
_create_detections(
|
|
xyxy=[[14, 15, 16, 17]],
|
|
),
|
|
{
|
|
"metadata": {"sensor_id": 104, "location": "west"},
|
|
"tags": ["not-urgent", "done"],
|
|
},
|
|
"test_detections_complex_data.csv",
|
|
[
|
|
[
|
|
"x_min",
|
|
"y_min",
|
|
"x_max",
|
|
"y_max",
|
|
"class_id",
|
|
"confidence",
|
|
"tracker_id",
|
|
"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
|
|
],
|
|
)
|
|
def test_csv_sink_manual(
|
|
detections: sv.Detections,
|
|
custom_data: dict[str, Any],
|
|
second_detections: sv.Detections,
|
|
second_custom_data: dict[str, Any],
|
|
file_name: str,
|
|
expected_result: list[list[Any]],
|
|
) -> None:
|
|
sink = sv.CSVSink(file_name)
|
|
sink.open()
|
|
sink.append(detections, custom_data)
|
|
sink.append(second_detections, second_custom_data)
|
|
sink.close()
|
|
|
|
assert_csv_equal(file_name, expected_result)
|
|
|
|
|
|
def assert_csv_equal(file_name, expected_rows) -> None:
|
|
with open(file_name, 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)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("value", "i", "n", "expected"),
|
|
[
|
|
(["x"], 0, 1, "x"),
|
|
(["a", "b", "c"], 0, 2, ["a", "b", "c"]),
|
|
([42, "hello", None], 1, 3, "hello"),
|
|
([["a", "b"], ["c", "d"]], 0, 2, ["a", "b"]),
|
|
("ab", 0, 2, "ab"),
|
|
(("z",), 0, 1, "z"),
|
|
],
|
|
)
|
|
def test_csv_sink_slice_value(value: Any, i: int, n: int, expected: Any) -> None:
|
|
assert CSVSink._slice_value(value, i, n) == expected
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("custom_data", "expected_value"),
|
|
[
|
|
pytest.param(
|
|
{"embedding": np.array([1, 2, 3])},
|
|
np.array([1, 2, 3]),
|
|
id="custom_data_ndarray_length_mismatch",
|
|
)
|
|
],
|
|
)
|
|
def test_csv_sink_broadcasts_ndarray_when_length_mismatches_detection_count(
|
|
custom_data: dict[str, Any] | None, expected_value: np.ndarray
|
|
) -> None:
|
|
"""Mismatched ndarray data is broadcast instead of indexed per row."""
|
|
detections = sv.Detections(
|
|
xyxy=np.array([[0, 0, 10, 10], [20, 20, 30, 30]]),
|
|
)
|
|
|
|
rows = CSVSink.parse_detection_data(detections, custom_data=custom_data)
|
|
|
|
assert len(rows) == 2
|
|
np.testing.assert_array_equal(rows[0]["embedding"], expected_value)
|
|
np.testing.assert_array_equal(rows[1]["embedding"], expected_value)
|