allowing to serialise Detections to a JSON file
This commit is contained in:
Adonai Vera
parent
87a4927d03
commit
6e32cdb55a
|
|
@ -1,10 +1,77 @@
|
|||
import json
|
||||
import csv
|
||||
from pathlib import Path
|
||||
from typing import List, Optional, Union
|
||||
from typing import List, Optional, Union, Dict, Any
|
||||
from supervision.detection.core import Detections
|
||||
|
||||
import numpy as np
|
||||
import yaml
|
||||
|
||||
class JSONSink:
|
||||
"""
|
||||
A utility class for saving detection data to a JSON file. This class is designed to
|
||||
efficiently serialize detection objects into a JSON format, allowing for the inclusion of
|
||||
bounding box coordinates and additional attributes like confidence, class ID, and tracker ID.
|
||||
|
||||
The class supports the capability to include custom data alongside the detection fields,
|
||||
providing flexibility for logging various types of information in a structured JSON format.
|
||||
|
||||
Args:
|
||||
filename (str): The name of the JSON file where the detections will be stored.
|
||||
Defaults to 'output.json'.
|
||||
|
||||
Usage:
|
||||
```python
|
||||
from supervision.utils.detections import Detections
|
||||
# Initialize JSONSink with a filename
|
||||
json_sink = JSONSink('my_detections.json')
|
||||
|
||||
# Assuming detections is an instance of Detections containing detection data
|
||||
detections = Detections(...)
|
||||
|
||||
# Open the JSONSink context, append detection data, and close the file automatically
|
||||
with json_sink as sink:
|
||||
sink.append(detections, custom_data={'frame': 1})
|
||||
```
|
||||
"""
|
||||
def __init__(self, filename: str = 'output.json'):
|
||||
self.filename: str = filename
|
||||
self.file: Optional[open] = None
|
||||
self.data: List[Dict[str, Any]] = []
|
||||
|
||||
def __enter__(self) -> 'JSONSink':
|
||||
self.open()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type: Optional[type], exc_val: Optional[Exception], exc_tb: Optional[Any]) -> None:
|
||||
self.write_and_close()
|
||||
|
||||
def open(self) -> None:
|
||||
self.file = open(self.filename, 'w')
|
||||
|
||||
def write_and_close(self) -> None:
|
||||
if self.file:
|
||||
json.dump(self.data, self.file, indent=4)
|
||||
self.file.close()
|
||||
|
||||
def append(self, detections: Detections, custom_data: Dict[str, Any] = None) -> None:
|
||||
for i in range(len(detections.xyxy)):
|
||||
detection_data = {
|
||||
'x_min': int(detections.xyxy[i][0]),
|
||||
'y_min': int(detections.xyxy[i][1]),
|
||||
'x_max': int(detections.xyxy[i][2]),
|
||||
'y_max': int(detections.xyxy[i][3]),
|
||||
'class_id': int(detections.class_id[i]),
|
||||
'confidence': float(detections.confidence[i]),
|
||||
'tracker_id': int(detections.tracker_id[i])
|
||||
}
|
||||
|
||||
for key, value in detections.data.items():
|
||||
detection_data[key] = value[i] if hasattr(value, '__getitem__') else value
|
||||
if custom_data:
|
||||
detection_data.update(custom_data)
|
||||
|
||||
self.data.append(detection_data)
|
||||
|
||||
class NumpyJsonEncoder(json.JSONEncoder):
|
||||
def default(self, obj):
|
||||
|
|
|
|||
|
|
@ -0,0 +1,110 @@
|
|||
import os
|
||||
import csv
|
||||
import pytest
|
||||
import numpy as np
|
||||
import json
|
||||
from supervision.utils.file import JSONSink
|
||||
from supervision.detection.core import Detections
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def detection_instances():
|
||||
# Setup detection instances as per the provided example
|
||||
detections = Detections(
|
||||
xyxy=np.array([[10, 20, 30, 40], [50, 60, 70, 80]]),
|
||||
confidence=np.array([0.7, 0.8]),
|
||||
class_id=np.array([0, 0]),
|
||||
tracker_id=np.array([0, 1]),
|
||||
data={'class_name': ['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': ['car', 'car']}
|
||||
)
|
||||
|
||||
custom_data = {'frame_number': 42}
|
||||
second_custom_data = {'frame_number': 43}
|
||||
|
||||
return detections, custom_data, second_detections, second_custom_data
|
||||
|
||||
def test_json_sink(detection_instances):
|
||||
detections, custom_data, second_detections, second_custom_data = detection_instances
|
||||
json_filename = "test_detections.json"
|
||||
expected_data = [
|
||||
{
|
||||
"x_min": 10, "y_min": 20, "x_max": 30, "y_max": 40,
|
||||
"class_id": 0, "confidence": 0.7, "tracker_id": 0, "class_name": "person",
|
||||
"frame_number": 42
|
||||
},
|
||||
{
|
||||
"x_min": 50, "y_min": 60, "x_max": 70, "y_max": 80,
|
||||
"class_id": 0, "confidence": 0.8, "tracker_id": 1, "class_name": "person",
|
||||
"frame_number": 42
|
||||
},
|
||||
{
|
||||
"x_min": 15, "y_min": 25, "x_max": 35, "y_max": 45,
|
||||
"class_id": 1, "confidence": 0.6, "tracker_id": 2, "class_name": "car",
|
||||
"frame_number": 43
|
||||
},
|
||||
{
|
||||
"x_min": 55, "y_min": 65, "x_max": 75, "y_max": 85,
|
||||
"class_id": 1, "confidence": 0.9, "tracker_id": 3, "class_name": "car",
|
||||
"frame_number": 43
|
||||
}
|
||||
]
|
||||
|
||||
# Using the JSONSink class to write the detection data to a JSON file
|
||||
with JSONSink(filename=json_filename) as sink:
|
||||
sink.append(detections, custom_data)
|
||||
sink.append(second_detections, second_custom_data)
|
||||
|
||||
# Read back the JSON file and verify its contents
|
||||
with open(json_filename, 'r') as file:
|
||||
data = json.load(file)
|
||||
assert data == expected_data, f"Data in JSON file did not match expected output: {data} != {expected_data}"
|
||||
|
||||
# Clean up by removing the test JSON file
|
||||
os.remove(json_filename)
|
||||
|
||||
def test_csv_sink_manual(detection_instances):
|
||||
detections, custom_data, second_detections, second_custom_data = detection_instances
|
||||
json_filename = "test_detections.json"
|
||||
expected_data = [
|
||||
{
|
||||
"x_min": 10, "y_min": 20, "x_max": 30, "y_max": 40,
|
||||
"class_id": 0, "confidence": 0.7, "tracker_id": 0, "class_name": "person",
|
||||
"frame_number": 42
|
||||
},
|
||||
{
|
||||
"x_min": 50, "y_min": 60, "x_max": 70, "y_max": 80,
|
||||
"class_id": 0, "confidence": 0.8, "tracker_id": 1, "class_name": "person",
|
||||
"frame_number": 42
|
||||
},
|
||||
{
|
||||
"x_min": 15, "y_min": 25, "x_max": 35, "y_max": 45,
|
||||
"class_id": 1, "confidence": 0.6, "tracker_id": 2, "class_name": "car",
|
||||
"frame_number": 43
|
||||
},
|
||||
{
|
||||
"x_min": 55, "y_min": 65, "x_max": 75, "y_max": 85,
|
||||
"class_id": 1, "confidence": 0.9, "tracker_id": 3, "class_name": "car",
|
||||
"frame_number": 43
|
||||
}
|
||||
]
|
||||
|
||||
sink = JSONSink(filename=json_filename)
|
||||
sink.open()
|
||||
sink.append(detections, custom_data)
|
||||
sink.append(second_detections, second_custom_data)
|
||||
sink.write_and_close()
|
||||
|
||||
# Read back the JSON file and verify its contents
|
||||
with open(json_filename, 'r') as file:
|
||||
data = json.load(file)
|
||||
assert data == expected_data, f"Data in JSON file did not match expected output: {data} != {expected_data}"
|
||||
|
||||
# Clean up by removing the test JSON file
|
||||
os.remove(json_filename)
|
||||
Loading…
Reference in New Issue