superiorvision 0.30.0.dev0__py3-none-any.whl

supervision/tracker/byte_tracker/core.py

@@ -0,0 +1,448 @@
+import numpy as np
+import numpy.typing as npt
+import torch
+
+from supervision.detection.core import Detections
+from supervision.detection.utils import box_iou_batch
+from supervision.tracker.byte_tracker import matching
+from supervision.tracker.byte_tracker.kalman_filter import KalmanFilter
+from supervision.tracker.byte_tracker.single_object_track import STrack, TrackState
+from supervision.tracker.byte_tracker.utils import IdCounter
+from supervision.utils.deprecate import TargetMode, deprecated_class
+
+
+def _valid_tracking_tensors(
+    tensors: torch.Tensor,
+) -> torch.Tensor:
+    """Identify finite tensors with positive-width and positive-height boxes."""
+    bboxes = tensors[:, :4]
+    widths = bboxes[:, 2] - bboxes[:, 0]
+    heights = bboxes[:, 3] - bboxes[:, 1]
+    return torch.isfinite(tensors).all(dim=1) & (widths > 0) & (heights > 0)
+
+
+@deprecated_class(
+    target=TargetMode.NOTIFY,
+    deprecated_in="0.28.0",
+    remove_in="0.31.0",
+)
+class ByteTrack:
+    """
+    Initialize the ByteTrack object.
+
+    !!! deprecated "Deprecated"
+
+        `ByteTrack` is deprecated since `supervision-0.28.0` and will be removed in
+        `supervision-0.31.0`. Use `ByteTrackTracker` from the private `tracktors`
+        fork instead. Note: the update method is renamed from
+        `update_with_detections()` to `update()`.
+
+    <video controls>
+        <source src="https://media.roboflow.com/supervision/video-examples/how-to/track-objects/annotate-video-with-traces.mp4" type="video/mp4">
+    </video>
+
+    Args:
+        track_activation_threshold: Detection confidence threshold
+            for track activation. Increasing track_activation_threshold improves accuracy
+            and stability but might miss true detections. Decreasing it increases
+            completeness but risks introducing noise and instability.
+        lost_track_buffer: Number of frames to buffer when a track is lost.
+            Increasing lost_track_buffer enhances occlusion handling, significantly
+            reducing the likelihood of track fragmentation or disappearance caused
+            by brief detection gaps.
+        minimum_matching_threshold: Threshold for matching tracks with detections.
+            Decreasing minimum_matching_threshold improves accuracy but risks fragmentation.
+            Increasing it improves completeness but risks false positives and drift.
+        frame_rate: The frame rate of the video. Accepts float values (e.g. 23.976,
+            29.97) for accurate lost-track-buffer calculation.
+        minimum_consecutive_frames: Number of consecutive frames that an object must
+            be tracked before it is considered a 'valid' track.
+            Increasing minimum_consecutive_frames prevents the creation of accidental tracks from
+            false detection or double detection, but risks missing shorter tracks.
+    """  # noqa: E501 // docs
+
+    def __init__(
+        self,
+        track_activation_threshold: float = 0.25,
+        lost_track_buffer: int = 30,
+        minimum_matching_threshold: float = 0.8,
+        frame_rate: float = 30,
+        minimum_consecutive_frames: int = 1,
+    ) -> None:
+        self.track_activation_threshold = track_activation_threshold
+        self.minimum_matching_threshold = minimum_matching_threshold
+
+        self.frame_id = 0
+        self.det_thresh = self.track_activation_threshold + 0.1
+        if self.det_thresh > 1.0:
+            self.det_thresh = self.track_activation_threshold
+        self.max_time_lost = int(frame_rate / 30.0 * lost_track_buffer)
+        self.minimum_consecutive_frames = minimum_consecutive_frames
+        self.kalman_filter = KalmanFilter()
+        self.shared_kalman = KalmanFilter()
+
+        self.tracked_tracks: list[STrack] = []
+        self.lost_tracks: list[STrack] = []
+        self.removed_tracks: list[STrack] = []
+
+        # Warning, possible bug: If you also set internal_id to start at 1,
+        # all traces will be connected across objects.
+        self.internal_id_counter = IdCounter()
+        self.external_id_counter = IdCounter(start_id=1)
+
+    def update_with_detections(self, detections: Detections) -> Detections:
+        """
+        Updates the tracker with the provided detections and returns the updated
+        detection results.
+
+        Args:
+            detections: The detections to pass through the tracker.
+
+        Example:
+            ```python
+            import supervision as sv
+            from ultralytics import YOLO
+
+            model = YOLO("<MODEL_PATH>")
+            tracker = sv.ByteTrack()
+
+            box_annotator = sv.BoxAnnotator()
+            label_annotator = sv.LabelAnnotator()
+
+            def callback(frame: np.ndarray, index: int) -> np.ndarray:
+                results = model(frame)[0]
+                detections = sv.Detections.from_ultralytics(results)
+                detections = tracker.update_with_detections(detections)
+
+                labels = [f"#{tracker_id}" for tracker_id in detections.tracker_id]
+
+                annotated_frame = box_annotator.annotate(
+                    scene=frame.copy(), detections=detections)
+                annotated_frame = label_annotator.annotate(
+                    scene=annotated_frame, detections=detections, labels=labels)
+                return annotated_frame
+
+            sv.process_video(
+                source_path="<SOURCE_VIDEO_PATH>",
+                target_path="<TARGET_VIDEO_PATH>",
+                callback=callback
+            )
+            ```
+        """
+        if detections.confidence is None:
+            raise ValueError("Detections confidence must be provided for tracking.")
+
+        tensors = torch.cat((detections.xyxy, detections.confidence[:, None]), dim=1)
+        tracks = self.update_with_tensors(tensors=tensors)
+
+        if len(tracks) > 0:
+            detection_bounding_boxes = tensors[:, :4]
+            track_bounding_boxes = torch.stack(
+                [
+                    track.tlbr.to(
+                        dtype=detection_bounding_boxes.dtype,
+                        device=detection_bounding_boxes.device,
+                    )
+                    for track in tracks
+                ]
+            )
+
+            ious = box_iou_batch(detection_bounding_boxes, track_bounding_boxes)
+
+            iou_costs = 1 - ious
+
+            matches, _, _ = matching.linear_assignment(iou_costs, 0.5)
+            tracked_detections = detections.select(slice(None))
+            tracked_detections.tracker_id = torch.full(
+                (len(detections),),
+                -1,
+                dtype=torch.long,
+                device=detections.xyxy.device,
+            )
+            for i_detection, i_track in matches:
+                detection_index = int(i_detection.item())
+                track_index = int(i_track.item())
+                tracked_detections.tracker_id[detection_index] = tracks[
+                    track_index
+                ].external_track_id
+
+            return tracked_detections.select(tracked_detections.tracker_id != -1)
+
+        else:
+            empty = detections.select(
+                torch.zeros(
+                    len(detections), dtype=torch.bool, device=detections.xyxy.device
+                )
+            )
+            empty.tracker_id = torch.empty(
+                0, dtype=torch.long, device=detections.xyxy.device
+            )
+            return empty
+
+    def reset(self) -> None:
+        """
+        Resets the internal state of the ByteTrack tracker.
+
+        This method clears the tracking data, including tracked, lost,
+        and removed tracks, as well as resetting the frame counter. It's
+        particularly useful when processing multiple videos sequentially,
+        ensuring the tracker starts with a clean state for each new video.
+        """
+        self.frame_id = 0
+        self.internal_id_counter.reset()
+        self.external_id_counter.reset()
+        self.tracked_tracks = []
+        self.lost_tracks = []
+        self.removed_tracks = []
+
+    def update_with_tensors(
+        self, tensors: npt.NDArray[np.float32] | torch.Tensor
+    ) -> list[STrack]:
+        """
+        Updates the tracker with the provided tensors and returns the updated tracks.
+
+        Args:
+            tensors: The new tensors to update with.
+
+        Returns:
+            Updated tracks.
+        """
+        tensors = torch.as_tensor(tensors)
+        if not tensors.is_floating_point():
+            tensors = tensors.to(dtype=torch.float32)
+        self.frame_id += 1
+        activated_starcks = []
+        refind_stracks = []
+        lost_stracks = []
+        removed_stracks = []
+
+        tensors = tensors[_valid_tracking_tensors(tensors)]
+        scores = tensors[:, 4]
+        bboxes = tensors[:, :4]
+
+        remain_inds = scores >= self.track_activation_threshold
+        inds_low = scores > 0.1
+        inds_high = scores < self.track_activation_threshold
+
+        inds_second = torch.logical_and(inds_low, inds_high)
+        dets_second = bboxes[inds_second]
+        dets = bboxes[remain_inds]
+        scores_keep = scores[remain_inds]
+        scores_second = scores[inds_second]
+
+        if len(dets) > 0:
+            """Detections"""
+            detections = [
+                STrack(
+                    STrack.tlbr_to_tlwh(tlbr),
+                    score_keep,
+                    self.minimum_consecutive_frames,
+                    self.shared_kalman,
+                    self.internal_id_counter,
+                    self.external_id_counter,
+                )
+                for (tlbr, score_keep) in zip(dets, scores_keep)
+            ]
+        else:
+            detections = []
+
+        """ Add newly detected tracklets to tracked_stracks"""
+        unconfirmed = []
+        tracked_stracks: list[STrack] = []
+
+        for track in self.tracked_tracks:
+            if not track.is_activated:
+                unconfirmed.append(track)
+            else:
+                tracked_stracks.append(track)
+
+        """ Step 2: First association, with high score detection boxes"""
+        strack_pool = joint_tracks(tracked_stracks, self.lost_tracks)
+        # Predict the current location with KF
+        STrack.multi_predict(strack_pool, self.shared_kalman)
+        dists = matching.iou_distance(strack_pool, detections)
+
+        dists = matching.fuse_score(dists, detections)
+        matches, u_track, u_detection = matching.linear_assignment(
+            dists, thresh=self.minimum_matching_threshold
+        )
+
+        for itracked, idet in matches:
+            itracked, idet = int(itracked.item()), int(idet.item())
+            track = strack_pool[itracked]
+            det = detections[idet]
+            if track.state == TrackState.Tracked:
+                track.update(detections[idet], self.frame_id)
+                activated_starcks.append(track)
+            else:
+                track.re_activate(det, self.frame_id)
+                refind_stracks.append(track)
+
+        """ Step 3: Second association, with low score detection boxes"""
+        # association the untrack to the low score detections
+        if len(dets_second) > 0:
+            """Detections"""
+            detections_second = [
+                STrack(
+                    STrack.tlbr_to_tlwh(tlbr),
+                    score_second,
+                    self.minimum_consecutive_frames,
+                    self.shared_kalman,
+                    self.internal_id_counter,
+                    self.external_id_counter,
+                )
+                for (tlbr, score_second) in zip(dets_second, scores_second)
+            ]
+        else:
+            detections_second = []
+        r_tracked_stracks = [
+            strack_pool[i]
+            for i in u_track
+            if strack_pool[i].state == TrackState.Tracked
+        ]
+        dists = matching.iou_distance(r_tracked_stracks, detections_second)
+        matches, u_track, _u_detection_second = matching.linear_assignment(
+            dists, thresh=0.5
+        )
+        for itracked, idet in matches:
+            itracked, idet = int(itracked.item()), int(idet.item())
+            track = r_tracked_stracks[itracked]
+            det = detections_second[idet]
+            if track.state == TrackState.Tracked:
+                track.update(det, self.frame_id)
+                activated_starcks.append(track)
+            else:
+                track.re_activate(det, self.frame_id)
+                refind_stracks.append(track)
+
+        for it in u_track:
+            track = r_tracked_stracks[it]
+            if not track.state == TrackState.Lost:
+                track.state = TrackState.Lost
+                lost_stracks.append(track)
+
+        """Deal with unconfirmed tracks, usually tracks with only one beginning frame"""
+        detections = [detections[i] for i in u_detection]
+        dists = matching.iou_distance(unconfirmed, detections)
+
+        dists = matching.fuse_score(dists, detections)
+        matches, u_unconfirmed, u_detection = matching.linear_assignment(
+            dists, thresh=0.7
+        )
+        for itracked, idet in matches:
+            itracked, idet = int(itracked.item()), int(idet.item())
+            unconfirmed[itracked].update(detections[idet], self.frame_id)
+            activated_starcks.append(unconfirmed[itracked])
+        for it in u_unconfirmed:
+            track = unconfirmed[it]
+            track.state = TrackState.Removed
+            removed_stracks.append(track)
+
+        """ Step 4: Init new stracks"""
+        for inew in u_detection:
+            track = detections[inew]
+            if track.score < self.det_thresh:
+                continue
+            track.activate(self.kalman_filter, self.frame_id)
+            activated_starcks.append(track)
+        """ Step 5: Update state"""
+        for track in self.lost_tracks:
+            if self.frame_id - track.frame_id > self.max_time_lost:
+                track.state = TrackState.Removed
+                removed_stracks.append(track)
+
+        self.tracked_tracks = [
+            t for t in self.tracked_tracks if t.state == TrackState.Tracked
+        ]
+        self.tracked_tracks = joint_tracks(self.tracked_tracks, activated_starcks)
+        self.tracked_tracks = joint_tracks(self.tracked_tracks, refind_stracks)
+        self.lost_tracks = sub_tracks(self.lost_tracks, self.tracked_tracks)
+        self.lost_tracks.extend(lost_stracks)
+        self.lost_tracks = sub_tracks(self.lost_tracks, self.removed_tracks)
+        self.removed_tracks = removed_stracks
+        self.tracked_tracks, self.lost_tracks = remove_duplicate_tracks(
+            self.tracked_tracks, self.lost_tracks
+        )
+        output_stracks = [track for track in self.tracked_tracks if track.is_activated]
+
+        return output_stracks
+
+
+def joint_tracks(
+    track_list_a: list[STrack], track_list_b: list[STrack]
+) -> list[STrack]:
+    """
+    Joins two lists of tracks, ensuring that the resulting list does not
+    contain tracks with duplicate internal_track_id values.
+
+    Args:
+        track_list_a: First list of tracks.
+        track_list_b: Second list of tracks.
+
+    Returns:
+        Combined list of tracks from track_list_a and track_list_b
+            without duplicate internal_track_id values.
+    """
+    seen_track_ids = set()
+    result = []
+
+    for track in track_list_a + track_list_b:
+        if track.internal_track_id not in seen_track_ids:
+            seen_track_ids.add(track.internal_track_id)
+            result.append(track)
+
+    return result
+
+
+def sub_tracks(track_list_a: list[STrack], track_list_b: list[STrack]) -> list[STrack]:
+    """
+    Returns a list of tracks from track_list_a after removing any tracks
+    that share the same internal_track_id with tracks in track_list_b.
+
+    Args:
+        track_list_a: List of tracks.
+        track_list_b: List of tracks to be subtracted from track_list_a.
+    Returns:
+        List of remaining tracks from track_list_a after subtraction.
+    """
+    tracks = {track.internal_track_id: track for track in track_list_a}
+    track_ids_b = {track.internal_track_id for track in track_list_b}
+
+    for track_id in track_ids_b:
+        tracks.pop(track_id, None)
+
+    return list(tracks.values())
+
+
+def remove_duplicate_tracks(
+    tracks_a: list[STrack], tracks_b: list[STrack]
+) -> tuple[list[STrack], list[STrack]]:
+    pairwise_distance = matching.iou_distance(tracks_a, tracks_b)
+    matching_pairs = (
+        torch.where(pairwise_distance < 0.05)
+        if isinstance(pairwise_distance, torch.Tensor)
+        else np.where(pairwise_distance < 0.05)
+    )
+
+    duplicates_a, duplicates_b = set(), set()
+    for track_index_a, track_index_b in zip(*matching_pairs):
+        if isinstance(track_index_a, torch.Tensor):
+            track_index_a, track_index_b = (
+                int(track_index_a.item()),
+                int(track_index_b.item()),
+            )
+        time_a = tracks_a[track_index_a].frame_id - tracks_a[track_index_a].start_frame
+        time_b = tracks_b[track_index_b].frame_id - tracks_b[track_index_b].start_frame
+        if time_a > time_b:
+            duplicates_b.add(track_index_b)
+        else:
+            duplicates_a.add(track_index_a)
+
+    result_a = [
+        track for index, track in enumerate(tracks_a) if index not in duplicates_a
+    ]
+    result_b = [
+        track for index, track in enumerate(tracks_b) if index not in duplicates_b
+    ]
+
+    return result_a, result_b

