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

better `Detections.merge` validation

This commit is contained in:
SkalskiP 2023-12-29 18:24:24 +01:00
parent 099d36011d
commit 5984d3374c
3 changed files with 92 additions and 39 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

63
supervision/detection/core.py
View File

@ -1,7 +1,7 @@
from __future__ import annotations
from dataclasses import astuple, dataclass, field
from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
from dataclasses import dataclass, field
from typing import Dict, Iterator, List, Optional, Tuple, Union
import numpy as np
@ -12,13 +12,8 @@ from supervision.detection.utils import (
merge_data,
non_max_suppression,
process_roboflow_result,
validate_class_id,
validate_confidence,
validate_data,
validate_mask,
validate_tracker_id,
validate_xyxy,
xywh_to_xyxy,
validate_detections_fields,
)
from supervision.geometry.core import Position
@ -52,13 +47,14 @@ class Detections:
data: Dict[str, Union[np.ndarray, List]] = field(default_factory=dict)
def __post_init__(self):
n = len(self.xyxy)
validate_xyxy(xyxy=self.xyxy, n=n)
validate_mask(mask=self.mask, n=n)
validate_class_id(class_id=self.class_id, n=n)
validate_confidence(confidence=self.confidence, n=n)
validate_tracker_id(tracker_id=self.tracker_id, n=n)
validate_data(data=self.data, n=n)
validate_detections_fields(
xyxy=self.xyxy,
mask=self.mask,
confidence=self.confidence,
class_id=self.class_id,
tracker_id=self.tracker_id,
data=self.data
)
def __len__(self):
"""
@ -682,20 +678,33 @@ class Detections:
if len(detections_list) == 0:
return Detections.empty()
detections_tuples_list = [astuple(detection) for detection in detections_list]
xyxy, mask, confidence, class_id, tracker_id, data = [
list(field_values) for field_values in zip(*detections_tuples_list)
]
for detections in detections_list:
validate_detections_fields(
xyxy=detections.xyxy,
mask=detections.mask,
confidence=detections.confidence,
class_id=detections.class_id,
tracker_id=detections.tracker_id,
data=detections.data
)
def all_not_none(item_list: List[Any]):
return all(x is not None for x in item_list)
xyxy = np.vstack([d.xyxy for d in detections_list])
xyxy = np.vstack(xyxy)
mask = np.vstack(mask) if all_not_none(mask) else None
confidence = np.hstack(confidence) if all_not_none(confidence) else None
class_id = np.hstack(class_id) if all_not_none(class_id) else None
tracker_id = np.hstack(tracker_id) if all_not_none(tracker_id) else None
data = merge_data(data)
def stack_or_none(name: str):
if all(d.__getattribute__(name) is None for d in detections_list):
return None
if any(d.__getattribute__(name) is None for d in detections_list):
raise ValueError(f"All or none of the '{name}' fields must be None")
return np.vstack([d.__getattribute__(name) for d in detections_list]) \
if name == 'mask' else np.hstack(
[d.__getattribute__(name) for d in detections_list])
mask = stack_or_none('mask')
confidence = stack_or_none('confidence')
class_id = stack_or_none('class_id')
tracker_id = stack_or_none('tracker_id')
data = merge_data([d.data for d in detections_list])
return cls(
xyxy=xyxy,

53
supervision/detection/utils.py
View File

@ -472,45 +472,65 @@ def calculate_masks_centroids(masks: np.ndarray) -> np.ndarray:
return np.column_stack((centroid_x, centroid_y)).astype(int)
def validate_xyxy(xyxy: Any, n: int) -> None:
is_valid = isinstance(xyxy, np.ndarray) and xyxy.shape == (n, 4)
def validate_xyxy(xyxy: Any) -> None:
expected_shape = "(_, 4)"
actual_shape = str(getattr(xyxy, 'shape', None))
is_valid = isinstance(xyxy, np.ndarray) and xyxy.ndim == 2 and xyxy.shape[1] == 4
if not is_valid:
raise ValueError("xyxy must be 2d np.ndarray with (n, 4) shape")
raise ValueError(
f"xyxy must be a 2D np.ndarray with shape {expected_shape}, but got shape "
f"{actual_shape}")
def validate_mask(mask: Any, n: int) -> None:
expected_shape = f"({n}, H, W)"
actual_shape = str(getattr(mask, 'shape', None))
is_valid = mask is None or (
isinstance(mask, np.ndarray) and len(mask.shape) == 3 and mask.shape[0] == n
)
if not is_valid:
raise ValueError("mask must be 3d np.ndarray with (n, H, W) shape")
raise ValueError(
f"mask must be a 3D np.ndarray with shape {expected_shape}, but got shape "
f"{actual_shape}")
def validate_class_id(class_id: Any, n: int) -> None:
expected_shape = f"({n},)"
actual_shape = str(getattr(class_id, 'shape', None))
is_valid = class_id is None or (
isinstance(class_id, np.ndarray) and class_id.shape == (n,)
)
if not is_valid:
raise ValueError("class_id must be None or 1d np.ndarray with (n,) shape")
raise ValueError(
f"class_id must be a 1D np.ndarray with shape {expected_shape}, but got "
f"shape {actual_shape}")
def validate_confidence(confidence: Any, n: int) -> None:
expected_shape = f"({n},)"
actual_shape = str(getattr(confidence, 'shape', None))
is_valid = confidence is None or (
isinstance(confidence, np.ndarray) and confidence.shape == (n,)
)
if not is_valid:
raise ValueError("confidence must be None or 1d np.ndarray with (n,) shape")
raise ValueError(
f"confidence must be a 1D np.ndarray with shape {expected_shape}, but got "
f"shape {actual_shape}")
def validate_tracker_id(tracker_id: Any, n: int) -> None:
expected_shape = f"({n},)"
actual_shape = str(getattr(tracker_id, 'shape', None))
is_valid = tracker_id is None or (
isinstance(tracker_id, np.ndarray) and tracker_id.shape == (n,)
)
if not is_valid:
raise ValueError("tracker_id must be None or 1d np.ndarray with (n,) shape")
raise ValueError(
f"tracker_id must be a 1D np.ndarray with shape {expected_shape}, but got "
f"shape {actual_shape}")
def validate_data(data: Dict[str, Union[np.ndarray, List]], n: int) -> None:
def validate_data(data: Dict[str, Any], n: int) -> None:
for key, value in data.items():
if isinstance(value, list):
if len(value) != n:
@ -526,6 +546,23 @@ def validate_data(data: Dict[str, Union[np.ndarray, List]], n: int) -> None:
raise ValueError(f"Value for key '{key}' must be a list or np.ndarray")
def validate_detections_fields(
xyxy: Any,
mask: Any,
class_id: Any,
confidence: Any,
tracker_id: Any,
data: Dict[str, Any]
) -> None:
validate_xyxy(xyxy)
n = len(xyxy)
validate_mask(mask, n)
validate_class_id(class_id, n)
validate_confidence(confidence, n)
validate_tracker_id(tracker_id, n)
validate_data(data, n)
def is_data_equal(data_a: Dict[str, np.ndarray], data_b: Dict[str, np.ndarray]) -> bool:
"""
Compares the data payloads of two Detections instances.

15
test/detection/test_core.py
View File

@ -143,17 +143,24 @@ def test_getitem(
[
([], Detections.empty(), DoesNotRaise()), # empty detections list
(
[Detections.empty()],
[
Detections.empty()
],
Detections.empty(),
DoesNotRaise(),
), # single empty detections
(
[mock_detections(xyxy=[[10, 10, 20, 20]])],
[
mock_detections(xyxy=[[10, 10, 20, 20]])
],
mock_detections(xyxy=[[10, 10, 20, 20]]),
DoesNotRaise(),
), # single detection with xyxy field
(
[mock_detections(xyxy=[[10, 10, 20, 20]]), Detections.empty()],
[
mock_detections(xyxy=[[10, 10, 20, 20]]),
mock_detections(xyxy=np.empty((0, 4), dtype=np.float32)),
],
mock_detections(xyxy=[[10, 10, 20, 20]]),
DoesNotRaise(),
), # single detection with xyxy field + empty detection
@ -171,7 +178,7 @@ def test_getitem(
mock_detections(xyxy=[[20, 20, 30, 30]]),
],
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
DoesNotRaise(),
pytest.raises(ValueError),
), # detection with xyxy, class_id fields + detection with xyxy field
(
[