From 5984d3374cd46cded2e7c8455bf9749076fe2bf2 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Fri, 29 Dec 2023 18:24:24 +0100 Subject: [PATCH] better `Detections.merge` validation --- supervision/detection/core.py | 63 +++++++++++++++++++--------------- supervision/detection/utils.py | 53 +++++++++++++++++++++++----- test/detection/test_core.py | 15 +++++--- 3 files changed, 92 insertions(+), 39 deletions(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 8269571a..1ce5599f 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -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, diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index 5cccdb3f..81f5983f 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -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. diff --git a/test/detection/test_core.py b/test/detection/test_core.py index 40ab22a6..84cf8a98 100644 --- a/test/detection/test_core.py +++ b/test/detection/test_core.py @@ -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 ( [