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

ultralytics/trackers/basetrack.py

@@ -2,13 +2,13 @@
 """Module defines the base classes and structures for object tracking in YOLO."""
 
 from collections import OrderedDict
+from typing import Any
 
 import numpy as np
 
 
 class TrackState:
-    """
-    Enumeration class representing the possible states of an object being tracked.
+    """Enumeration class representing the possible states of an object being tracked.
 
     Attributes:
         New (int): State when the object is newly detected.
@@ -29,8 +29,7 @@ class TrackState:
 
 
 class BaseTrack:
-    """
-    Base class for object tracking, providing foundational attributes and methods.
+    """Base class for object tracking, providing foundational attributes and methods.
 
     Attributes:
         _count (int): Class-level counter for unique track IDs.
@@ -66,15 +65,7 @@ class BaseTrack:
     _count = 0
 
     def __init__(self):
-        """
-        Initialize a new track with a unique ID and foundational tracking attributes.
-
-        Examples:
-            Initialize a new track
-            >>> track = BaseTrack()
-            >>> print(track.track_id)
-            0
-        """
+        """Initialize a new track with a unique ID and foundational tracking attributes."""
         self.track_id = 0
         self.is_activated = False
         self.state = TrackState.New
@@ -88,37 +79,37 @@ class BaseTrack:
         self.location = (np.inf, np.inf)
 
     @property
-    def end_frame(self):
-        """Returns the ID of the most recent frame where the object was tracked."""
+    def end_frame(self) -> int:
+        """Return the ID of the most recent frame where the object was tracked."""
         return self.frame_id
 
     @staticmethod
-    def next_id():
+    def next_id() -> int:
         """Increment and return the next unique global track ID for object tracking."""
         BaseTrack._count += 1
         return BaseTrack._count
 
-    def activate(self, *args):
-        """Activates the track with provided arguments, initializing necessary attributes for tracking."""
+    def activate(self, *args: Any) -> None:
+        """Activate the track with provided arguments, initializing necessary attributes for tracking."""
         raise NotImplementedError
 
-    def predict(self):
-        """Predicts the next state of the track based on the current state and tracking model."""
+    def predict(self) -> None:
+        """Predict the next state of the track based on the current state and tracking model."""
         raise NotImplementedError
 
-    def update(self, *args, **kwargs):
-        """Updates the track with new observations and data, modifying its state and attributes accordingly."""
+    def update(self, *args: Any, **kwargs: Any) -> None:
+        """Update the track with new observations and data, modifying its state and attributes accordingly."""
         raise NotImplementedError
 
-    def mark_lost(self):
-        """Marks the track as lost by updating its state to TrackState.Lost."""
+    def mark_lost(self) -> None:
+        """Mark the track as lost by updating its state to TrackState.Lost."""
         self.state = TrackState.Lost
 
-    def mark_removed(self):
-        """Marks the track as removed by setting its state to TrackState.Removed."""
+    def mark_removed(self) -> None:
+        """Mark the track as removed by setting its state to TrackState.Removed."""
         self.state = TrackState.Removed
 
     @staticmethod
-    def reset_id():
+    def reset_id() -> None:
         """Reset the global track ID counter to its initial value."""
         BaseTrack._count = 0

ultralytics/trackers/bot_sort.py

@@ -1,6 +1,9 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+from __future__ import annotations
+
 from collections import deque
+from typing import Any
 
 import numpy as np
 import torch
@@ -16,8 +19,7 @@ from .utils.kalman_filter import KalmanFilterXYWH
 
 
 class BOTrack(STrack):
-    """
-    An extended version of the STrack class for YOLO, adding object tracking features.
+    """An extended version of the STrack class for YOLO, adding object tracking features.
 
     This class extends the STrack class to include additional functionalities for object tracking, such as feature
     smoothing, Kalman filter prediction, and reactivation of tracks.
@@ -51,26 +53,27 @@ class BOTrack(STrack):
 
     shared_kalman = KalmanFilterXYWH()
 
-    def __init__(self, tlwh, score, cls, feat=None, feat_history=50):
-        """
-        Initialize a BOTrack object with temporal parameters, such as feature history, alpha, and current features.
+    def __init__(
+        self, xywh: np.ndarray, score: float, cls: int, feat: np.ndarray | None = None, feat_history: int = 50
+    ):
+        """Initialize a BOTrack object with temporal parameters, such as feature history, alpha, and current features.
 
         Args:
