dgenerate-ultralytics-headless 8.3.137__py3-none-any.whl → 8.3.224__py3-none-any.whl

ultralytics/trackers/byte_tracker.py

@@ -1,5 +1,9 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+from __future__ import annotations
+
+from typing import Any
+
 import numpy as np
 
 from ..utils import LOGGER
@@ -10,8 +14,7 @@ from .utils.kalman_filter import KalmanFilterXYAH
 
 
 class STrack(BaseTrack):
-    """
-    Single object tracking representation that uses Kalman filtering for state estimation.
+    """Single object tracking representation that uses Kalman filtering for state estimation.
 
     This class is responsible for storing all the information regarding individual tracklets and performs state updates
     and predictions based on Kalman filter.
@@ -29,16 +32,17 @@ class STrack(BaseTrack):
         idx (int): Index or identifier for the object.
         frame_id (int): Current frame ID.
         start_frame (int): Frame where the object was first detected.
+        angle (float | None): Optional angle information for oriented bounding boxes.
 
     Methods:
-        predict(): Predict the next state of the object using Kalman filter.
-        multi_predict(stracks): Predict the next states for multiple tracks.
-        multi_gmc(stracks, H): Update multiple track states using a homography matrix.
-        activate(kalman_filter, frame_id): Activate a new tracklet.
-        re_activate(new_track, frame_id, new_id): Reactivate a previously lost tracklet.
-        update(new_track, frame_id): Update the state of a matched track.
-        convert_coords(tlwh): Convert bounding box to x-y-aspect-height format.
-        tlwh_to_xyah(tlwh): Convert tlwh bounding box to xyah format.
+        predict: Predict the next state of the object using Kalman filter.
+        multi_predict: Predict the next states for multiple tracks.
+        multi_gmc: Update multiple track states using a homography matrix.
+        activate: Activate a new tracklet.
+        re_activate: Reactivate a previously lost tracklet.
+        update: Update the state of a matched track.
+        convert_coords: Convert bounding box to x-y-aspect-height format.
+        tlwh_to_xyah: Convert tlwh bounding box to xyah format.
 
     Examples:
         Initialize and activate a new track
@@ -48,13 +52,12 @@ class STrack(BaseTrack):
 
     shared_kalman = KalmanFilterXYAH()
 
-    def __init__(self, xywh, score, cls):
-        """
-        Initialize a new STrack instance.
+    def __init__(self, xywh: list[float], score: float, cls: Any):
+        """Initialize a new STrack instance.
 
         Args:
-            xywh (List[float]): Bounding box coordinates and dimensions in the format (x, y, w, h, [a], idx), where
-                (x, y) is the center, (w, h) are width and height, [a] is optional aspect ratio, and idx is the id.
+            xywh (list[float]): Bounding box coordinates and dimensions in the format (x, y, w, h, [a], idx), where (x,
+                y) is the center, (w, h) are width and height, [a] is optional aspect ratio, and idx is the id.
             score (float): Confidence score of the detection.
             cls (Any): Class label for the detected object.
 
@@ -79,14 +82,14 @@ class STrack(BaseTrack):
         self.angle = xywh[4] if len(xywh) == 6 else None
 
     def predict(self):
-        """Predicts the next state (mean and covariance) of the object using the Kalman filter."""
+        """Predict the next state (mean and covariance) of the object using the Kalman filter."""
         mean_state = self.mean.copy()
         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):
+    def multi_predict(stracks: list[STrack]):
         """Perform multi-object predictive tracking using Kalman filter for the provided list of STrack instances."""
         if len(stracks) <= 0:
             return
@@ -101,9 +104,9 @@ class STrack(BaseTrack):
             stracks[i].covariance = cov
 
     @staticmethod
-    def multi_gmc(stracks, H=np.eye(2, 3)):
+    def multi_gmc(stracks: list[STrack], H: np.ndarray = np.eye(2, 3)):
         """Update state tracks positions and covariances using a homography matrix for multiple tracks."""
-        if len(stracks) > 0:
+        if stracks:
             multi_mean = np.asarray([st.mean.copy() for st in stracks])
             multi_covariance = np.asarray([st.covariance for st in stracks])
 
@@ -119,7 +122,7 @@ class STrack(BaseTrack):
                 stracks[i].mean = mean
                 stracks[i].covariance = cov
 
