superiorvision 0.30.0.dev0__py3-none-any.whl

supervision/tracker/byte_tracker/matching.py

@@ -0,0 +1,211 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+import torch
+from scipy.optimize import linear_sum_assignment
+
+# The public utility dispatches on tensor inputs and keeps IoU on the source
+# device.  The older ``iou_and_nms`` implementation is NumPy-only, which made
+# the tracker silently fall back to host arrays after its Kalman prediction.
+from supervision.detection.utils import box_iou_batch
+
+if TYPE_CHECKING:
+    from supervision.tracker.byte_tracker.core import STrack
+
+
+def indices_to_matches(
+    cost_matrix: np.ndarray, indices: np.ndarray, thresh: float
+) -> tuple[np.ndarray, tuple, tuple]:
+    matched_cost = cost_matrix[tuple(zip(*indices))]
+    matched_mask = matched_cost <= thresh
+
+    matches = indices[matched_mask]
+    unmatched_a = tuple(set(range(cost_matrix.shape[0])) - set(matches[:, 0]))
+    unmatched_b = tuple(set(range(cost_matrix.shape[1])) - set(matches[:, 1]))
+    return matches, unmatched_a, unmatched_b
+
+
+def linear_assignment(
+    cost_matrix: np.ndarray | torch.Tensor, thresh: float
+) -> tuple[np.ndarray, tuple[int], tuple[int, int]]:
+    if isinstance(cost_matrix, torch.Tensor):
+        return _linear_assignment_tensor(cost_matrix, thresh)
+    if cost_matrix.size == 0:
+        return (
+            np.empty((0, 2), dtype=int),
+            tuple(range(cost_matrix.shape[0])),
+            tuple(range(cost_matrix.shape[1])),
+        )
+
+    assignment_cost_matrix = cost_matrix.copy()
+    assignment_cost_matrix[assignment_cost_matrix > thresh] = thresh + 1e-4
+    row_ind, col_ind = linear_sum_assignment(assignment_cost_matrix)
+    indices = np.column_stack((row_ind, col_ind))
+
+    return indices_to_matches(cost_matrix, indices, thresh)
+
+
+def _linear_assignment_tensor(
+    cost_matrix: torch.Tensor, thresh: float
+) -> tuple[torch.Tensor, tuple[int, ...], tuple[int, ...]]:
+    """Hungarian assignment without materialising the cost matrix on the host.
+
+    ByteTrack owns Python track objects, so the final unmatched identifiers are
+    intentionally Python tuples. Costs and the assignment itself stay on the
+    input tensor's device.
+    """
+    rows, cols = cost_matrix.shape
+    if rows == 0 or cols == 0:
+        return (
+            torch.empty((0, 2), dtype=torch.long, device=cost_matrix.device),
+            tuple(range(rows)),
+            tuple(range(cols)),
+        )
+    size = max(rows, cols)
+    fill = cost_matrix.new_tensor(thresh + 1e-4)
+    cost = cost_matrix.new_full((size, size), thresh + 1e-4)
+    cost[:rows, :cols] = torch.minimum(cost_matrix, fill)
+    # Standard Kuhn-Munkres augmenting-path algorithm. Potentials, predecessor
+    # state, and candidate costs stay as tensors on the source device. Each
+    # iteration reduces a whole candidate row at once; only the chosen column
+    # index crosses into Python to drive the finite control loop.
+    u = cost.new_zeros(size + 1)
+    v = cost.new_zeros(size + 1)
+    p = torch.zeros(size + 1, dtype=torch.long, device=cost.device)
+    way = torch.zeros(size + 1, dtype=torch.long, device=cost.device)
+    for i in range(1, size + 1):
+        p[0] = i
+        j0 = 0
+        minv = cost.new_full((size + 1,), torch.inf)
+        used = torch.zeros(size + 1, dtype=torch.bool, device=cost.device)
+        while True:
+            used[j0] = True
+            i0 = p[j0]
+            unused = ~used[1:]
+            cur = cost[i0 - 1] - u[i0] - v[1:]
+            better = unused & (cur < minv[1:])
+            minv[1:] = torch.where(better, cur, minv[1:])
+            way[1:] = torch.where(better, way.new_full((size,), j0), way[1:])
+
+            candidates = torch.where(
+                unused, minv[1:], minv.new_full((size,), torch.inf)
+            )
+            delta, next_column = torch.min(candidates, dim=0)
+            j1 = int(next_column.item()) + 1
+
+            used_columns = torch.nonzero(used, as_tuple=False).flatten()
+            u[p[used_columns]] += delta
+            v[used_columns] -= delta
+            minv[~used] -= delta
+            j0 = j1
+            if int(p[j0].item()) == 0:
+                break
+        while True:
+            j1 = int(way[j0].item())
+            p[j0] = p[j1]
+            j0 = j1
+            if j0 == 0:
+                break
+    pairs = torch.stack(
+        (
+            p[1:].clone() - 1,
+            torch.arange(size, device=cost.device),
+        ),
+        dim=1,
+    )
+    pairs = pairs[(pairs[:, 0] < rows) & (pairs[:, 1] < cols)]
+    matched = pairs[cost_matrix[pairs[:, 0], pairs[:, 1]] <= thresh]
+    matched = matched[torch.argsort(matched[:, 0])]
+    matched_rows = torch.zeros(rows, dtype=torch.bool, device=cost.device)
+    matched_cols = torch.zeros(cols, dtype=torch.bool, device=cost.device)
+    matched_rows[matched[:, 0]] = True
+    matched_cols[matched[:, 1]] = True
+    return (
+        matched,
+        tuple(torch.nonzero(~matched_rows, as_tuple=False).flatten().cpu().tolist()),
+        tuple(torch.nonzero(~matched_cols, as_tuple=False).flatten().cpu().tolist()),
+    )
+
+
+def iou_distance(
+    atracks: list[STrack], btracks: list[STrack]
+) -> np.ndarray | torch.Tensor:
+    if (
+        len(atracks) > 0
+        and (
+            isinstance(atracks[0], torch.Tensor)
+            or isinstance(atracks[0].tlbr, torch.Tensor)
+        )
+    ) or (
+        len(btracks) > 0
+        and (
+            isinstance(btracks[0], torch.Tensor)
+            or isinstance(btracks[0].tlbr, torch.Tensor)
+        )
+    ):
+        first = atracks[0] if atracks else btracks[0]
+        first_box = first if isinstance(first, torch.Tensor) else first.tlbr
+        device = first_box.device
+        atlbrs = (
+            torch.stack(
+                [
+                    track if isinstance(track, torch.Tensor) else track.tlbr
+                    for track in atracks
+                ]
+            )
+            if atracks
+            else torch.empty((0, 4), device=device)
+        )
+        btlbrs = (
+            torch.stack(
+                [
+                    track if isinstance(track, torch.Tensor) else track.tlbr
+                    for track in btracks
+                ]
+            )
+            if btracks
+            else torch.empty((0, 4), device=device)
+        )
+        if not len(atlbrs) or not len(btlbrs):
+            return torch.empty((len(atlbrs), len(btlbrs)), device=device)
+        return 1 - box_iou_batch(atlbrs, btlbrs)
+    if (len(atracks) > 0 and isinstance(atracks[0], np.ndarray)) or (
+        len(btracks) > 0 and isinstance(btracks[0], np.ndarray)
+    ):
+        atlbrs = atracks
+        btlbrs = btracks
+    else:
+        atlbrs = [track.tlbr for track in atracks]
+        btlbrs = [track.tlbr for track in btracks]
+
+    _ious = np.zeros((len(atlbrs), len(btlbrs)), dtype=np.float32)
+    if _ious.size != 0:
+        _ious = box_iou_batch(np.asarray(atlbrs), np.asarray(btlbrs))
+    cost_matrix = 1 - _ious
+
+    return cost_matrix
+
+
+def fuse_score(
+    cost_matrix: np.ndarray | torch.Tensor, stracks: list[STrack]
+) -> np.ndarray | torch.Tensor:
+    if isinstance(cost_matrix, torch.Tensor):
+        if cost_matrix.numel() == 0:
+            return cost_matrix
+        scores = torch.stack(
+            [
+                torch.as_tensor(track.score, device=cost_matrix.device)
+                for track in stracks
+            ]
+        )
+        return 1 - (1 - cost_matrix) * scores.unsqueeze(0)
+    if cost_matrix.size == 0:
+        return cost_matrix
+    iou_sim = 1 - cost_matrix
+    det_scores = np.array([strack.score for strack in stracks])
+    det_scores = np.expand_dims(det_scores, axis=0).repeat(cost_matrix.shape[0], axis=0)
+    fuse_sim = iou_sim * det_scores
+    fuse_cost = 1 - fuse_sim
+    return fuse_cost

