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

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

This commit is contained in:
pre-commit-ci[bot] 2024-01-31 16:41:56 +00:00
parent 12568336ba
commit fd00bd4ba2
4 changed files with 69 additions and 33 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
supervision/__init__.py
View File

@ -36,10 +36,10 @@ from supervision.dataset.core import (
from supervision.detection.annotate import BoxAnnotator
from supervision.detection.core import Detections
from supervision.detection.line_counter import LineZone, LineZoneAnnotator
from supervision.detection.tools.csv_sink import CSVSink
from supervision.detection.tools.inference_slicer import InferenceSlicer
from supervision.detection.tools.polygon_zone import PolygonZone, PolygonZoneAnnotator
from supervision.detection.tools.smoother import DetectionsSmoother
from supervision.detection.tools.csv_sink import CSVSink
from supervision.detection.utils import (
box_iou_batch,
calculate_masks_centroids,

33
supervision/detection/tools/csv_sink.py
View File

@ -1,10 +1,8 @@
from __future__ import annotations
import csv
from pathlib import Path
from typing import Any, Dict, List, Optional, Union
from typing import Any, Dict, List, Optional
import numpy as np
from supervision.detection.core import Detections
@ -18,6 +16,7 @@ BASE_HEADER = [
"tracker_id",
]
class CSVSink:
"""
A utility class for saving detection data to a CSV file. This class is designed to
@ -71,8 +70,8 @@ class CSVSink:
callback=callback
)
csv_sink.close()
```
""" # noqa: E501 // docs
```
""" # noqa: E501 // docs
def __init__(self, filename: str = "output.csv"):
self.filename = filename
@ -102,7 +101,9 @@ class CSVSink:
self.file.close()
@staticmethod
def parse_detection_data(detections: Detections, custom_data: Dict[str, Any] = None) -> List[Dict[str, Any]]:
def parse_detection_data(
detections: Detections, custom_data: Dict[str, Any] = None
) -> List[Dict[str, Any]]:
parsed_rows = []
for i in range(len(detections.xyxy)):
row = {
@ -122,20 +123,26 @@ class CSVSink:
parsed_rows.append(row)
return parsed_rows
def append(self, detections: Detections, custom_data: Dict[str, Any] = None) -> None:
def append(
self, detections: Detections, custom_data: Dict[str, Any] = None
) -> None:
if not self.writer:
raise Exception(f"Cannot append to CSV: The file '{self.filename}' is not open.")
raise Exception(
f"Cannot append to CSV: The file '{self.filename}' is not open."
)
if not self.header_written:
self.write_header(detections, custom_data)
parsed_rows = CSVSink.parse_detection_data(detections, custom_data)
for row in parsed_rows:
self.writer.writerow([row.get(fieldname, "") for fieldname in self.fieldnames])
self.writer.writerow(
[row.get(fieldname, "") for fieldname in self.fieldnames]
)
def write_header(
self, detections: Detections, custom_data: Dict[str, Any]
) -> None:
dynamic_header = sorted(set(custom_data.keys()) | set(getattr(detections, "data", {}).keys()))
def write_header(self, detections: Detections, custom_data: Dict[str, Any]) -> None:
dynamic_header = sorted(
set(custom_data.keys()) | set(getattr(detections, "data", {}).keys())
)
self.fieldnames = BASE_HEADER + dynamic_header
self.writer.writerow(self.fieldnames)
self.header_written = True

2
supervision/utils/file.py
View File

@ -146,4 +146,4 @@ def save_yaml_file(data: dict, file_path: str) -> None:
"""
with open(file_path, "w") as outfile:
yaml.dump(data, outfile, sort_keys=False, default_flow_style=None)
yaml.dump(data, outfile, sort_keys=False, default_flow_style=None)

65
test/detection/test_csv.py
View File

@ -1,9 +1,12 @@
import os
import csv
import pytest
import os
import numpy as np
from supervision.detection.core import Detections
import pytest
import supervision as sv
from supervision.detection.core import Detections
@pytest.fixture(scope="module")
def detection_instances():
@ -13,31 +16,42 @@ def detection_instances():
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'])}
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'])}
data={"class_name": np.array(["car", "car"])},
)
custom_data = {'frame_number': 42}
second_custom_data = {'frame_number': 43}
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"],
[
"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]
[55, 65, 75, 85, 1, 0.9, 3, "car", 43],
]
# Using the CSVSink class to write the detection data to a CSV file
@ -46,23 +60,36 @@ def test_csv_sink(detection_instances):
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:
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]}"
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"],
[
"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]
[55, 65, 75, 85, 1, 0.9, 3, "car", 43],
]
# Using the CSVSink class to write the detection data to a CSV file
@ -73,10 +100,12 @@ def test_csv_sink_manual(detection_instances):
sink.close()
# Read back the CSV file and verify its contents
with open(csv_filename, mode='r', newline='') as file:
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]}"
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)