-    def activate(self, kalman_filter, frame_id):
+    def activate(self, kalman_filter: KalmanFilterXYAH, frame_id: int):
         """Activate a new tracklet using the provided Kalman filter and initialize its state and covariance."""
         self.kalman_filter = kalman_filter
         self.track_id = self.next_id()
@@ -132,8 +135,8 @@ class STrack(BaseTrack):
         self.frame_id = frame_id
         self.start_frame = frame_id
 
-    def re_activate(self, new_track, frame_id, new_id=False):
-        """Reactivates a previously lost track using new detection data and updates its state and attributes."""
+    def re_activate(self, new_track: STrack, frame_id: int, new_id: bool = False):
+        """Reactivate a previously lost track using new detection data and update its state and attributes."""
         self.mean, self.covariance = self.kalman_filter.update(
             self.mean, self.covariance, self.convert_coords(new_track.tlwh)
         )
@@ -148,9 +151,8 @@ class STrack(BaseTrack):
         self.angle = new_track.angle
         self.idx = new_track.idx
 
-    def update(self, new_track, frame_id):
-        """
-        Update the state of a matched track.
+    def update(self, new_track: STrack, frame_id: int):
+        """Update the state of a matched track.
 
         Args:
             new_track (STrack): The new track containing updated information.
@@ -177,13 +179,13 @@ class STrack(BaseTrack):
         self.angle = new_track.angle
         self.idx = new_track.idx
 
-    def convert_coords(self, tlwh):
+    def convert_coords(self, tlwh: np.ndarray) -> np.ndarray:
         """Convert a bounding box's top-left-width-height format to its x-y-aspect-height equivalent."""
         return self.tlwh_to_xyah(tlwh)
 
     @property
-    def tlwh(self):
-        """Returns the bounding box in top-left-width-height format from the current state estimate."""
+    def tlwh(self) -> np.ndarray:
+        """Get the bounding box in top-left-width-height format from the current state estimate."""
         if self.mean is None:
             return self._tlwh.copy()
         ret = self.mean[:4].copy()
@@ -192,14 +194,14 @@ class STrack(BaseTrack):
         return ret
 
     @property
-    def xyxy(self):
-        """Converts bounding box from (top left x, top left y, width, height) to (min x, min y, max x, max y) format."""
+    def xyxy(self) -> np.ndarray:
+        """Convert bounding box from (top left x, top left y, width, height) to (min x, min y, max x, max y) format."""
         ret = self.tlwh.copy()
         ret[2:] += ret[:2]
         return ret
 
     @staticmethod
-    def tlwh_to_xyah(tlwh):
+    def tlwh_to_xyah(tlwh: np.ndarray) -> np.ndarray:
         """Convert bounding box from tlwh format to center-x-center-y-aspect-height (xyah) format."""
         ret = np.asarray(tlwh).copy()
         ret[:2] += ret[2:] / 2
@@ -207,58 +209,58 @@ class STrack(BaseTrack):
         return ret
 
     @property
-    def xywh(self):
-        """Returns the current position of the bounding box in (center x, center y, width, height) format."""
+    def xywh(self) -> np.ndarray:
+        """Get the current position of the bounding box in (center x, center y, width, height) format."""
         ret = np.asarray(self.tlwh).copy()
         ret[:2] += ret[2:] / 2
         return ret
 
     @property
-    def xywha(self):
-        """Returns position in (center x, center y, width, height, angle) format, warning if angle is missing."""
+    def xywha(self) -> np.ndarray:
+        """Get position in (center x, center y, width, height, angle) format, warning if angle is missing."""
         if self.angle is None:
             LOGGER.warning("`angle` attr not found, returning `xywh` instead.")
             return self.xywh
         return np.concatenate([self.xywh, self.angle[None]])
 
     @property
-    def result(self):
-        """Returns the current tracking results in the appropriate bounding box format."""
+    def result(self) -> list[float]:
+        """Get the current tracking results in the appropriate bounding box format."""
         coords = self.xyxy if self.angle is None else self.xywha
-        return coords.tolist() + [self.track_id, self.score, self.cls, self.idx]
+        return [*coords.tolist(), self.track_id, self.score, self.cls, self.idx]
 
-    def __repr__(self):
-        """Returns a string representation of the STrack object including start frame, end frame, and track ID."""
+    def __repr__(self) -> str:
+        """Return a string representation of the STrack object including start frame, end frame, and track ID."""
         return f"OT_{self.track_id}_({self.start_frame}-{self.end_frame})"
 
 
 class BYTETracker:
-    """
-    BYTETracker: A tracking algorithm built on top of YOLOv8 for object detection and tracking.
+    """BYTETracker: A tracking algorithm built on top of YOLOv8 for object detection and tracking.
 
-    This class encapsulates the functionality for initializing, updating, and managing the tracks for detected objects in a
-    video sequence. It maintains the state of tracked, lost, and removed tracks over frames, utilizes Kalman filtering for
-    predicting the new object locations, and performs data association.
+    This class encapsulates the functionality for initializing, updating, and managing the tracks for detected objects
+    in a video sequence. It maintains the state of tracked, lost, and removed tracks over frames, utilizes Kalman
+    filtering for predicting the new object locations, and performs data association.
 
     Attributes:
-        tracked_stracks (List[STrack]): List of successfully activated tracks.
-        lost_stracks (List[STrack]): List of lost tracks.
-        removed_stracks (List[STrack]): List of removed tracks.
+        tracked_stracks (list[STrack]): List of successfully activated tracks.
+        lost_stracks (list[STrack]): List of lost tracks.
+        removed_stracks (list[STrack]): List of removed tracks.
         frame_id (int): The current frame ID.
         args (Namespace): Command-line arguments.
         max_time_lost (int): The maximum frames for a track to be considered as 'lost'.
         kalman_filter (KalmanFilterXYAH): Kalman Filter object.
 
     Methods:
-        update(results, img=None): Updates object tracker with new detections.
-        get_kalmanfilter(): Returns a Kalman filter object for tracking bounding boxes.
-        init_track(dets, scores, cls, img=None): Initialize object tracking with detections.
-        get_dists(tracks, detections): Calculates the distance between tracks and detections.
-        multi_predict(tracks): Predicts the location of tracks.
-        reset_id(): Resets the ID counter of STrack.
-        joint_stracks(tlista, tlistb): Combines two lists of stracks.
-        sub_stracks(tlista, tlistb): Filters out the stracks present in the second list from the first list.
-        remove_duplicate_stracks(stracksa, stracksb): Removes duplicate stracks based on IoU.
+        update: Update object tracker with new detections.
+        get_kalmanfilter: Return a Kalman filter object for tracking bounding boxes.
+        init_track: Initialize object tracking with detections.
+        get_dists: Calculate the distance between tracks and detections.
+        multi_predict: Predict the location of tracks.
+        reset_id: Reset the ID counter of STrack.
+        reset: Reset the tracker by clearing all tracks.
+        joint_stracks: Combine two lists of stracks.
+        sub_stracks: Filter out the stracks present in the second list from the first list.
+        remove_duplicate_stracks: Remove duplicate stracks based on IoU.
 
     Examples:
         Initialize BYTETracker and update with detection results