supervision/tracker/byte_tracker/single_object_track.py

@@ -0,0 +1,194 @@
+from __future__ import annotations
+
+from enum import Enum
+
+import numpy as np
+import numpy.typing as npt
+import torch
+
+from supervision.tracker.byte_tracker.kalman_filter import KalmanFilter
+from supervision.tracker.byte_tracker.utils import IdCounter
+
+
+class TrackState(Enum):
+    New = 0
+    Tracked = 1
+    Lost = 2
+    Removed = 3
+
+
+class STrack:
+    def __init__(
+        self,
+        tlwh: npt.NDArray[np.float32] | torch.Tensor,
+        score: float | torch.Tensor,
+        minimum_consecutive_frames: int,
+        shared_kalman: KalmanFilter,
+        internal_id_counter: IdCounter,
+        external_id_counter: IdCounter,
+    ):
+        self.state = TrackState.New
+        self.is_activated = False
+        self.start_frame = 0
+        self.frame_id = 0
+
+        self._tlwh = torch.as_tensor(tlwh, dtype=torch.float32)
+        self.kalman_filter: KalmanFilter | None = None
+        self.shared_kalman = shared_kalman
+        self.mean: torch.Tensor | None = None
+        self.covariance: torch.Tensor | None = None
+        self.is_activated = False
+
+        self.score = torch.as_tensor(
+            score, dtype=self._tlwh.dtype, device=self._tlwh.device
+        )
+        self.tracklet_len = 0
+
+        self.minimum_consecutive_frames = minimum_consecutive_frames
+
+        self.internal_id_counter = internal_id_counter
+        self.external_id_counter = external_id_counter
+        self.internal_track_id = self.internal_id_counter.NO_ID
+        self.external_track_id = self.external_id_counter.NO_ID
+
+    def predict(self) -> None:
+        assert self.mean is not None
+        assert self.covariance is not None
+        assert self.kalman_filter is not None
+        mean_state = self.mean.clone()
+        if self.state != TrackState.Tracked:
+            mean_state[7] = 0
+        self.mean, self.covariance = self.kalman_filter.predict(
+            mean_state, self.covariance
+        )
+
+    @staticmethod
+    def multi_predict(stracks: list[STrack], shared_kalman: KalmanFilter) -> None:
+        if len(stracks) > 0:
+            multi_mean_states = []
+            multi_covariance = []
+            for i, st in enumerate(stracks):
+                assert st.mean is not None
+                assert st.covariance is not None
+                multi_mean_states.append(st.mean.clone())
+                multi_covariance.append(st.covariance)
+                if st.state != TrackState.Tracked:
+                    multi_mean_states[i][7] = 0
+
+            predicted_mean, predicted_covariance = shared_kalman.multi_predict(
+                torch.stack(multi_mean_states), torch.stack(multi_covariance)
+            )
+            for i, (mean, cov) in enumerate(zip(predicted_mean, predicted_covariance)):
+                stracks[i].mean = mean
+                stracks[i].covariance = cov
+
+    def activate(self, kalman_filter: KalmanFilter, frame_id: int) -> None:
+        """Start a new tracklet after its first detection."""
+        self.kalman_filter = kalman_filter
+        self.internal_track_id = self.internal_id_counter.new_id()
+        self.mean, self.covariance = self.kalman_filter.initiate(
+            self.tlwh_to_xyah(self._tlwh)
+        )
+
+        self.tracklet_len = 1
+        self.state = TrackState.Tracked
+        if frame_id == 1 and self.tracklet_len >= self.minimum_consecutive_frames:
+            self.is_activated = True
+            self.external_track_id = self.external_id_counter.new_id()
+
+        self.frame_id = frame_id
+        self.start_frame = frame_id
+
+    def re_activate(self, new_track: STrack, frame_id: int) -> None:
+        assert self.kalman_filter is not None
+        assert self.mean is not None
+        assert self.covariance is not None
+        self.mean, self.covariance = self.kalman_filter.update(
+            self.mean, self.covariance, self.tlwh_to_xyah(new_track.tlwh)
+        )
+        self.tracklet_len = 0
+        self.state = TrackState.Tracked
+
+        self.frame_id = frame_id
+        self.score = new_track.score
+
+    def update(self, new_track: STrack, frame_id: int) -> None:
+        """
+        Update a matched track.
+
+        Args:
+            new_track: The new track data.
+            frame_id: The current frame ID.
+        """
+        assert self.kalman_filter is not None
+        assert self.mean is not None
+        assert self.covariance is not None
+        self.frame_id = frame_id
+        self.tracklet_len += 1
+
+        new_tlwh = new_track.tlwh
+        self.mean, self.covariance = self.kalman_filter.update(
+            self.mean, self.covariance, self.tlwh_to_xyah(new_tlwh)
+        )
+        self.state = TrackState.Tracked
+        if self.tracklet_len >= self.minimum_consecutive_frames:
+            self.is_activated = True
+            if self.external_track_id == self.external_id_counter.NO_ID:
+                self.external_track_id = self.external_id_counter.new_id()
+
+        self.score = new_track.score
+
+    @property
+    def tlwh(self) -> torch.Tensor:
+        """Get current position in bounding box format `(top left x, top left y,
+        width, height)`.
+        """
+        if self.mean is None:
+            return self._tlwh.clone()
+        ret = self.mean[:4].clone()
+        ret[2] *= ret[3]
+        ret[:2] -= ret[2:] / 2
+        return ret
+
+    @property
+    def tlbr(self) -> torch.Tensor:
+        """Convert bounding box to format `(min x, min y, max x, max y)`, i.e.,
+        `(top left, bottom right)`.
+        """
+        ret = self.tlwh.clone()
+        ret[2:] += ret[:2]
+        return ret
+
+    @staticmethod
+    def tlwh_to_xyah(
+        tlwh: npt.NDArray[np.float32] | torch.Tensor,
+    ) -> torch.Tensor:
+        """Convert bounding box to format `(center x, center y, aspect ratio,
+        height)`, where the aspect ratio is `width / height`.
+        """
+        ret = torch.as_tensor(tlwh, dtype=torch.float32).clone()
+        ret[:2] += ret[2:] / 2
+        ret[2] /= ret[3]
+        return ret
+
+    def to_xyah(self) -> torch.Tensor:
+        return self.tlwh_to_xyah(self.tlwh)
+
+    @staticmethod
+    def tlbr_to_tlwh(
+        tlbr: npt.NDArray[np.float32] | torch.Tensor,
+    ) -> torch.Tensor:
+        ret = torch.as_tensor(tlbr, dtype=torch.float32).clone()
+        ret[2:] -= ret[:2]
+        return ret
+
+    @staticmethod
+    def tlwh_to_tlbr(
+        tlwh: npt.NDArray[np.float32] | torch.Tensor,
+    ) -> torch.Tensor:
+        ret = torch.as_tensor(tlwh, dtype=torch.float32).clone()
+        ret[2:] += ret[:2]
+        return ret
+
+    def __repr__(self) -> str:
+        return f"OT_{self.internal_track_id}_({self.start_frame}-{self.frame_id})"