-            tlwh (np.ndarray): Bounding box coordinates in tlwh format (top left x, top left y, width, height).
+            xywh (np.ndarray): Bounding box coordinates in xywh format (center x, center y, width, height).
             score (float): Confidence score of the detection.
             cls (int): Class ID of the detected object.
-            feat (np.ndarray | None): Feature vector associated with the detection.
+            feat (np.ndarray, optional): Feature vector associated with the detection.
             feat_history (int): Maximum length of the feature history deque.
 
         Examples:
             Initialize a BOTrack object with bounding box, score, class ID, and feature vector
-            >>> tlwh = np.array([100, 50, 80, 120])
+            >>> xywh = np.array([100, 150, 60, 50])
             >>> score = 0.9
             >>> cls = 1
             >>> feat = np.random.rand(128)
-            >>> bo_track = BOTrack(tlwh, score, cls, feat)
+            >>> bo_track = BOTrack(xywh, score, cls, feat)
         """
-        super().__init__(tlwh, score, cls)
+        super().__init__(xywh, score, cls)
 
         self.smooth_feat = None
         self.curr_feat = None
@@ -79,7 +82,7 @@ class BOTrack(STrack):
         self.features = deque([], maxlen=feat_history)
         self.alpha = 0.9
 
-    def update_features(self, feat):
+    def update_features(self, feat: np.ndarray) -> None:
         """Update the feature vector and apply exponential moving average smoothing."""
         feat /= np.linalg.norm(feat)
         self.curr_feat = feat
@@ -90,7 +93,7 @@ class BOTrack(STrack):
         self.features.append(feat)
         self.smooth_feat /= np.linalg.norm(self.smooth_feat)
 
-    def predict(self):
+    def predict(self) -> None:
         """Predict the object's future state using the Kalman filter to update its mean and covariance."""
         mean_state = self.mean.copy()
         if self.state != TrackState.Tracked:
@@ -99,20 +102,20 @@ class BOTrack(STrack):
 
         self.mean, self.covariance = self.kalman_filter.predict(mean_state, self.covariance)
 
-    def re_activate(self, new_track, frame_id, new_id=False):
+    def re_activate(self, new_track: BOTrack, frame_id: int, new_id: bool = False) -> None:
         """Reactivate a track with updated features and optionally assign a new ID."""
         if new_track.curr_feat is not None:
             self.update_features(new_track.curr_feat)
         super().re_activate(new_track, frame_id, new_id)
 
-    def update(self, new_track, frame_id):
+    def update(self, new_track: BOTrack, frame_id: int) -> None:
         """Update the track with new detection information and the current frame ID."""
         if new_track.curr_feat is not None:
             self.update_features(new_track.curr_feat)
         super().update(new_track, frame_id)
 
     @property
-    def tlwh(self):
+    def tlwh(self) -> np.ndarray:
         """Return the current bounding box position in `(top left x, top left y, width, height)` format."""
         if self.mean is None:
             return self._tlwh.copy()
@@ -121,7 +124,7 @@ class BOTrack(STrack):
         return ret
 
     @staticmethod
-    def multi_predict(stracks):
+    def multi_predict(stracks: list[BOTrack]) -> None:
         """Predict the mean and covariance for multiple object tracks using a shared Kalman filter."""
         if len(stracks) <= 0:
             return
@@ -136,12 +139,12 @@ class BOTrack(STrack):
             stracks[i].mean = mean
             stracks[i].covariance = cov
 
-    def convert_coords(self, tlwh):
+    def convert_coords(self, tlwh: np.ndarray) -> np.ndarray:
         """Convert tlwh bounding box coordinates to xywh format."""
         return self.tlwh_to_xywh(tlwh)
 
     @staticmethod
-    def tlwh_to_xywh(tlwh):
+    def tlwh_to_xywh(tlwh: np.ndarray) -> np.ndarray:
         """Convert bounding box from tlwh (top-left-width-height) to xywh (center-x-center-y-width-height) format."""
         ret = np.asarray(tlwh).copy()
         ret[:2] += ret[2:] / 2
@@ -149,8 +152,7 @@ class BOTrack(STrack):
 
 
 class BOTSORT(BYTETracker):
-    """
-    An extended version of the BYTETracker class for YOLO, designed for object tracking with ReID and GMC algorithm.
+    """An extended version of the BYTETracker class for YOLO, designed for object tracking with ReID and GMC algorithm.
 
     Attributes:
         proximity_thresh (float): Threshold for spatial proximity (IoU) between tracks and detections.