@@ -267,9 +269,8 @@ class BYTETracker:
         >>> tracked_objects = tracker.update(results)
     """
 
-    def __init__(self, args, frame_rate=30):
-        """
-        Initialize a BYTETracker instance for object tracking.
+    def __init__(self, args, frame_rate: int = 30):
+        """Initialize a BYTETracker instance for object tracking.
 
         Args:
             args (Namespace): Command-line arguments containing tracking parameters.
@@ -290,8 +291,8 @@ class BYTETracker:
         self.kalman_filter = self.get_kalmanfilter()
         self.reset_id()
 
-    def update(self, results, img=None, feats=None):
-        """Updates the tracker with new detections and returns the current list of tracked objects."""
+    def update(self, results, img: np.ndarray | None = None, feats: np.ndarray | None = None) -> np.ndarray:
+        """Update the tracker with new detections and return the current list of tracked objects."""
         self.frame_id += 1
         activated_stracks = []
         refind_stracks = []
@@ -299,24 +300,19 @@ class BYTETracker:
         removed_stracks = []
 
         scores = results.conf
-        bboxes = results.xywhr if hasattr(results, "xywhr") else results.xywh
-        # Add index
-        bboxes = np.concatenate([bboxes, np.arange(len(bboxes)).reshape(-1, 1)], axis=-1)
-        cls = results.cls
-
         remain_inds = scores >= self.args.track_high_thresh
         inds_low = scores > self.args.track_low_thresh
         inds_high = scores < self.args.track_high_thresh
 
         inds_second = inds_low & inds_high
-        dets_second = bboxes[inds_second]
-        dets = bboxes[remain_inds]
-        scores_keep = scores[remain_inds]
-        scores_second = scores[inds_second]
-        cls_keep = cls[remain_inds]
-        cls_second = cls[inds_second]
-
-        detections = self.init_track(dets, scores_keep, cls_keep, img if feats is None else feats)
+        results_second = results[inds_second]
+        results = results[remain_inds]
+        feats_keep = feats_second = img
+        if feats is not None and len(feats):
+            feats_keep = feats[remain_inds]
+            feats_second = feats[inds_second]
+
+        detections = self.init_track(results, feats_keep)
         # Add newly detected tracklets to tracked_stracks
         unconfirmed = []
         tracked_stracks = []  # type: list[STrack]
@@ -332,7 +328,7 @@ class BYTETracker:
         if hasattr(self, "gmc") and img is not None:
             # use try-except here to bypass errors from gmc module
             try:
-                warp = self.gmc.apply(img, dets)
+                warp = self.gmc.apply(img, results.xyxy)
             except Exception:
                 warp = np.eye(2, 3)
             STrack.multi_gmc(strack_pool, warp)
@@ -351,11 +347,11 @@ class BYTETracker:
                 track.re_activate(det, self.frame_id, new_id=False)
                 refind_stracks.append(track)
         # Step 3: Second association, with low score detection boxes association the untrack to the low score detections
-        detections_second = self.init_track(dets_second, scores_second, cls_second, img if feats is None else feats)
+        detections_second = self.init_track(results_second, feats_second)
         r_tracked_stracks = [strack_pool[i] for i in u_track if strack_pool[i].state == TrackState.Tracked]
         # TODO
         dists = matching.iou_distance(r_tracked_stracks, detections_second)
-        matches, u_track, u_detection_second = matching.linear_assignment(dists, thresh=0.5)
+        matches, u_track, _u_detection_second = matching.linear_assignment(dists, thresh=0.5)
         for itracked, idet in matches:
             track = r_tracked_stracks[itracked]
             det = detections_second[idet]
@@ -408,32 +404,36 @@ class BYTETracker:
 
         return np.asarray([x.result for x in self.tracked_stracks if x.is_activated], dtype=np.float32)
 
-    def get_kalmanfilter(self):
-        """Returns a Kalman filter object for tracking bounding boxes using KalmanFilterXYAH."""
+    def get_kalmanfilter(self) -> KalmanFilterXYAH:
+        """Return a Kalman filter object for tracking bounding boxes using KalmanFilterXYAH."""
         return KalmanFilterXYAH()
 
-    def init_track(self, dets, scores, cls, img=None):
-        """Initializes object tracking with given detections, scores, and class labels using the STrack algorithm."""
-        return [STrack(xyxy, s, c) for (xyxy, s, c) in zip(dets, scores, cls)] if len(dets) else []  # detections
+    def init_track(self, results, img: np.ndarray | None = None) -> list[STrack]:
+        """Initialize object tracking with given detections, scores, and class labels using the STrack algorithm."""
+        if len(results) == 0:
+            return []
+        bboxes = results.xywhr if hasattr(results, "xywhr") else results.xywh
+        bboxes = np.concatenate([bboxes, np.arange(len(bboxes)).reshape(-1, 1)], axis=-1)
+        return [STrack(xywh, s, c) for (xywh, s, c) in zip(bboxes, results.conf, results.cls)]
 
-    def get_dists(self, tracks, detections):
-        """Calculates the distance between tracks and detections using IoU and optionally fuses scores."""
+    def get_dists(self, tracks: list[STrack], detections: list[STrack]) -> np.ndarray:
+        """Calculate the distance between tracks and detections using IoU and optionally fuse scores."""
         dists = matching.iou_distance(tracks, detections)
         if self.args.fuse_score:
             dists = matching.fuse_score(dists, detections)
         return dists
 
-    def multi_predict(self, tracks):
+    def multi_predict(self, tracks: list[STrack]):
         """Predict the next states for multiple tracks using Kalman filter."""
         STrack.multi_predict(tracks)
 
     @staticmethod
     def reset_id():
-        """Resets the ID counter for STrack instances to ensure unique track IDs across tracking sessions."""
+        """Reset the ID counter for STrack instances to ensure unique track IDs across tracking sessions."""
         STrack.reset_id()
 
     def reset(self):
-        """Resets the tracker by clearing all tracked, lost, and removed tracks and reinitializing the Kalman filter."""
+        """Reset the tracker by clearing all tracked, lost, and removed tracks and reinitializing the Kalman filter."""
         self.tracked_stracks = []  # type: list[STrack]
         self.lost_stracks = []  # type: list[STrack]
         self.removed_stracks = []  # type: list[STrack]
@@ -442,8 +442,8 @@ class BYTETracker:
         self.reset_id()
 
     @staticmethod
-    def joint_stracks(tlista, tlistb):
-        """Combines two lists of STrack objects into a single list, ensuring no duplicates based on track IDs."""
+    def joint_stracks(tlista: list[STrack], tlistb: list[STrack]) -> list[STrack]:
+        """Combine two lists of STrack objects into a single list, ensuring no duplicates based on track IDs."""
         exists = {}
         res = []
         for t in tlista:
@@ -457,14 +457,14 @@ class BYTETracker:
         return res
 
     @staticmethod
-    def sub_stracks(tlista, tlistb):
-        """Filters out the stracks present in the second list from the first list."""
+    def sub_stracks(tlista: list[STrack], tlistb: list[STrack]) -> list[STrack]:
+        """Filter out the stracks present in the second list from the first list."""
         track_ids_b = {t.track_id for t in tlistb}
         return [t for t in tlista if t.track_id not in track_ids_b]
 
     @staticmethod
-    def remove_duplicate_stracks(stracksa, stracksb):
-        """Removes duplicate stracks from two lists based on Intersection over Union (IoU) distance."""
+    def remove_duplicate_stracks(stracksa: list[STrack], stracksb: list[STrack]) -> tuple[list[STrack], list[STrack]]:
+        """Remove duplicate stracks from two lists based on Intersection over Union (IoU) distance."""
         pdist = matching.iou_distance(stracksa, stracksb)
         pairs = np.where(pdist < 0.15)
         dupa, dupb = [], []

ultralytics/trackers/track.py

@@ -16,19 +16,14 @@ TRACKER_MAP = {"bytetrack": BYTETracker, "botsort": BOTSORT}
 
 
 def on_predict_start(predictor: object, persist: bool = False) -> None:
-    """
-    Initialize trackers for object tracking during prediction.
+    """Initialize trackers for object tracking during prediction.
 
     Args:
-        predictor (object): The predictor object to initialize trackers for.
-        persist (bool): Whether to persist the trackers if they already exist.
-
-    Raises:
-        AssertionError: If the tracker_type is not 'bytetrack' or 'botsort'.
-        ValueError: If the task is 'classify' as classification doesn't support tracking.
+        predictor (ultralytics.engine.predictor.BasePredictor): The predictor object to initialize trackers for.
+        persist (bool, optional): Whether to persist the trackers if they already exist.
 
     Examples:
-        Initialize trackers for a predictor object:
+        Initialize trackers for a predictor object
         >>> predictor = SomePredictorClass()
         >>> on_predict_start(predictor, persist=True)
     """