supervision/tracker/byte_tracker/kalman_filter.py

@@ -0,0 +1,186 @@
+import numpy as np
+import numpy.typing as npt
+import torch
+
+TensorLike = npt.NDArray[np.floating] | torch.Tensor
+
+
+def _as_float_tensor(
+    value: TensorLike,
+    *,
+    device: torch.device | None = None,
+    dtype: torch.dtype | None = None,
+) -> torch.Tensor:
+    tensor = torch.as_tensor(value, device=device)
+    if dtype is not None:
+        return tensor.to(dtype=dtype)
+    return tensor if tensor.is_floating_point() else tensor.to(dtype=torch.float32)
+
+
+class KalmanFilter:
+    """
+    A simple Kalman filter for tracking bounding boxes in image space.
+
+    The 8-dimensional state space is (x, y, a, h, vx, vy, va, vh), where
+    (x, y) is the bounding box center, a is the aspect ratio (w/h), h is
+    the height, and their respective velocities.
+
+    Object motion follows a constant velocity model. The bounding box location
+    (x, y, a, h) is taken as direct observation of the state space (linear
+    observation model).
+    """
+
+    def __init__(self) -> None:
+        self._std_weight_position: float = 1.0 / 20
+        self._std_weight_velocity: float = 1.0 / 160
+
+    def initiate(self, measurement: TensorLike) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Create track from unassociated measurement.
+
+        Args:
+            measurement: The initial measurement vector.
+
+        Returns:
+            The mean vector and covariance matrix of the new track.
+        """
+        measurement = _as_float_tensor(measurement)
+        mean = torch.cat((measurement, torch.zeros_like(measurement)))
+        std = torch.stack(
+            (
+                2 * self._std_weight_position * measurement[3],
+                2 * self._std_weight_position * measurement[3],
+                measurement.new_tensor(1e-2),
+                2 * self._std_weight_position * measurement[3],
+                10 * self._std_weight_velocity * measurement[3],
+                10 * self._std_weight_velocity * measurement[3],
+                measurement.new_tensor(1e-5),
+                10 * self._std_weight_velocity * measurement[3],
+            )
+        )
+        return mean, torch.diag(std.square())
+
+    def predict(
+        self, mean: TensorLike, covariance: TensorLike
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Run Kalman filter prediction step.
+
+        Args:
+            mean: The object state mean at the previous time step.
+            covariance: The object state covariance at the previous time step.
+
+        Returns:
+            The mean vector and covariance matrix of the predicted state.
+        """
+        mean = _as_float_tensor(mean)
+        covariance = _as_float_tensor(covariance, device=mean.device, dtype=mean.dtype)
+        motion = torch.eye(8, dtype=mean.dtype, device=mean.device)
+        motion[range(4), range(4, 8)] = 1
+        std = torch.stack(
+            (
+                self._std_weight_position * mean[3],
+                self._std_weight_position * mean[3],
+                mean.new_tensor(1e-2),
+                self._std_weight_position * mean[3],
+                self._std_weight_velocity * mean[3],
+                self._std_weight_velocity * mean[3],
+                mean.new_tensor(1e-5),
+                self._std_weight_velocity * mean[3],
+            )
+        )
+        return mean @ motion.T, motion @ covariance @ motion.T + torch.diag(
+            std.square()
+        )
+
+    def project(
+        self, mean: TensorLike, covariance: TensorLike
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Project state distribution to measurement space.
+
+        Args:
+            mean: The state's mean vector.
+            covariance: The state's covariance matrix.
+
+        Returns:
+            The projected mean and covariance matrix of the given state estimate.
+        """
+        mean = _as_float_tensor(mean)
+        covariance = _as_float_tensor(covariance, device=mean.device, dtype=mean.dtype)
+        update = torch.eye(4, 8, dtype=mean.dtype, device=mean.device)
+        std = torch.stack(
+            (
+                self._std_weight_position * mean[3],
+                self._std_weight_position * mean[3],
+                mean.new_tensor(1e-1),
+                self._std_weight_position * mean[3],
+            )
+        )
+        return update @ mean, update @ covariance @ update.T + torch.diag(std.square())
+
+    def multi_predict(
+        self, mean: TensorLike, covariance: TensorLike
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Run Kalman filter prediction step (Vectorized version).
+        Args:
+            mean: The object state means at the previous time step.
+            covariance: The object state covariances at the previous time step.
+
+        Returns:
+            The mean vector and covariance matrix of the predicted state.
+        """
+        mean = _as_float_tensor(mean)
+        covariance = _as_float_tensor(covariance, device=mean.device, dtype=mean.dtype)
+        motion = torch.eye(8, dtype=mean.dtype, device=mean.device)
+        motion[range(4), range(4, 8)] = 1
+        height = mean[:, 3]
+        std = torch.stack(
+            (
+                self._std_weight_position * height,
+                self._std_weight_position * height,
+                torch.full_like(height, 1e-2),
+                self._std_weight_position * height,
+                self._std_weight_velocity * height,
+                self._std_weight_velocity * height,
+                torch.full_like(height, 1e-5),
+                self._std_weight_velocity * height,
+            ),
+            dim=1,
+        )
+        return (
+            mean @ motion.T,
+            motion @ covariance @ motion.T + torch.diag_embed(std.square()),
+        )
+
+    def update(
+        self,
+        mean: TensorLike,
+        covariance: TensorLike,
+        measurement: TensorLike,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Run Kalman filter correction step.
+
+        Args:
+            mean: The predicted state's mean vector.
+            covariance: The state's covariance matrix.
+            measurement: The measurement vector.
+
+        Returns:
+            The measurement-corrected state distribution.
+        """
+        mean = _as_float_tensor(mean)
+        covariance = _as_float_tensor(covariance, device=mean.device, dtype=mean.dtype)
+        measurement = _as_float_tensor(
+            measurement, device=mean.device, dtype=mean.dtype
+        )
+        projected_mean, projected_cov = self.project(mean, covariance)
+        update = torch.eye(4, 8, dtype=mean.dtype, device=mean.device)
+        kalman_gain = torch.linalg.solve(projected_cov, (covariance @ update.T).T).T
+        innovation = measurement - projected_mean
+        return (
+            mean + innovation @ kalman_gain.T,
+            covariance - kalman_gain @ projected_cov @ kalman_gain.T,
+        )