@@ -172,16 +174,15 @@ class BOTSORT(BYTETracker):
         >>> bot_sort.init_track(dets, scores, cls, img)
         >>> bot_sort.multi_predict(tracks)
 
-    Note:
+    Notes:
         The class is designed to work with a YOLO object detection model and supports ReID only if enabled via args.
     """
 
-    def __init__(self, args, frame_rate=30):
-        """
-        Initialize BOTSORT object with ReID module and GMC algorithm.
+    def __init__(self, args: Any, frame_rate: int = 30):
+        """Initialize BOTSORT object with ReID module and GMC algorithm.
 
         Args:
-            args (object): Parsed command-line arguments containing tracking parameters.
+            args (Any): Parsed command-line arguments containing tracking parameters.
             frame_rate (int): Frame rate of the video being processed.
 
         Examples:
@@ -203,21 +204,23 @@ class BOTSORT(BYTETracker):
             else None
         )
 
-    def get_kalmanfilter(self):
+    def get_kalmanfilter(self) -> KalmanFilterXYWH:
         """Return an instance of KalmanFilterXYWH for predicting and updating object states in the tracking process."""
         return KalmanFilterXYWH()
 
-    def init_track(self, dets, scores, cls, img=None):
+    def init_track(self, results, img: np.ndarray | None = None) -> list[BOTrack]:
         """Initialize object tracks using detection bounding boxes, scores, class labels, and optional ReID features."""
-        if len(dets) == 0:
+        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)
         if self.args.with_reid and self.encoder is not None:
-            features_keep = self.encoder(img, dets)
-            return [BOTrack(xyxy, s, c, f) for (xyxy, s, c, f) in zip(dets, scores, cls, features_keep)]  # detections
+            features_keep = self.encoder(img, bboxes)
+            return [BOTrack(xywh, s, c, f) for (xywh, s, c, f) in zip(bboxes, results.conf, results.cls, features_keep)]
         else:
-            return [BOTrack(xyxy, s, c) for (xyxy, s, c) in zip(dets, scores, cls)]  # detections
+            return [BOTrack(xywh, s, c) for (xywh, s, c) in zip(bboxes, results.conf, results.cls)]
 
-    def get_dists(self, tracks, detections):
+    def get_dists(self, tracks: list[BOTrack], detections: list[BOTrack]) -> np.ndarray:
         """Calculate distances between tracks and detections using IoU and optionally ReID embeddings."""
         dists = matching.iou_distance(tracks, detections)
         dists_mask = dists > (1 - self.proximity_thresh)
@@ -232,11 +235,11 @@ class BOTSORT(BYTETracker):
             dists = np.minimum(dists, emb_dists)
         return dists
 
-    def multi_predict(self, tracks):
+    def multi_predict(self, tracks: list[BOTrack]) -> None:
         """Predict the mean and covariance of multiple object tracks using a shared Kalman filter."""
         BOTrack.multi_predict(tracks)
 
-    def reset(self):
+    def reset(self) -> None:
         """Reset the BOTSORT tracker to its initial state, clearing all tracked objects and internal states."""
         super().reset()
         self.gmc.reset_params()
@@ -245,16 +248,22 @@ class BOTSORT(BYTETracker):
 class ReID:
     """YOLO model as encoder for re-identification."""
 
-    def __init__(self, model):
-        """Initialize encoder for re-identification."""
+    def __init__(self, model: str):
+        """Initialize encoder for re-identification.
+
+        Args:
+            model (str): Path to the YOLO model for re-identification.
+        """
         from ultralytics import YOLO
 
         self.model = YOLO(model)
-        self.model(embed=[len(self.model.model.model) - 2 if ".pt" in model else -1], verbose=False)  # initialize
+        self.model(embed=[len(self.model.model.model) - 2 if ".pt" in model else -1], verbose=False, save=False)  # init
 
-    def __call__(self, img, dets):
+    def __call__(self, img: np.ndarray, dets: np.ndarray) -> list[np.ndarray]:
         """Extract embeddings for detected objects."""
-        feats = self.model([save_one_box(det, img, save=False) for det in xywh2xyxy(torch.from_numpy(dets[:, :4]))])
+        feats = self.model.predictor(
+            [save_one_box(det, img, save=False) for det in xywh2xyxy(torch.from_numpy(dets[:, :4]))]
+        )
         if len(feats) != dets.shape[0] and feats[0].shape[0] == dets.shape[0]:
             feats = feats[0]  # batched prediction with non-PyTorch backend
         return [f.cpu().numpy() for f in feats]