@@ -74,12 +69,11 @@ def on_predict_start(predictor: object, persist: bool = False) -> None:
 
 
 def on_predict_postprocess_end(predictor: object, persist: bool = False) -> None:
-    """
-    Postprocess detected boxes and update with object tracking.
+    """Postprocess detected boxes and update with object tracking.
 
     Args:
         predictor (object): The predictor object containing the predictions.
-        persist (bool): Whether to persist the trackers if they already exist.
+        persist (bool, optional): Whether to persist the trackers if they already exist.
 
     Examples:
         Postprocess predictions and update with tracking
@@ -96,8 +90,6 @@ def on_predict_postprocess_end(predictor: object, persist: bool = False) -> None
             predictor.vid_path[i if is_stream else 0] = vid_path
 
         det = (result.obb if is_obb else result.boxes).cpu().numpy()
-        if len(det) == 0:
-            continue
         tracks = tracker.update(det, result.orig_img, getattr(result, "feats", None))
         if len(tracks) == 0:
             continue
@@ -109,8 +101,7 @@ def on_predict_postprocess_end(predictor: object, persist: bool = False) -> None
 
 
 def register_tracker(model: object, persist: bool) -> None:
-    """
-    Register tracking callbacks to the model for object tracking during prediction.
+    """Register tracking callbacks to the model for object tracking during prediction.
 
     Args:
         model (object): The model object to register tracking callbacks for.