wehub-resource-sync
9194ef5abd
Docs/Test Workflow / Test docs build (push) Failing after 0s
Check links & references / links-check (push) Failing after 1s
Pytest/Test Workflow / Import Test and Pytest Run (ubuntu-latest, 3.10) (push) Failing after 0s
Pytest/Test Workflow / Import Test and Pytest Run (ubuntu-latest, 3.11) (push) Failing after 0s
PR Conflict Labeler / main (push) Failing after 2s
Pytest/Test Workflow / Import Test and Pytest Run (ubuntu-latest, 3.12) (push) Failing after 2s
Pytest/Test Workflow / Import Test and Pytest Run (ubuntu-latest, 3.13) (push) Failing after 0s
Pytest/Test Workflow / Build this Package (push) Failing after 5s
Pytest/Test Workflow / Import Test and Pytest Run (macos-latest, 3.10) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (macos-latest, 3.11) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (macos-latest, 3.12) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (macos-latest, 3.13) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (windows-latest, 3.10) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (windows-latest, 3.11) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (windows-latest, 3.12) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (windows-latest, 3.13) (push) Has been cancelled
Pytest/Test Workflow / testing-guardian (push) Has been cancelled
659 lines
21 KiB
Python
659 lines
21 KiB
Python
"""
|
|
Helper functions and utilities for testing the `supervision` library.
|
|
|
|
This module provides convenient factory functions for creating `Detections`
|
|
and `KeyPoints` objects from simple list-based inputs, as well as utilities
|
|
for generating synthetic test data and performing custom assertions.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import io
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
from supervision.detection.core import Detections
|
|
from supervision.key_points.core import KeyPoints
|
|
|
|
|
|
def make_panoptic_png(segment_map: np.ndarray) -> bytes:
|
|
"""Encode a segment-ID array as a 24-bit RGBA PNG byte string.
|
|
|
|
Segment IDs are stored in RGB little-endian order so tests can cover
|
|
panoptic labels above 255 without collisions.
|
|
"""
|
|
segment_map_u32 = np.asarray(segment_map, dtype=np.uint32)
|
|
arr = np.zeros((*segment_map_u32.shape, 4), dtype=np.uint8)
|
|
arr[:, :, 0] = (segment_map_u32 & 0xFF).astype(np.uint8)
|
|
arr[:, :, 1] = ((segment_map_u32 >> 8) & 0xFF).astype(np.uint8)
|
|
arr[:, :, 2] = ((segment_map_u32 >> 16) & 0xFF).astype(np.uint8)
|
|
arr[:, :, 3] = 255
|
|
buf = io.BytesIO()
|
|
Image.fromarray(arr).save(buf, format="PNG")
|
|
return buf.getvalue()
|
|
|
|
|
|
def _create_detections(
|
|
xyxy: list[list[float]],
|
|
mask: list[np.ndarray] | None = None,
|
|
confidence: list[float] | None = None,
|
|
class_id: list[int] | None = None,
|
|
tracker_id: list[int] | None = None,
|
|
data: dict[str, list[Any]] | None = None,
|
|
) -> Detections:
|
|
"""
|
|
Create a Detections object from list-based inputs.
|
|
|
|
This is a helper function primarily used for testing purposes to quickly
|
|
instantiate a Detections object without manually converting lists to numpy arrays.
|
|
|
|
Args:
|
|
xyxy: Bounding boxes in `(x_min, y_min, x_max, y_max)`
|
|
format.
|
|
mask: Binary masks for each detection.
|
|
confidence: Confidence scores for each detection.
|
|
class_id: Class identifiers for each detection.
|
|
tracker_id: Tracker identifiers for each detection.
|
|
data: Additional data to be associated with
|
|
each detection.
|
|
|
|
Returns:
|
|
A Detections object containing the provided data.
|
|
|
|
Examples:
|
|
>>> import numpy as np
|
|
>>> detections = _create_detections(
|
|
... xyxy=[[0, 0, 10, 10], [20, 20, 30, 30]],
|
|
... confidence=[0.5, 0.8],
|
|
... class_id=[0, 1]
|
|
... )
|
|
>>> detections.xyxy
|
|
array([[ 0., 0., 10., 10.],
|
|
[20., 20., 30., 30.]], dtype=float32)
|
|
>>> detections.confidence
|
|
array([0.5, 0.8], dtype=float32)
|
|
>>> detections.class_id
|
|
array([0, 1])
|
|
"""
|
|
|
|
def convert_data(data: dict[str, list[Any]]):
|
|
return {k: np.array(v) for k, v in data.items()}
|
|
|
|
return Detections(
|
|
xyxy=np.array(xyxy, dtype=np.float32),
|
|
mask=(mask if mask is None else np.array(mask, dtype=bool)),
|
|
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)
|
|
),
|
|
data=convert_data(data) if data else {},
|
|
)
|
|
|
|
|
|
def _create_key_points(
|
|
xy: list[list[list[float]]],
|
|
confidence: list[list[float]] | None = None,
|
|
class_id: list[int] | None = None,
|
|
detection_confidence: list[float] | None = None,
|
|
visible: list[list[bool]] | None = None,
|
|
data: dict[str, list[Any]] | None = None,
|
|
) -> KeyPoints:
|
|
"""
|
|
Create a KeyPoints object from list-based inputs.
|
|
|
|
This is a helper function primarily used for testing purposes to quickly
|
|
instantiate a KeyPoints object without manually converting lists to numpy arrays.
|
|
|
|
Args:
|
|
xy: Keypoint coordinates in `(x, y)` format for
|
|
each detection.
|
|
confidence: Per-keypoint confidence scores.
|
|
class_id: Class identifiers for each keypoint set.
|
|
detection_confidence: Detection-level confidence scores.
|
|
data: Additional data to be associated with
|
|
each keypoint set.
|
|
|
|
Returns:
|
|
A KeyPoints object containing the provided data.
|
|
|
|
Examples:
|
|
>>> import numpy as np
|
|
>>> key_points = _create_key_points(
|
|
... xy=[[[0, 0], [10, 10]], [[20, 20], [30, 30]]],
|
|
... confidence=[[0.5, 0.8], [0.9, 0.1]],
|
|
... class_id=[0, 1]
|
|
... )
|
|
>>> key_points.xy
|
|
array([[[ 0., 0.],
|
|
[10., 10.]],
|
|
<BLANKLINE>
|
|
[[20., 20.],
|
|
[30., 30.]]], dtype=float32)
|
|
>>> key_points.keypoint_confidence
|
|
array([[0.5, 0.8],
|
|
[0.9, 0.1]], dtype=float32)
|
|
>>> key_points.class_id
|
|
array([0, 1])
|
|
"""
|
|
|
|
def convert_data(data: dict[str, list[Any]]):
|
|
return {k: np.array(v) for k, v in data.items()}
|
|
|
|
return KeyPoints(
|
|
xy=np.array(xy, dtype=np.float32),
|
|
keypoint_confidence=(
|
|
confidence if confidence is None else np.array(confidence, dtype=np.float32)
|
|
),
|
|
detection_confidence=(
|
|
detection_confidence
|
|
if detection_confidence is None
|
|
else np.array(detection_confidence, dtype=np.float32)
|
|
),
|
|
visible=(visible if visible is None else np.array(visible, dtype=bool)),
|
|
class_id=(class_id if class_id is None else np.array(class_id, dtype=int)),
|
|
data=convert_data(data) if data else {},
|
|
)
|
|
|
|
|
|
def _generate_random_boxes(
|
|
count: int,
|
|
image_size: tuple[int, int] = (1920, 1080),
|
|
min_box_size: int = 20,
|
|
max_box_size: int = 200,
|
|
seed: int | None = None,
|
|
) -> np.ndarray:
|
|
"""
|
|
Generate random bounding boxes within given image dimensions and size constraints.
|
|
|
|
Creates `count` bounding boxes randomly positioned and sized, ensuring each
|
|
stays within image bounds and has width and height in the specified range.
|
|
|
|
Args:
|
|
count: Number of random bounding boxes to generate.
|
|
image_size: Image size as `(width, height)`.
|
|
min_box_size: Minimum side length (pixels) for generated boxes.
|
|
max_box_size: Maximum side length (pixels) for generated boxes.
|
|
seed: Optional random seed for reproducibility.
|
|
|
|
Returns:
|
|
Array of shape `(count, 4)` with bounding boxes as
|
|
`(x_min, y_min, x_max, y_max)`.
|
|
|
|
Examples:
|
|
>>> boxes = _generate_random_boxes(
|
|
... count=2, image_size=(1000, 1000),
|
|
... min_box_size=10, max_box_size=20, seed=42)
|
|
>>> boxes.shape
|
|
(2, 4)
|
|
>>> boxes
|
|
array([[843.36676, 687.33374, 861.1063 , 701.72253],
|
|
[752.81146, 770.53467, 763.75323, 790.2909 ]], dtype=float32)
|
|
"""
|
|
rng = np.random.default_rng(seed)
|
|
|
|
img_w, img_h = image_size
|
|
out = np.zeros((count, 4), dtype=np.float32)
|
|
|
|
for i in range(count):
|
|
w = rng.uniform(min_box_size, max_box_size)
|
|
h = rng.uniform(min_box_size, max_box_size)
|
|
|
|
x_min = rng.uniform(0, img_w - w)
|
|
y_min = rng.uniform(0, img_h - h)
|
|
x_max = x_min + w
|
|
y_max = y_min + h
|
|
|
|
out[i] = (x_min, y_min, x_max, y_max)
|
|
|
|
return out
|
|
|
|
|
|
def assert_almost_equal(actual, expected, tolerance=1e-5) -> None:
|
|
"""
|
|
Assert that two values are equal within a specified tolerance.
|
|
|
|
Args:
|
|
actual: The value to check.
|
|
expected: The expected value.
|
|
tolerance: The maximum allowed difference between `actual`
|
|
and `expected`.
|
|
|
|
Examples:
|
|
>>> assert_almost_equal(0.500001, 0.5)
|
|
>>> assert_almost_equal(0.6, 0.5, tolerance=0.2)
|
|
>>> assert_almost_equal(0.6, 0.5)
|
|
Traceback (most recent call last):
|
|
...
|
|
AssertionError: Expected 0.5, but got 0.6.
|
|
"""
|
|
assert abs(actual - expected) < tolerance, f"Expected {expected}, but got {actual}."
|
|
|
|
|
|
def assert_image_mostly_same(
|
|
original: np.ndarray, annotated: np.ndarray, similarity_threshold: float = 0.9
|
|
) -> None:
|
|
"""
|
|
Assert that the annotated image is mostly the same as the original.
|
|
|
|
Args:
|
|
original: Original image
|
|
annotated: Annotated image
|
|
similarity_threshold:
|
|
Minimum percentage of pixels that should be the same (0.0 to 1.0)
|
|
"""
|
|
# Check that images have the same shape
|
|
assert original.shape == annotated.shape
|
|
|
|
# Calculate number of identical pixels
|
|
identical_pixels = np.sum(np.all(original == annotated, axis=-1))
|
|
total_pixels = original.shape[0] * original.shape[1]
|
|
similarity = identical_pixels / total_pixels
|
|
|
|
# Check that at least similarity_threshold of pixels are identical
|
|
assert similarity >= similarity_threshold, (
|
|
f"Images are only {similarity:.1%} similar, "
|
|
f"which is below the {similarity_threshold:.1%} threshold"
|
|
)
|
|
|
|
# Check that the image is not completely identical
|
|
assert not np.array_equal(original, annotated), "Images are completely identical"
|
|
|
|
|
|
class _FakeTensor:
|
|
"""Minimal tensor wrapper for cpu().numpy() and int()."""
|
|
|
|
def __init__(self, arr: np.ndarray) -> None:
|
|
self._arr = np.asarray(arr)
|
|
|
|
def cpu(self) -> _FakeTensor:
|
|
return self
|
|
|
|
def numpy(self) -> np.ndarray:
|
|
return self._arr
|
|
|
|
def int(self) -> _FakeTensor:
|
|
return _FakeTensor(self._arr.astype(int))
|
|
|
|
|
|
class _FakeYOLOv5Results:
|
|
"""YOLOv5-like results exposing pred list."""
|
|
|
|
def __init__(self, pred0: np.ndarray) -> None:
|
|
self.pred = [_FakeTensor(pred0)]
|
|
|
|
|
|
class _FakeUltralyticsBoxes:
|
|
"""Ultralytics-like Boxes exposing xyxy/conf/cls and optional id."""
|
|
|
|
def __init__(
|
|
self,
|
|
xyxy: np.ndarray,
|
|
conf: np.ndarray,
|
|
cls: np.ndarray,
|
|
id_: np.ndarray | None = None,
|
|
) -> None:
|
|
self.xyxy = _FakeTensor(xyxy)
|
|
self.conf = _FakeTensor(conf)
|
|
self.cls = _FakeTensor(cls)
|
|
self.id = _FakeTensor(id_) if id_ is not None else None
|
|
|
|
|
|
class _FakeUltralyticsResults:
|
|
"""Ultralytics-like results container used by from_ultralytics."""
|
|
|
|
def __init__(self, boxes, names: dict[int, str], length: int = 0) -> None:
|
|
self.boxes = boxes
|
|
self.names = names
|
|
self.obb = None
|
|
self.masks = None
|
|
self._length = length
|
|
|
|
def __len__(self) -> int:
|
|
return self._length
|
|
|
|
|
|
class _FakeYoloNasPrediction:
|
|
"""YOLO-NAS-like prediction struct."""
|
|
|
|
def __init__(self, bboxes_xyxy, confidence, labels) -> None:
|
|
self.bboxes_xyxy = bboxes_xyxy
|
|
self.confidence = confidence
|
|
self.labels = labels
|
|
|
|
|
|
class _FakeYoloNasResults:
|
|
"""YOLO-NAS-like results exposing prediction."""
|
|
|
|
def __init__(self, prediction: _FakeYoloNasPrediction) -> None:
|
|
self.prediction = prediction
|
|
|
|
|
|
class _FakeYoloNasKeyPoint:
|
|
"""YOLO-NAS-like key point struct."""
|
|
|
|
def __init__(self, poses, labels=None) -> None:
|
|
self.poses = np.array(poses, dtype=np.float32)
|
|
if labels is not None:
|
|
self.labels = np.array(labels, dtype=int)
|
|
|
|
|
|
class _FakeYoloNasKeyPointResults:
|
|
"""YOLO-NAS-like results exposing key points."""
|
|
|
|
def __init__(self, prediction: _FakeYoloNasKeyPoint, class_names=None) -> None:
|
|
self.prediction = prediction
|
|
self.class_names = class_names
|
|
|
|
|
|
class _FakeMediapipeLandmark:
|
|
def __init__(self, x, y, visibility=1.0) -> None:
|
|
self.x = x
|
|
self.y = y
|
|
self.visibility = visibility
|
|
|
|
|
|
class _FakeMediapipePose:
|
|
def __init__(self, landmarks: list[_FakeMediapipeLandmark]) -> None:
|
|
self.landmark = landmarks
|
|
|
|
|
|
class _FakeMediapipeResults:
|
|
def __init__(
|
|
self,
|
|
pose_landmarks: list[list[_FakeMediapipeLandmark]]
|
|
| _FakeMediapipePose
|
|
| None = None,
|
|
face_landmarks: _FakeMediapipeLandmark | None = None,
|
|
multi_face_landmarks: list[_FakeMediapipeLandmark] | None = None,
|
|
) -> None:
|
|
self.pose_landmarks = pose_landmarks
|
|
self.face_landmarks = face_landmarks
|
|
self.multi_face_landmarks = multi_face_landmarks
|
|
|
|
|
|
def create_yolo_dataset(
|
|
dataset_dir: str,
|
|
num_images: int = 15,
|
|
image_size: tuple[int, int, int] = (640, 640, 3),
|
|
classes: list[str] | None = None,
|
|
objects_per_image_range: tuple[int, int] = (2, 4),
|
|
seed: int = 42,
|
|
) -> dict[str, Any]:
|
|
"""
|
|
Create a synthetic YOLO-format dataset on disk.
|
|
|
|
Generates dummy images with YOLO-format annotations, `data.yaml` file,
|
|
and directory structure suitable for testing dataset loading.
|
|
|
|
Args:
|
|
dataset_dir: Root directory path for the dataset.
|
|
num_images: Number of images to generate.
|
|
image_size: Image dimensions as `(width, height, channels)`.
|
|
classes: List of class names. Defaults to `["class_0", "class_1"]`.
|
|
objects_per_image_range: Range of objects per image as `(min, max)`.
|
|
Actual count will cycle through this range.
|
|
seed: Random seed for reproducibility.
|
|
|
|
Returns:
|
|
Dictionary containing:
|
|
- `tmpdir`: Root dataset directory path
|
|
- `images_dir`: Images directory path
|
|
- `labels_dir`: Labels directory path
|
|
- `data_yaml_path`: `data.yaml` file path
|
|
- `num_images`: Number of images created
|
|
- `image_size`: Image dimensions
|
|
- `image_annotations`: List of annotations per image
|
|
|
|
Examples:
|
|
>>> from pathlib import Path
|
|
>>> import tempfile
|
|
>>> tmpdir = Path(tempfile.mkdtemp())
|
|
>>> dataset_info = create_yolo_dataset(str(tmpdir), num_images=5)
|
|
>>> dataset_info["num_images"]
|
|
5
|
|
>>> len(list(Path(dataset_info["images_dir"]).glob("*.jpg")))
|
|
5
|
|
"""
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
|
|
if classes is None:
|
|
classes = ["class_0", "class_1"]
|
|
|
|
np.random.seed(seed)
|
|
|
|
dataset_path = Path(dataset_dir)
|
|
images_dir = dataset_path / "images"
|
|
labels_dir = dataset_path / "labels"
|
|
images_dir.mkdir(parents=True, exist_ok=True)
|
|
labels_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
min_objects, max_objects = objects_per_image_range
|
|
num_classes = len(classes)
|
|
image_annotations = []
|
|
|
|
for i in range(num_images):
|
|
# Create dummy image
|
|
img_path = images_dir / f"image_{i:03d}.jpg"
|
|
img = np.zeros(image_size, dtype=np.uint8)
|
|
cv2.imwrite(str(img_path), img)
|
|
|
|
# Determine number of objects for this image
|
|
num_objects = min_objects + (i % (max_objects - min_objects + 1))
|
|
yolo_lines = []
|
|
objects = []
|
|
|
|
for j in range(num_objects):
|
|
class_id = j % num_classes
|
|
# Random positions with spacing to avoid overlap
|
|
x_center = 0.15 + (j * 0.25) + np.random.uniform(-0.05, 0.05)
|
|
y_center = 0.15 + (j * 0.2) + np.random.uniform(-0.05, 0.05)
|
|
width = 0.12
|
|
height = 0.12
|
|
|
|
# Clip to valid range [0, 1]
|
|
x_center = np.clip(x_center, width / 2, 1 - width / 2)
|
|
y_center = np.clip(y_center, height / 2, 1 - height / 2)
|
|
|
|
yolo_lines.append(
|
|
f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n"
|
|
)
|
|
objects.append((class_id, x_center, y_center, width, height))
|
|
|
|
# Write YOLO annotation file
|
|
label_path = labels_dir / f"image_{i:03d}.txt"
|
|
label_path.write_text("".join(yolo_lines))
|
|
image_annotations.append(objects)
|
|
|
|
# Create data.yaml
|
|
data_yaml_path = dataset_path / "data.yaml"
|
|
yaml_content = "names:\n" + "\n".join(f"- {cls}" for cls in classes) + "\n"
|
|
data_yaml_path.write_text(yaml_content)
|
|
|
|
return {
|
|
"tmpdir": dataset_path,
|
|
"images_dir": str(images_dir),
|
|
"labels_dir": str(labels_dir),
|
|
"data_yaml_path": str(data_yaml_path),
|
|
"num_images": num_images,
|
|
"image_size": image_size,
|
|
"image_annotations": image_annotations,
|
|
}
|
|
|
|
|
|
class _FakeDetachTensor:
|
|
"""Fake torch.Tensor supporting the cpu().detach().numpy() call chain."""
|
|
|
|
def __init__(self, arr: np.ndarray) -> None:
|
|
self._arr = np.asarray(arr)
|
|
|
|
def cpu(self) -> _FakeDetachTensor:
|
|
"""Return self to allow chaining."""
|
|
return self
|
|
|
|
def detach(self) -> _FakeDetachTensor:
|
|
"""Return self to allow chaining."""
|
|
return self
|
|
|
|
def numpy(self) -> np.ndarray:
|
|
"""Return underlying array."""
|
|
return self._arr
|
|
|
|
|
|
class _FakeDetectron2Boxes:
|
|
"""Fake Detectron2 Boxes exposing .tensor for the cpu().numpy() chain."""
|
|
|
|
def __init__(self, xyxy: np.ndarray) -> None:
|
|
self.tensor = _FakeTensor(xyxy)
|
|
|
|
|
|
class _FakeDetectron2Instances:
|
|
"""Fake Detectron2 Instances: pred_boxes, scores, pred_classes, optional masks."""
|
|
|
|
def __init__(
|
|
self,
|
|
xyxy: np.ndarray,
|
|
scores: np.ndarray,
|
|
class_ids: np.ndarray,
|
|
masks: np.ndarray | None = None,
|
|
) -> None:
|
|
self.pred_boxes = _FakeDetectron2Boxes(xyxy)
|
|
self.scores = _FakeTensor(scores)
|
|
self.pred_classes = _FakeTensor(class_ids)
|
|
if masks is not None:
|
|
self.pred_masks = _FakeTensor(masks)
|
|
|
|
|
|
class _FakeMMDetPredInstances:
|
|
"""Fake MMDetection pred_instances supporting the 'masks' in membership check."""
|
|
|
|
def __init__(
|
|
self,
|
|
xyxy: np.ndarray,
|
|
scores: np.ndarray,
|
|
labels: np.ndarray,
|
|
masks: np.ndarray | None = None,
|
|
) -> None:
|
|
self.bboxes = _FakeTensor(xyxy)
|
|
self.scores = _FakeTensor(scores)
|
|
self.labels = _FakeTensor(labels)
|
|
self._masks: np.ndarray | None = masks
|
|
if masks is not None:
|
|
self.masks = _FakeTensor(masks)
|
|
|
|
def __contains__(self, key: str) -> bool:
|
|
"""Return True for 'masks' only when masks were provided at construction."""
|
|
return key == "masks" and self._masks is not None
|
|
|
|
|
|
class _FakeMMDetResults:
|
|
"""Fake MMDetection inference result wrapping pred_instances."""
|
|
|
|
def __init__(self, pred_instances: _FakeMMDetPredInstances) -> None:
|
|
self.pred_instances = pred_instances
|
|
|
|
|
|
class _FakeDeepSparseResults:
|
|
"""Fake DeepSparse inference result with list attributes boxes, scores, labels."""
|
|
|
|
def __init__(
|
|
self,
|
|
boxes: list[np.ndarray],
|
|
scores: list[np.ndarray],
|
|
labels: list[np.ndarray],
|
|
) -> None:
|
|
self.boxes = boxes
|
|
self.scores = scores
|
|
self.labels = labels
|
|
|
|
|
|
class _FakeNCNNRect:
|
|
"""Fake ncnn Rect with x, y, w, h as numpy float32 scalars supporting .astype()."""
|
|
|
|
def __init__(self, x: float, y: float, w: float, h: float) -> None:
|
|
self.x = np.float32(x)
|
|
self.y = np.float32(y)
|
|
self.w = np.float32(w)
|
|
self.h = np.float32(h)
|
|
|
|
|
|
class _FakeNCNNObject:
|
|
"""Fake ncnn detected object with rect, prob, label."""
|
|
|
|
def __init__(
|
|
self, x: float, y: float, w: float, h: float, prob: float, label: int
|
|
) -> None:
|
|
self.rect = _FakeNCNNRect(x, y, w, h)
|
|
self.prob = prob
|
|
self.label = label
|
|
|
|
|
|
def create_predictions_with_class_iou_tests(
|
|
gt_detections: Detections, num_classes: int
|
|
) -> Detections:
|
|
"""
|
|
Create predictions that test IoU+class matching behavior.
|
|
|
|
For each ground truth detection, creates predictions with different patterns:
|
|
- Pattern 0 (i%3==0): Correct match (same bbox, correct class)
|
|
- Pattern 1 (i%3==1): Wrong class with perfect IoU + correct class with offset
|
|
- Pattern 2 (i%3==2): Correct class with slight offset
|
|
|
|
This tests that predictions with wrong class don't match even with high IoU,
|
|
which is the key fix in the confusion matrix calculation.
|
|
|
|
Args:
|
|
gt_detections: Ground truth detections to create predictions for
|
|
num_classes: Total number of classes in the dataset
|
|
|
|
Returns:
|
|
Detections object with predictions designed to test IoU+class matching
|
|
"""
|
|
if len(gt_detections) == 0:
|
|
# No ground truth, return a single false positive
|
|
return _create_detections(
|
|
xyxy=[[10, 10, 50, 50]], class_id=[0], confidence=[0.9]
|
|
)
|
|
|
|
pred_boxes = []
|
|
pred_classes = []
|
|
pred_confs = []
|
|
|
|
for i, (box, cls) in enumerate(zip(gt_detections.xyxy, gt_detections.class_id)):
|
|
if i % 3 == 0:
|
|
# Pattern 1: Correct match
|
|
pred_boxes.append(box)
|
|
pred_classes.append(cls)
|
|
pred_confs.append(0.95)
|
|
|
|
elif i % 3 == 1:
|
|
# Pattern 2: Test the fix - add wrong class prediction with perfect IoU,
|
|
# then correct class with slightly offset bbox
|
|
wrong_cls = (cls + 1) % num_classes
|
|
pred_boxes.append(box) # Perfect IoU
|
|
pred_classes.append(wrong_cls) # Wrong class
|
|
pred_confs.append(0.90)
|
|
|
|
# Add correct class with slight offset
|
|
offset_box = box + np.array([2, 2, 2, 2], dtype=np.float32)
|
|
pred_boxes.append(offset_box)
|
|
pred_classes.append(cls) # Correct class
|
|
pred_confs.append(0.85)
|
|
|
|
else:
|
|
# Pattern 3: Correct match with slight offset
|
|
offset_box = box + np.array([1, 1, 1, 1], dtype=np.float32)
|
|
pred_boxes.append(offset_box)
|
|
pred_classes.append(cls)
|
|
pred_confs.append(0.92)
|
|
|
|
return _create_detections(
|
|
xyxy=pred_boxes, class_id=pred_classes, confidence=pred_confs
|
|
)
|