supervision/tracker/byte_tracker/utils.py

@@ -0,0 +1,34 @@
+class IdCounter:
+    def __init__(self, start_id: int = 0) -> None:
+        """
+        Initialize the ID counter.
+
+        Args:
+            start_id: The starting integer for the counter.
+
+        Raises:
+            ValueError: If start_id is less than or equal to -1.
+        """
+        self.start_id = start_id
+        if self.start_id <= self.NO_ID:
+            raise ValueError(f"start_id must be greater than {self.NO_ID}")
+        self.reset()
+
+    def reset(self) -> None:
+        """Reset the counter to the initial start_id."""
+        self._id = self.start_id
+
+    def new_id(self) -> int:
+        """
+        Get the current ID and increment the counter.
+
+        Returns:
+            The newly assigned ID.
+        """
+        returned_id = self._id
+        self._id += 1
+        return returned_id
+
+    @property
+    def NO_ID(self) -> int:
+        return -1

supervision/utils/conversion.py

@@ -0,0 +1,223 @@
+import functools
+from collections.abc import Callable
+from typing import Any, TypeVar, cast
+
+import cv2
+import numpy as np
+import numpy.typing as npt
+from PIL import Image
+
+from supervision.draw.base import ImageType
+from supervision.utils.deprecate import deprecated, void
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def ensure_cv2_image_for_class_method(
+    annotate_func: F,
+) -> F:
+    """
+    Decorates `BaseAnnotator.annotate` implementations, converts scene to
+    an image type used internally by the annotators, converts back when annotation
+    is complete.
+
+    Assumes the annotators modify the scene in-place.
+
+    Raises:
+        TypeError: If `scene` is not a `numpy.ndarray` or `PIL.Image.Image`.
+    """
+
+    @functools.wraps(annotate_func)
+    def wrapper(self: Any, scene: ImageType, *args: Any, **kwargs: Any) -> Any:
+        if isinstance(scene, np.ndarray):
+            return annotate_func(self, scene, *args, **kwargs)
+
+        if isinstance(scene, Image.Image):
+            scene_np = pillow_to_cv2(scene)
+            annotated_np = annotate_func(self, scene_np, *args, **kwargs)
+            scene.paste(cv2_to_pillow(annotated_np))
+            return scene
+
+        raise TypeError(f"Unsupported image type: {type(scene)}")
+
+    return cast(F, wrapper)
+
+
+@deprecated(  # type: ignore[untyped-decorator]
+    target=ensure_cv2_image_for_class_method,
+    deprecated_in="0.27.0",
+    remove_in="0.31.0",
+)
+def ensure_cv2_image_for_annotation(
+    annotate_func: F,
+) -> F:
+    return cast(F, void(annotate_func))
+
+
+def ensure_cv2_image_for_standalone_function(
+    image_processing_fun: F,
+) -> F:
+    """
+    Decorates image processing functions that accept np.ndarray, converting `image` to
+    np.ndarray, converts back when processing is complete.
+
+    Assumes the annotators do NOT modify the scene in-place.
+
+    Raises:
+        TypeError: If `image` is not a `numpy.ndarray` or `PIL.Image.Image`.
+    """
+
+    @functools.wraps(image_processing_fun)
+    def wrapper(image: ImageType, *args: Any, **kwargs: Any) -> Any:
+        if isinstance(image, np.ndarray):
+            return image_processing_fun(image, *args, **kwargs)
+
+        if isinstance(image, Image.Image):
+            scene = pillow_to_cv2(image)
+            annotated = image_processing_fun(scene, *args, **kwargs)
+            return cv2_to_pillow(annotated)
+
+        raise TypeError(f"Unsupported image type: {type(image)}")
+
+    return cast(F, wrapper)
+
+
+def ensure_pil_image_for_class_method(
+    annotate_func: F,
+) -> F:
+    """
+    Decorates image processing functions that accept np.ndarray, converting `image` to
+    PIL image, converts back when processing is complete.
+
+    Assumes the annotators modify the scene in-place.
+
+    Raises:
+        TypeError: If `scene` is not a `numpy.ndarray` or `PIL.Image.Image`.
+    """
+
+    @functools.wraps(annotate_func)
+    def wrapper(self: Any, scene: ImageType, *args: Any, **kwargs: Any) -> Any:
+        if isinstance(scene, np.ndarray):
+            scene_pil = cv2_to_pillow(scene)
+            annotated_pil = annotate_func(self, scene_pil, *args, **kwargs)
+            np.copyto(scene, pillow_to_cv2(annotated_pil))
+            return scene
+
+        if isinstance(scene, Image.Image):
+            return cast(ImageType, annotate_func(self, scene, *args, **kwargs))
+
+        raise TypeError(f"Unsupported image type: {type(scene)}")
+
+    return cast(F, wrapper)
+
+
+@deprecated(  # type: ignore[untyped-decorator]
+    target=ensure_pil_image_for_class_method,
+    deprecated_in="0.27.0",
+    remove_in="0.31.0",
+)
+def ensure_pil_image_for_annotation(
+    annotate_func: F,
+) -> F:
+    return cast(F, void(annotate_func))
+
+
+@deprecated(  # type: ignore[untyped-decorator]
+    target=ensure_cv2_image_for_standalone_function,
+    deprecated_in="0.27.0",
+    remove_in="0.31.0",
+)
+def ensure_cv2_image_for_processing(
+    image_processing_fun: F,
+) -> F:
+    return cast(F, void(image_processing_fun))
+
+
+def images_to_cv2(
+    images: list[npt.NDArray[np.uint8] | Image.Image],
+) -> list[npt.NDArray[np.uint8]]:
+    """
+    Converts images provided either as Pillow images or OpenCV
+    images into OpenCV format.
+
+    Args:
+        images: Images to be converted
+
+    Returns:
+        List of input images in OpenCV format
+            (with order preserved).
+
+    """
+    result: list[npt.NDArray[np.uint8]] = []
+    for image in images:
+        if isinstance(image, Image.Image):
+            result.append(pillow_to_cv2(image))
+        else:
+            result.append(image)
+    return result
+
+
+def pillow_to_cv2(image: Image.Image) -> npt.NDArray[np.uint8]:
+    """
+    Converts Pillow image into OpenCV image, handling RGB -> BGR
+    conversion. Palette images are first expanded to RGB so palette indices are
+    resolved to their actual colors.
+
+    Args:
+        image: Pillow image in RGB, grayscale, or palette mode.
+
+    Returns:
+        Input image converted to OpenCV format.
+
+    Examples:
+        ```pycon
+        >>> from PIL import Image
+        >>> from supervision.utils.conversion import pillow_to_cv2
+        >>> image = Image.new("RGB", (10, 10), color=(255, 0, 0))
+        >>> scene = pillow_to_cv2(image)
+        >>> scene.shape
+        (10, 10, 3)
+        >>> scene[0, 0].tolist()
+        [0, 0, 255]
+
+        ```
+    """
+    if image.mode == "P":
+        image = image.convert("RGB")
+
+    scene = np.array(image)
+    if scene.ndim == 2:
+        return cast(npt.NDArray[np.uint8], scene.astype(np.uint8, copy=False))
+    scene = cv2.cvtColor(scene, cv2.COLOR_RGB2BGR)
+    # cvtColor already returns uint8 here, so astype is a no-op other than the
+    # full-image copy it forces; copy=False keeps the dtype guard without it.
+    return cast(npt.NDArray[np.uint8], scene.astype(np.uint8, copy=False))
+
+
+def cv2_to_pillow(image: npt.NDArray[np.uint8]) -> Image.Image:
+    """
+    Converts OpenCV image into Pillow image, handling BGR -> RGB
+    conversion.
+
+    Args:
+        image: OpenCV image (in BGR format).
+
+    Returns:
+        Input image converted to Pillow format.
+
+    Examples:
+        ```pycon
+        >>> import numpy as np
+        >>> from supervision.utils.conversion import cv2_to_pillow
+        >>> scene = np.zeros((10, 10, 3), dtype=np.uint8)
+        >>> scene[:, :, 2] = 255
+        >>> image = cv2_to_pillow(scene)
+        >>> image.size
+        (10, 10)
+        >>> image.getpixel((0, 0))
+        (255, 0, 0)
+
+        ```
+    """
+    rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    return Image.fromarray(rgb_image)