Skip to content

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 np.ndarray representing ground-truth boxes. shape = (N, 4) where N is number of true objects.

required
boxes_detection np.ndarray

2D np.ndarray representing detection boxes. shape = (M, 4) where M is number of detected objects.

required

Returns:

Type Description
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.

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 np.array of shape (N, W, H) containing 2D bool masks

required

Returns:

Type Description
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

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")