Files
wehub-resource-sync 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
chore: import upstream snapshot with attribution
2026-07-13 12:06:10 +08:00

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)