xiaoxue_admin
/
supervision
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(sinks): slice list and tuple custom_data values per row (#2216) 

* test: add regression tests for list/tuple custom_data slicing
* fix: slice list and tuple custom_data values per row
* docs: document custom_data slicing contract in append() docstrings
* docs: add docstring to _slice_value in CSVSink and JSONSink
* docs: add docstring to parse_detection_data in CSVSink and JSONSink
* test: add test for detections.data with plain Python list values
* test: add _slice_value edge-case unit tests
* docs: add per-row slicing note to CSVSink and JSONSink class docstrings

---------

Co-authored-by: jirka <6035284+Borda@users.noreply.github.com>
Co-authored-by: Claude Code <noreply@anthropic.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
This commit is contained in:
Junghwan 2026-04-18 07:24:36 +09:00 committed by GitHub
parent 5d165e04c2
commit 06beb6ff2f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 297 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

56
src/supervision/detection/tools/csv_sink.py
View File

@ -39,6 +39,9 @@ class CSVSink:
CSVSink allows passing custom data alongside detection fields, providing
flexibility for logging various types of information.
When a list or tuple value in custom_data (or detections.data) has the
same length as the detection count, each element is written to the
corresponding detection row; any other value is broadcast to all rows.
Args:
file_name: The name of the CSV file where the detections will be stored.
@ -113,17 +116,56 @@ class CSVSink:
self.file.close()
@staticmethod
def _slice_value(value: Any, i: int) -> Any:
def _slice_value(value: Any, i: int, n: int) -> Any:
"""
Return the i-th element when the value stores per-detection data.
Dispatch rules:
- np.ndarray with ndim == 0: return as-is for broadcasting
- np.ndarray with ndim >= 1: return value[i]
- list or tuple with len equal to n: return value[i]
- any other type: return as-is for broadcasting
Args:
value: Custom-data field value.
i: Zero-based detection index.
n: Total number of detections.
Returns:
Element at position i if value is a per-detection sequence,
otherwise value unchanged.
"""
if isinstance(value, np.ndarray):
return value if value.ndim == 0 else value[i]
if isinstance(value, (list, tuple)) and len(value) == n:
return value[i]
return value
@staticmethod
def parse_detection_data(
detections: Detections, custom_data: dict[str, Any] | None = None
) -> list[dict[str, Any]]:
"""
Convert detections and optional custom data into per-detection rows.
Builds one dictionary per detection containing bounding box coordinates,
detection attributes, and any values from ``detections.data`` or
``custom_data``. List and tuple values in ``custom_data`` with length
equal to ``len(detections.xyxy)`` are sliced one element per row; all
other values are broadcast to every row.
Args:
detections: Detection data to serialize into row dictionaries.
custom_data: Optional extra fields to include in each row.
Returns:
A list of dictionaries, one per detection, containing ``xyxy``
coordinates, ``class_id``, ``confidence``, ``tracker_id``, and any
values from ``detections.data`` or ``custom_data``.
"""
parsed_rows = []
for i in range(len(detections.xyxy)):
n = len(detections.xyxy)
for i in range(n):
row = {
"x_min": detections.xyxy[i][0],
"y_min": detections.xyxy[i][1],
@ -142,11 +184,11 @@ class CSVSink:
if hasattr(detections, "data"):
for key, value in detections.data.items():
row[key] = CSVSink._slice_value(value, i)
row[key] = CSVSink._slice_value(value, i, n)
if custom_data:
for key, value in custom_data.items():
row[key] = CSVSink._slice_value(value, i)
row[key] = CSVSink._slice_value(value, i, n)
parsed_rows.append(row)
return parsed_rows
@ -159,7 +201,11 @@ class CSVSink:
Args:
detections: The detection data.
custom_data: Custom data to include.
custom_data: Custom data to include. Scalars, dictionaries, and
other non-sequence values are broadcast to every detection in
this batch. NumPy arrays, lists, and tuples with length equal
to ``len(detections)`` are sliced per detection; other lists
and tuples are broadcast unchanged.
"""
if not self.writer:
raise Exception(

56
src/supervision/detection/tools/json_sink.py
View File

@ -21,6 +21,9 @@ class JSONSink:
JSONSink allows passing custom data alongside detection fields, providing
flexibility for logging various types of information.
When a list or tuple value in custom_data (or detections.data) has the
same length as the detection count, each element is written to the
corresponding detection row; any other value is broadcast to all rows.
Args:
file_name: The name of the JSON file where the detections will be stored.
@ -85,17 +88,56 @@ class JSONSink:
self.file.close()
@staticmethod
def _slice_value(value: Any, i: int) -> Any:
def _slice_value(value: Any, i: int, n: int) -> Any:
"""
Return the i-th element when the value stores per-detection data.
Dispatch rules:
- np.ndarray with ndim == 0: return as-is for broadcasting
- np.ndarray with ndim >= 1: return value[i]
- list or tuple with len equal to n: return value[i]
- any other type: return as-is for broadcasting
Args:
value: Custom-data field value.
i: Zero-based detection index.
n: Total number of detections.
Returns:
Element at position i if value is a per-detection sequence,
otherwise value unchanged.
"""
if isinstance(value, np.ndarray):
return value if value.ndim == 0 else value[i]
if isinstance(value, (list, tuple)) and len(value) == n:
return value[i]
return value
@staticmethod
def parse_detection_data(
detections: Detections, custom_data: dict[str, Any] | None = None
) -> list[dict[str, Any]]:
"""
Convert detections and optional custom data into per-detection rows.
Builds one dictionary per detection containing bounding box coordinates,
detection attributes, and any values from ``detections.data`` or
``custom_data``. List and tuple values in ``custom_data`` with length
equal to ``len(detections.xyxy)`` are sliced one element per row; all
other values are broadcast to every row.
Args:
detections: Detection data to serialize into row dictionaries.
custom_data: Optional extra fields to include in each row.
Returns:
A list of dictionaries, one per detection, containing ``xyxy``
coordinates, ``class_id``, ``confidence``, ``tracker_id``, and any
values from ``detections.data`` or ``custom_data``.
"""
parsed_rows = []
for i in range(len(detections.xyxy)):
n = len(detections.xyxy)
for i in range(n):
row = {
"x_min": float(detections.xyxy[i][0]),
"y_min": float(detections.xyxy[i][1]),
@ -114,11 +156,11 @@ class JSONSink:
if hasattr(detections, "data"):
for key, value in detections.data.items():
row[key] = str(JSONSink._slice_value(value, i))
row[key] = str(JSONSink._slice_value(value, i, n))
if custom_data:
for key, value in custom_data.items():
v = JSONSink._slice_value(value, i)
v = JSONSink._slice_value(value, i, n)
row[key] = str(v) if isinstance(value, np.ndarray) else v
parsed_rows.append(row)
@ -132,7 +174,11 @@ class JSONSink:
Args:
detections: The detection data.
custom_data: Custom data to include.
custom_data: Custom data to include. Scalars, dictionaries, and
other non-sequence values are broadcast to every detection in
this batch. NumPy arrays, lists, and tuples with length equal
to ``len(detections)`` are sliced per detection; other lists
and tuples are broadcast unchanged.
"""
parsed_rows = JSONSink.parse_detection_data(detections, custom_data)
self.data.extend(parsed_rows)

92
tests/detection/test_csv.py
View File

@ -1,3 +1,5 @@
from __future__ import annotations
import csv
import os
from typing import Any
@ -6,6 +8,7 @@ import numpy as np
import pytest
import supervision as sv
from supervision.detection.tools.csv_sink import CSVSink
from tests.helpers import _create_detections
@ -225,19 +228,85 @@ from tests.helpers import _create_detections
["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],
custom_data: dict[str, Any] | None,
second_detections: sv.Detections,
second_custom_data: dict[str, Any],
second_custom_data: dict[str, Any] | None,
file_name: str,
expected_result: list[list[Any]],
) -> None:
with sv.CSVSink(file_name) as sink:
sink.append(detections, custom_data)
sink.append(second_detections, second_custom_data)
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)
@ -455,3 +524,18 @@ def assert_csv_equal(file_name, expected_rows):
)
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

