Detection module¶
core
¶
Detections
dataclass
¶
A class for storing and manipulating detection results in a standardized format.
Source code in supertracker/detection/core.py
@dataclass
class Detections:
"""
A class for storing and manipulating detection results in a standardized format.
"""
xyxy: np.ndarray
mask: Optional[np.ndarray] = None
confidence: Optional[np.ndarray] = None
class_id: Optional[np.ndarray] = None
tracker_id: Optional[np.ndarray] = None
data: Dict[str, Union[np.ndarray, List]] = field(default_factory=dict)
metadata: Dict[str, Any] = field(default_factory=dict)
def __post_init__(self):
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):
"""
Returns the number of detections in the Detections object.
"""
return len(self.xyxy)
def __iter__(
self,
) -> Iterator[
Tuple[
np.ndarray,
Optional[np.ndarray],
Optional[float],
Optional[int],
Optional[int],
Dict[str, Union[np.ndarray, List]],
]
]:
"""
Iterates over the Detections object.
"""
for i in range(len(self.xyxy)):
yield (
self.xyxy[i],
self.mask[i] if self.mask is not None else None,
self.confidence[i] if self.confidence is not None else None,
self.class_id[i] if self.class_id is not None else None,
self.tracker_id[i] if self.tracker_id is not None else None,
get_data_item(self.data, i),
)
def __eq__(self, other: Detections):
return all(
[
np.array_equal(self.xyxy, other.xyxy),
np.array_equal(self.mask, other.mask),
np.array_equal(self.class_id, other.class_id),
np.array_equal(self.confidence, other.confidence),
np.array_equal(self.tracker_id, other.tracker_id),
is_data_equal(self.data, other.data),
is_metadata_equal(self.metadata, other.metadata),
]
)
@classmethod
def from_yolo(cls, custom_results) -> Detections:
"""
Creates a `sv.Detections` instance from a custom Results class that mimics
the ultralytics Results format.
Args:
custom_results: The output Results instance from custom detector
Returns:
Detections: A new Detections object.
"""
# Handle OBB case first if available
if hasattr(custom_results, "obb") and custom_results.obb is not None:
class_id = custom_results.obb.cls.astype(int)
class_names = np.array([custom_results.names[i] for i in class_id])
oriented_box_coordinates = custom_results.obb.xyxyxyxy
return cls(
xyxy=custom_results.obb.xyxy,
confidence=custom_results.obb.conf,
class_id=class_id,
tracker_id=custom_results.obb.id.astype(int) if custom_results.obb.id is not None else None,
data={
ORIENTED_BOX_COORDINATES: oriented_box_coordinates,
CLASS_NAME_DATA_FIELD: class_names,
},
)
# Handle regular detection case
if hasattr(custom_results, "boxes") and custom_results.boxes is not None:
class_id = custom_results.boxes.cls.astype(int)
class_names = np.array([custom_results.names[i] for i in class_id])
return cls(
xyxy=custom_results.boxes.xyxy,
confidence=custom_results.boxes.conf,
class_id=class_id,
mask=custom_results.masks.data if hasattr(custom_results, "masks") and custom_results.masks is not None else None,
tracker_id=custom_results.boxes.id.astype(int) if hasattr(custom_results.boxes, "id") and custom_results.boxes.id is not None else None,
data={CLASS_NAME_DATA_FIELD: class_names},
)
# Handle case with only masks
if hasattr(custom_results, "masks") and custom_results.masks is not None:
masks = custom_results.masks.data
return cls(
xyxy=mask_to_xyxy(masks),
mask=masks,
class_id=np.arange(len(custom_results)),
)
# Return empty detections if no supported format is found
return cls.empty()
@classmethod
def empty(cls) -> Detections:
"""
Create an empty Detections object with no bounding boxes,
confidences, or class IDs.
Returns:
(Detections): An empty Detections object.
"""
return cls(
xyxy=np.empty((0, 4), dtype=np.float32),
confidence=np.array([], dtype=np.float32),
class_id=np.array([], dtype=int),
)
def is_empty(self) -> bool:
"""
Returns `True` if the `Detections` object is considered empty.
"""
empty_detections = Detections.empty()
empty_detections.data = self.data
empty_detections.metadata = self.metadata
return self == empty_detections
def __getitem__(
self, index: Union[int, slice, List[int], np.ndarray, str]
) -> Union[Detections, List, np.ndarray, None]:
"""
Get a subset of the Detections object or access an item from its data field.
"""
if isinstance(index, str):
return self.data.get(index)
if self.is_empty():
return self
if isinstance(index, int):
index = [index]
return Detections(
xyxy=self.xyxy[index],
mask=self.mask[index] if self.mask is not None else None,
confidence=self.confidence[index] if self.confidence is not None else None,
class_id=self.class_id[index] if self.class_id is not None else None,
tracker_id=self.tracker_id[index] if self.tracker_id is not None else None,
data=get_data_item(self.data, index),
metadata=self.metadata,
)
def __setitem__(self, key: str, value: Union[np.ndarray, List]):
"""
Set a value in the data dictionary of the Detections object.
"""
if not isinstance(value, (np.ndarray, list)):
raise TypeError("Value must be a np.ndarray or a list")
if isinstance(value, list):
value = np.array(value)
self.data[key] = value
__getitem__(self, index)
special
¶
Get a subset of the Detections object or access an item from its data field.
Source code in supertracker/detection/core.py
def __getitem__(
self, index: Union[int, slice, List[int], np.ndarray, str]
) -> Union[Detections, List, np.ndarray, None]:
"""
Get a subset of the Detections object or access an item from its data field.
"""
if isinstance(index, str):
return self.data.get(index)
if self.is_empty():
return self
if isinstance(index, int):
index = [index]
return Detections(
xyxy=self.xyxy[index],
mask=self.mask[index] if self.mask is not None else None,
confidence=self.confidence[index] if self.confidence is not None else None,
class_id=self.class_id[index] if self.class_id is not None else None,
tracker_id=self.tracker_id[index] if self.tracker_id is not None else None,
data=get_data_item(self.data, index),
metadata=self.metadata,
)
__iter__(self)
special
¶
Iterates over the Detections object.
Source code in supertracker/detection/core.py
def __iter__(
self,
) -> Iterator[
Tuple[
np.ndarray,
Optional[np.ndarray],
Optional[float],
Optional[int],
Optional[int],
Dict[str, Union[np.ndarray, List]],
]
]:
"""
Iterates over the Detections object.
"""
for i in range(len(self.xyxy)):
yield (
self.xyxy[i],
self.mask[i] if self.mask is not None else None,
self.confidence[i] if self.confidence is not None else None,
self.class_id[i] if self.class_id is not None else None,
self.tracker_id[i] if self.tracker_id is not None else None,
get_data_item(self.data, i),
)
__len__(self)
special
¶
Returns the number of detections in the Detections object.
Source code in supertracker/detection/core.py
def __len__(self):
"""
Returns the number of detections in the Detections object.
"""
return len(self.xyxy)
__setitem__(self, key, value)
special
¶
Set a value in the data dictionary of the Detections object.
Source code in supertracker/detection/core.py
def __setitem__(self, key: str, value: Union[np.ndarray, List]):
"""
Set a value in the data dictionary of the Detections object.
"""
if not isinstance(value, (np.ndarray, list)):
raise TypeError("Value must be a np.ndarray or a list")
if isinstance(value, list):
value = np.array(value)
self.data[key] = value
empty()
classmethod
¶
Create an empty Detections object with no bounding boxes, confidences, or class IDs.
Returns:
| Type | Description |
|---|---|
(Detections) |
An empty Detections object. |
Source code in supertracker/detection/core.py
@classmethod
def empty(cls) -> Detections:
"""
Create an empty Detections object with no bounding boxes,
confidences, or class IDs.
Returns:
(Detections): An empty Detections object.
"""
return cls(
xyxy=np.empty((0, 4), dtype=np.float32),
confidence=np.array([], dtype=np.float32),
class_id=np.array([], dtype=int),
)
from_yolo(custom_results)
classmethod
¶
Creates a sv.Detections instance from a custom Results class that mimics
the ultralytics Results format.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
custom_results |
The output Results instance from custom detector |
required |
Returns:
| Type | Description |
|---|---|
Detections |
A new Detections object. |
Source code in supertracker/detection/core.py
@classmethod
def from_yolo(cls, custom_results) -> Detections:
"""
Creates a `sv.Detections` instance from a custom Results class that mimics
the ultralytics Results format.
Args:
custom_results: The output Results instance from custom detector
Returns:
Detections: A new Detections object.
"""
# Handle OBB case first if available
if hasattr(custom_results, "obb") and custom_results.obb is not None:
class_id = custom_results.obb.cls.astype(int)
class_names = np.array([custom_results.names[i] for i in class_id])
oriented_box_coordinates = custom_results.obb.xyxyxyxy
return cls(
xyxy=custom_results.obb.xyxy,
confidence=custom_results.obb.conf,
class_id=class_id,
tracker_id=custom_results.obb.id.astype(int) if custom_results.obb.id is not None else None,
data={
ORIENTED_BOX_COORDINATES: oriented_box_coordinates,
CLASS_NAME_DATA_FIELD: class_names,
},
)
# Handle regular detection case
if hasattr(custom_results, "boxes") and custom_results.boxes is not None:
class_id = custom_results.boxes.cls.astype(int)
class_names = np.array([custom_results.names[i] for i in class_id])
return cls(
xyxy=custom_results.boxes.xyxy,
confidence=custom_results.boxes.conf,
class_id=class_id,
mask=custom_results.masks.data if hasattr(custom_results, "masks") and custom_results.masks is not None else None,
tracker_id=custom_results.boxes.id.astype(int) if hasattr(custom_results.boxes, "id") and custom_results.boxes.id is not None else None,
data={CLASS_NAME_DATA_FIELD: class_names},
)
# Handle case with only masks
if hasattr(custom_results, "masks") and custom_results.masks is not None:
masks = custom_results.masks.data
return cls(
xyxy=mask_to_xyxy(masks),
mask=masks,
class_id=np.arange(len(custom_results)),
)
# Return empty detections if no supported format is found
return cls.empty()
is_empty(self)
¶
Returns True if the Detections object is considered empty.
Source code in supertracker/detection/core.py
def is_empty(self) -> bool:
"""
Returns `True` if the `Detections` object is considered empty.
"""
empty_detections = Detections.empty()
empty_detections.data = self.data
empty_detections.metadata = self.metadata
return self == empty_detections
utils
¶
box_iou_batch(boxes_true, boxes_detection)
¶
Compute Intersection over Union (IoU) of two sets of bounding boxes -
boxes_true and boxes_detection. Both sets
of boxes are expected to be in (x_min, y_min, x_max, y_max) format.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
boxes_true |
np.ndarray |
2D |
required |
boxes_detection |
np.ndarray |
2D |
required |
Returns:
| Type | Description |
|---|---|
np.ndarray |
Pairwise IoU of boxes from |
Source code in supertracker/detection/utils.py
def box_iou_batch(boxes_true: np.ndarray, boxes_detection: np.ndarray) -> np.ndarray:
"""
Compute Intersection over Union (IoU) of two sets of bounding boxes -
`boxes_true` and `boxes_detection`. Both sets
of boxes are expected to be in `(x_min, y_min, x_max, y_max)` format.
Args:
boxes_true (np.ndarray): 2D `np.ndarray` representing ground-truth boxes.
`shape = (N, 4)` where `N` is number of true objects.
boxes_detection (np.ndarray): 2D `np.ndarray` representing detection boxes.
`shape = (M, 4)` where `M` is number of detected objects.
Returns:
np.ndarray: Pairwise IoU of boxes from `boxes_true` and `boxes_detection`.
`shape = (N, M)` where `N` is number of true objects and
`M` is number of detected objects.
"""
def box_area(box):
return (box[2] - box[0]) * (box[3] - box[1])
area_true = box_area(boxes_true.T)
area_detection = box_area(boxes_detection.T)
top_left = np.maximum(boxes_true[:, None, :2], boxes_detection[:, :2])
bottom_right = np.minimum(boxes_true[:, None, 2:], boxes_detection[:, 2:])
area_inter = np.prod(np.clip(bottom_right - top_left, a_min=0, a_max=None), 2)
ious = area_inter / (area_true[:, None] + area_detection - area_inter)
ious = np.nan_to_num(ious)
return ious
get_data_item(data, index)
¶
Retrieve a subset of the data dictionary based on the given index.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
data |
Dict[str, Union[numpy.ndarray, List]] |
The data dictionary of the Detections object. |
required |
index |
Union[int, slice, List[int], numpy.ndarray] |
The index or indices specifying the subset to retrieve. |
required |
Returns:
| Type | Description |
|---|---|
Dict[str, Union[numpy.ndarray, List]] |
A subset of the data dictionary corresponding to the specified index. |
Source code in supertracker/detection/utils.py
def get_data_item(
data: Dict[str, Union[np.ndarray, List]],
index: Union[int, slice, List[int], np.ndarray],
) -> Dict[str, Union[np.ndarray, List]]:
"""
Retrieve a subset of the data dictionary based on the given index.
Args:
data: The data dictionary of the Detections object.
index: The index or indices specifying the subset to retrieve.
Returns:
A subset of the data dictionary corresponding to the specified index.
"""
subset_data = {}
for key, value in data.items():
if isinstance(value, np.ndarray):
subset_data[key] = value[index]
elif isinstance(value, list):
if isinstance(index, slice):
subset_data[key] = value[index]
elif isinstance(index, list):
subset_data[key] = [value[i] for i in index]
elif isinstance(index, np.ndarray):
if index.dtype == bool:
subset_data[key] = [
value[i] for i, index_value in enumerate(index) if index_value
]
else:
subset_data[key] = [value[i] for i in index]
elif isinstance(index, int):
subset_data[key] = [value[index]]
else:
raise TypeError(f"Unsupported index type: {type(index)}")
else:
raise TypeError(f"Unsupported data type for key '{key}': {type(value)}")
return subset_data
is_data_equal(data_a, data_b)
¶
Compares the data payloads of two Detections instances.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
data_a, |
data_b |
The data payloads of the instances. |
required |
Returns:
| Type | Description |
|---|---|
bool |
True if the data payloads are equal, False otherwise. |
Source code in supertracker/detection/utils.py
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.
Args:
data_a, data_b: The data payloads of the instances.
Returns:
True if the data payloads are equal, False otherwise.
"""
return set(data_a.keys()) == set(data_b.keys()) and all(
np.array_equal(data_a[key], data_b[key]) for key in data_a
)
is_metadata_equal(metadata_a, metadata_b)
¶
Compares the metadata payloads of two Detections instances.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
metadata_a, |
metadata_b |
The metadata payloads of the instances. |
required |
Returns:
| Type | Description |
|---|---|
bool |
True if the metadata payloads are equal, False otherwise. |
Source code in supertracker/detection/utils.py
def is_metadata_equal(metadata_a: Dict[str, Any], metadata_b: Dict[str, Any]) -> bool:
"""
Compares the metadata payloads of two Detections instances.
Args:
metadata_a, metadata_b: The metadata payloads of the instances.
Returns:
True if the metadata payloads are equal, False otherwise.
"""
return set(metadata_a.keys()) == set(metadata_b.keys()) and all(
np.array_equal(metadata_a[key], metadata_b[key])
if (
isinstance(metadata_a[key], np.ndarray)
and isinstance(metadata_b[key], np.ndarray)
)
else metadata_a[key] == metadata_b[key]
for key in metadata_a
)
mask_to_xyxy(masks)
¶
Converts a 3D np.array of 2D bool masks into a 2D np.array of bounding boxes.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
masks |
np.ndarray |
A 3D |
required |
Returns:
| Type | Description |
|---|---|
np.ndarray |
A 2D |
Source code in supertracker/detection/utils.py
def mask_to_xyxy(masks: np.ndarray) -> np.ndarray:
"""
Converts a 3D `np.array` of 2D bool masks into a 2D `np.array` of bounding boxes.
Parameters:
masks (np.ndarray): A 3D `np.array` of shape `(N, W, H)`
containing 2D bool masks
Returns:
np.ndarray: A 2D `np.array` of shape `(N, 4)` containing the bounding boxes
`(x_min, y_min, x_max, y_max)` for each mask
"""
n = masks.shape[0]
xyxy = np.zeros((n, 4), dtype=int)
for i, mask in enumerate(masks):
rows, cols = np.where(mask)
if len(rows) > 0 and len(cols) > 0:
x_min, x_max = np.min(cols), np.max(cols)
y_min, y_max = np.min(rows), np.max(rows)
xyxy[i, :] = [x_min, y_min, x_max, y_max]
return xyxy
validate_detections_fields(xyxy, mask, confidence, class_id, tracker_id, data)
¶
Simplified validation for detection fields
Source code in supertracker/detection/utils.py
def validate_detections_fields(
xyxy: Any,
mask: Any,
confidence: Any,
class_id: Any,
tracker_id: Any,
data: Dict[str, Any],
) -> None:
"""Simplified validation for detection fields"""
if not isinstance(xyxy, np.ndarray) or xyxy.ndim != 2 or xyxy.shape[1] != 4:
raise ValueError("xyxy must be a 2D numpy array with shape (N, 4)")
n = len(xyxy)
# Validate mask if provided
if mask is not None and (not isinstance(mask, np.ndarray) or mask.ndim != 3 or mask.shape[0] != n):
raise ValueError(f"mask must be a 3D numpy array with shape ({n}, H, W)")
# Validate confidence if provided
if confidence is not None and (not isinstance(confidence, np.ndarray) or confidence.shape != (n,)):
raise ValueError(f"confidence must be a 1D numpy array with shape ({n},)")
# Validate class_id if provided
if class_id is not None and (not isinstance(class_id, np.ndarray) or class_id.shape != (n,)):
raise ValueError(f"class_id must be a 1D numpy array with shape ({n},)")
# Validate tracker_id if provided
if tracker_id is not None and (not isinstance(tracker_id, np.ndarray) or tracker_id.shape != (n,)):
raise ValueError(f"tracker_id must be a 1D numpy array with shape ({n},)")
# Validate data dictionary
if data:
for key, value in data.items():
if isinstance(value, list):
if len(value) != n:
raise ValueError(f"Length of list for key '{key}' must be {n}")
elif isinstance(value, np.ndarray):
if value.ndim == 1 and value.shape[0] != n:
raise ValueError(f"Shape of np.ndarray for key '{key}' must be ({n},)")
elif value.ndim > 1 and value.shape[0] != n:
raise ValueError(f"First dimension of np.ndarray for key '{key}' must have size {n}")
else:
raise ValueError(f"Value for key '{key}' must be a list or np.ndarray")