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

fix(pre_commit): 🎨 auto format pre-commit hooks

Browse Source
This commit is contained in:
pre-commit-ci[bot] committed 2024-02-01 22:56:02 +00:00
1 parent 788bc317fa
commit aba5c0a07b
3 files changed
+159 -56
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/detection/tools/save_detections.md
+1 -1
View File
@@ -9,4 +9,4 @@ status: new
<h2>CSV Sink</h2>
</div>
:::supervision.detection.tools.csv_sink.CSVSink
:::supervision.detection.tools.csv_sink.CSVSink
supervision/detection/tools/csv_sink.py
+14 -8
View File
@@ -4,7 +4,6 @@ import csv
import os
from typing import Any, Dict, List, Optional
import numpy as np
from supervision.detection.core import Detections
@@ -18,6 +17,7 @@ BASE_HEADER = [
"tracker_id",
]
class CSVSink:
"""
A utility class for saving detection data to a CSV file. This class is designed to
@@ -63,7 +63,7 @@ class CSVSink:
self.file: Optional[open] = None
self.writer: Optional[csv.writer] = None
self.header_written = False
self.field_names = []
self.field_names = []
def __enter__(self) -> CSVSink:
"""
@@ -104,7 +104,7 @@ class CSVSink:
parent_directory = os.path.dirname(self.file_name)
if parent_directory and not os.path.exists(parent_directory):
os.makedirs(parent_directory)
self.file = open(self.file_name, "w", newline="")
self.writer = csv.writer(self.file)
@@ -139,18 +139,24 @@ class CSVSink:
"y_min": detections.xyxy[i][1],
"x_max": detections.xyxy[i][2],
"y_max": detections.xyxy[i][3],
"class_id": "" if detections.class_id is None else str(detections.class_id[i]),
"confidence": "" if detections.confidence is None else str(detections.confidence[i]),
"tracker_id": "" if detections.tracker_id is None else str(detections.tracker_id[i]),
"class_id": ""
if detections.class_id is None
else str(detections.class_id[i]),
"confidence": ""
if detections.confidence is None
else str(detections.confidence[i]),
"tracker_id": ""
if detections.tracker_id is None
else str(detections.tracker_id[i]),
}
if hasattr(detections, "data"):
for key, value in detections.data.items():
if value.ndim == 0:
if value.ndim == 0:
row[key] = value
else:
row[key] = value[i]
if custom_data:
row.update(custom_data)
parsed_rows.append(row)
test/detection/test_csv.py
+144 -47
View File
@@ -1,6 +1,4 @@
import csv
import os
from contextlib import ExitStack as DoesNotRaise
from test.test_utils import mock_detections
from typing import Any, Dict, List
@@ -18,7 +16,7 @@ import supervision as sv
confidence=[0.7, 0.8],
class_id=[0, 0],
tracker_id=[0, 1],
data={"class_name": ["person", "person"]}
data={"class_name": ["person", "person"]},
),
{"frame_number": 42},
mock_detections(
@@ -26,7 +24,7 @@ import supervision as sv
confidence=[0.6, 0.9],
class_id=[1, 1],
tracker_id=[2, 3],
data={"class_name": ["car", "car"]}
data={"class_name": ["car", "car"]},
),
{"frame_number": 43},
"test_detections.csv",
@@ -46,19 +44,19 @@ import supervision as sv
["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
(
mock_detections(
xyxy=[[60, 70, 80, 90], [100, 110, 120, 130]],
tracker_id=[4, 5],
data={"class_name": ["bike", "dog"]}
data={"class_name": ["bike", "dog"]},
),
{"frame_number": 44},
mock_detections(
xyxy=[[65, 75, 85, 95], [105, 115, 125, 135]],
confidence=[0.5, 0.4],
data={"class_name": ["tree", "cat"]}
data={"class_name": ["tree", "cat"]},
),
{"frame_number": 45},
"test_detections_missing_fields.csv",
@@ -74,22 +72,22 @@ import supervision as sv
"class_name",
"frame_number",
],
["60.0", "70.0", "80.0", "90.0", "", "", "4", "bike", "44"],
["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
(
mock_detections(
xyxy=[[10, 11, 12, 13]],
confidence=[0.95],
data={"class_name": "unknown", "is_detected": True, "score": 1}
data={"class_name": "unknown", "is_detected": True, "score": 1},
),
{"frame_number": 46},
mock_detections(
xyxy=[[14, 15, 16, 17]],
data={"class_name": "artifact", "is_detected": False, "score": 0.85}
data={"class_name": "artifact", "is_detected": False, "score": 0.85},
),
{"frame_number": 47},
"test_detections_varied_data.csv",
@@ -107,19 +105,49 @@ import supervision as sv
"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
[
"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
(
mock_detections(
xyxy=[[20, 21, 22, 23]],
),
{"metadata": {"sensor_id": 101, "location": "north"}, "tags": ["urgent", "review"]},
{
"metadata": {"sensor_id": 101, "location": "north"},
"tags": ["urgent", "review"],
},
mock_detections(
xyxy=[[14, 15, 16, 17]],
),
{"metadata": {"sensor_id": 104, "location": "west"}, "tags": ["not-urgent", "done"]},
{
"metadata": {"sensor_id": 104, "location": "west"},
"tags": ["not-urgent", "done"],
},
"test_detections_complex_data.csv",
[
[
@@ -133,28 +161,47 @@ import supervision as sv
"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
[
"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(
detections: mock_detections,
custom_data: Dict[str, Any],
second_detections: mock_detections,
second_custom_data: Dict[str, Any],
file_name: str,
expected_result: List[List[Any]]
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)
assert_csv_equal(file_name, expected_result)
@pytest.mark.parametrize(
"detections, custom_data, second_detections, second_custom_data, file_name, expected_result",
[
@@ -164,7 +211,7 @@ def test_csv_sink(
confidence=[0.7, 0.8],
class_id=[0, 0],
tracker_id=[0, 1],
data={"class_name": ["person", "person"]}
data={"class_name": ["person", "person"]},
),
{"frame_number": 42},
mock_detections(
@@ -172,7 +219,7 @@ def test_csv_sink(
confidence=[0.6, 0.9],
class_id=[1, 1],
tracker_id=[2, 3],
data={"class_name": ["car", "car"]}
data={"class_name": ["car", "car"]},
),
{"frame_number": 43},
"test_detections.csv",
@@ -192,19 +239,19 @@ def test_csv_sink(
["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
(
mock_detections(
xyxy=[[60, 70, 80, 90], [100, 110, 120, 130]],
tracker_id=[4, 5],
data={"class_name": ["bike", "dog"]}
data={"class_name": ["bike", "dog"]},
),
{"frame_number": 44},
mock_detections(
xyxy=[[65, 75, 85, 95], [105, 115, 125, 135]],
confidence=[0.5, 0.4],
data={"class_name": ["tree", "cat"]}
data={"class_name": ["tree", "cat"]},
),
{"frame_number": 45},
"test_detections_missing_fields.csv",
@@ -220,22 +267,22 @@ def test_csv_sink(
"class_name",
"frame_number",
],
["60.0", "70.0", "80.0", "90.0", "", "", "4", "bike", "44"],
["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
(
mock_detections(
xyxy=[[10, 11, 12, 13]],
confidence=[0.95],
data={"class_name": "unknown", "is_detected": True, "score": 1}
data={"class_name": "unknown", "is_detected": True, "score": 1},
),
{"frame_number": 46},
mock_detections(
xyxy=[[14, 15, 16, 17]],
data={"class_name": "artifact", "is_detected": False, "score": 0.85}
data={"class_name": "artifact", "is_detected": False, "score": 0.85},
),
{"frame_number": 47},
"test_detections_varied_data.csv",
@@ -253,19 +300,49 @@ def test_csv_sink(
"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
[
"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
(
mock_detections(
xyxy=[[20, 21, 22, 23]],
),
{"metadata": {"sensor_id": 101, "location": "north"}, "tags": ["urgent", "review"]},
{
"metadata": {"sensor_id": 101, "location": "north"},
"tags": ["urgent", "review"],
},
mock_detections(
xyxy=[[14, 15, 16, 17]],
),
{"metadata": {"sensor_id": 104, "location": "west"}, "tags": ["not-urgent", "done"]},
{
"metadata": {"sensor_id": 104, "location": "west"},
"tags": ["not-urgent", "done"],
},
"test_detections_complex_data.csv",
[
[
@@ -279,13 +356,32 @@ def test_csv_sink(
"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
[
"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: mock_detections,
custom_data: Dict[str, Any],
@@ -302,12 +398,13 @@ def test_csv_sink_manual(
assert_csv_equal(file_name, expected_result)
def assert_csv_equal(file_name, expected_rows):
def assert_csv_equal(file_name, expected_rows):
with open(file_name, 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]}"
#os.remove(file_name)
# os.remove(file_name)