111
tests/detection/test_json.py
View File

@ -1,3 +1,5 @@
from __future__ import annotations
import json
import os
from typing import Any
@ -275,19 +277,120 @@ from tests.helpers import _create_detections
},
],
), # numpy array in custom_data sliced per detection row
(
_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")},
_create_detections(
xyxy=[[15, 25, 35, 45]],
confidence=[0.7],
class_id=[2],
),
{"ids": ["c"], "tags": ("z",)},
"test_detections_list_custom_data.json",
[
{
"x_min": 10,
"y_min": 20,
"x_max": 30,
"y_max": 40,
"class_id": 0,
"confidence": 0.8999999761581421,
"tracker_id": "",
"ids": "a",
"tags": "x",
},
{
"x_min": 50,
"y_min": 60,
"x_max": 70,
"y_max": 80,
"class_id": 1,
"confidence": 0.800000011920929,
"tracker_id": "",
"ids": "b",
"tags": "y",
},
{
"x_min": 15,
"y_min": 25,
"x_max": 35,
"y_max": 45,
"class_id": 2,
"confidence": 0.699999988079071,
"tracker_id": "",
"ids": "c",
"tags": "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.json",
[
{
"x_min": 10.0,
"y_min": 20.0,
"x_max": 30.0,
"y_max": 40.0,
"class_id": "",
"confidence": "",
"tracker_id": "",
"labels": "person",
},
{
"x_min": 50.0,
"y_min": 60.0,
"x_max": 70.0,
"y_max": 80.0,
"class_id": "",
"confidence": "",
"tracker_id": "",
"labels": "car",
},
{
"x_min": 15.0,
"y_min": 25.0,
"x_max": 35.0,
"y_max": 45.0,
"class_id": "",
"confidence": "",
"tracker_id": "",
"labels": "bus",
},
],
), # plain Python list in detections.data is sliced per row without custom_data
],
)
def test_json_sink(
detections: sv.Detections,
custom_data: dict[str, Any],
custom_data: dict[str, Any] | None,
second_detections: sv.Detections,
second_custom_data: dict[str, Any],
second_custom_data: dict[str, Any] | None,
file_name: str,
expected_result: list[list[Any]],
) -> None:
with sv.JSONSink(file_name) as sink:
sink.append(detections, custom_data)
sink.append(second_detections, second_custom_data)
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_json_equal(file_name, expected_result)