72 lines
2.2 KiB
Python
72 lines
2.2 KiB
Python
from typing import Dict, List
|
|
|
|
import numpy as np
|
|
|
|
from supervision.detection.core import Detections
|
|
from supervision.dataset.core import DetectionDataset
|
|
|
|
|
|
def mock_detections(
|
|
xyxy: List[List[float]],
|
|
confidence: List[float] = None,
|
|
class_id: List[int] = None,
|
|
tracker_id: List[int] = None,
|
|
) -> Detections:
|
|
return Detections(
|
|
xyxy=np.array(xyxy, dtype=np.float32),
|
|
confidence=confidence
|
|
if confidence is None
|
|
else np.array(confidence, dtype=np.float32),
|
|
class_id=class_id if class_id is None else np.array(class_id, dtype=int),
|
|
tracker_id=tracker_id
|
|
if tracker_id is None
|
|
else np.array(tracker_id, dtype=int),
|
|
)
|
|
|
|
|
|
def mock_detection_dataset(
|
|
images: Dict[str, np.ndarray],
|
|
annotations: Dict[str, Detections],
|
|
classes: List[str],
|
|
) -> DetectionDataset:
|
|
return DetectionDataset(classes=classes, images=images, annotations=annotations)
|
|
|
|
|
|
def dummy_detection_dataset():
|
|
img_paths = ["a.png", "b.png", "c.png"]
|
|
classes = ["a", "b", "c"]
|
|
imgs = [
|
|
np.random.randint(0, 255, size=(28, 28, 3), dtype=np.uint8),
|
|
np.random.randint(0, 255, size=(28, 28, 3), dtype=np.uint8),
|
|
np.random.randint(0, 255, size=(28, 28, 3), dtype=np.uint8),
|
|
]
|
|
detections = [
|
|
mock_detections(
|
|
xyxy=[[10, 10, 20, 20], [20, 20, 25, 25]],
|
|
class_id=[0, 1],
|
|
confidence=np.ones(2),
|
|
),
|
|
mock_detections(
|
|
xyxy=[[10, 10, 20, 20], [20, 20, 25, 25]],
|
|
class_id=[2, 1],
|
|
confidence=np.ones(2),
|
|
),
|
|
mock_detections(
|
|
xyxy=[[10, 10, 20, 20], [20, 20, 25, 25], [10, 10, 15, 15]],
|
|
class_id=[0, 2, 2],
|
|
confidence=np.ones(3),
|
|
),
|
|
]
|
|
annotations = dict(zip(img_paths, detections))
|
|
images = dict(zip(img_paths, imgs))
|
|
dataset = mock_detection_dataset(images, annotations, classes)
|
|
return dataset
|
|
|
|
|
|
def dummy_detection_dataset_with_map_img_to_annotation():
|
|
dataset = dummy_detection_dataset()
|
|
dataset.map_img_to_annotation = lambda img: dataset.annotations[
|
|
[k for k, v in dataset.images.items() if np.array_equal(v, img)][0]
|
|
]
|
|
return dataset
|