ultralytics 8.1.29__py3-none-any.whl → 8.3.63__py3-none-any.whl

ultralytics/solutions/trackzone.py

@@ -0,0 +1,68 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+import cv2
+import numpy as np
+
+from ultralytics.solutions.solutions import BaseSolution
+from ultralytics.utils.plotting import Annotator, colors
+
+
+class TrackZone(BaseSolution):
+    """
+    A class to manage region-based object tracking in a video stream.
+
+    This class extends the BaseSolution class and provides functionality for tracking objects within a specific region
+    defined by a polygonal area. Objects outside the region are excluded from tracking. It supports dynamic initialization
+    of the region, allowing either a default region or a user-specified polygon.
+
+    Attributes:
+        region (ndarray): The polygonal region for tracking, represented as a convex hull.
+
+    Methods:
+        trackzone: Processes each frame of the video, applying region-based tracking.
+
+    Examples:
+        >>> tracker = TrackZone()
+        >>> frame = cv2.imread("frame.jpg")
+        >>> processed_frame = tracker.trackzone(frame)
+        >>> cv2.imshow("Tracked Frame", processed_frame)
+    """
+
+    def __init__(self, **kwargs):
+        """Initializes the TrackZone class for tracking objects within a defined region in video streams."""
+        super().__init__(**kwargs)
+        default_region = [(150, 150), (1130, 150), (1130, 570), (150, 570)]
+        self.region = cv2.convexHull(np.array(self.region or default_region, dtype=np.int32))
+
+    def trackzone(self, im0):
+        """
+        Processes the input frame to track objects within a defined region.
+
+        This method initializes the annotator, creates a mask for the specified region, extracts tracks
+        only from the masked area, and updates tracking information. Objects outside the region are ignored.
+
+        Args:
+            im0 (numpy.ndarray): The input image or frame to be processed.
+
+        Returns:
+            (numpy.ndarray): The processed image with tracking id and bounding boxes annotations.
+
+        Examples:
+            >>> tracker = TrackZone()
+            >>> frame = cv2.imread("path/to/image.jpg")
+            >>> tracker.trackzone(frame)
+        """
+        self.annotator = Annotator(im0, line_width=self.line_width)  # Initialize annotator
+        # Create a mask for the region and extract tracks from the masked image
+        masked_frame = cv2.bitwise_and(im0, im0, mask=cv2.fillPoly(np.zeros_like(im0[:, :, 0]), [self.region], 255))
+        self.extract_tracks(masked_frame)
+
+        cv2.polylines(im0, [self.region], isClosed=True, color=(255, 255, 255), thickness=self.line_width * 2)
+
+        # Iterate over boxes, track ids, classes indexes list and draw bounding boxes
+        for box, track_id, cls in zip(self.boxes, self.track_ids, self.clss):
+            self.annotator.box_label(box, label=f"{self.names[cls]}:{track_id}", color=colors(track_id, True))
+
+        self.display_output(im0)  # display output with base class function
+
+        return im0  # return output image for more usage

ultralytics/trackers/__init__.py

@@ -1,4 +1,4 @@
-# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 from .bot_sort import BOTSORT
 from .byte_tracker import BYTETracker

ultralytics/trackers/basetrack.py

@@ -1,5 +1,5 @@
-# Ultralytics YOLO 🚀, AGPL-3.0 license
-"""This module defines the base classes and structures for object tracking in YOLO."""
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+"""Module defines the base classes and structures for object tracking in YOLO."""
 
 from collections import OrderedDict
 
@@ -15,6 +15,11 @@ class TrackState:
         Tracked (int): State when the object is successfully tracked in subsequent frames.
         Lost (int): State when the object is no longer tracked.
         Removed (int): State when the object is removed from tracking.
+
+    Examples:
+        >>> state = TrackState.New
+        >>> if state == TrackState.New:
+        >>>     print("Object is newly detected.")
     """
 
     New = 0
@@ -33,8 +38,8 @@ class BaseTrack:
         is_activated (bool): Flag indicating whether the track is currently active.
         state (TrackState): Current state of the track.
         history (OrderedDict): Ordered history of the track's states.
-        features (list): List of features extracted from the object for tracking.
-        curr_feature (any): The current feature of the object being tracked.
+        features (List): List of features extracted from the object for tracking.
+        curr_feature (Any): The current feature of the object being tracked.
         score (float): The confidence score of the tracking.
         start_frame (int): The frame number where tracking started.
         frame_id (int): The most recent frame ID processed by the track.
@@ -50,12 +55,26 @@ class BaseTrack:
         mark_lost: Marks the track as lost.
         mark_removed: Marks the track as removed.
         reset_id: Resets the global track ID counter.
+
+    Examples:
+        Initialize a new track and mark it as lost:
+        >>> track = BaseTrack()
+        >>> track.mark_lost()
+        >>> print(track.state)  # Output: 2 (TrackState.Lost)
     """
 
     _count = 0
 
     def __init__(self):
-        """Initializes a new track with unique ID and foundational tracking attributes."""
+        """
+        Initializes a new track with a unique ID and foundational tracking attributes.
+
+        Examples:
+            Initialize a new track
+            >>> track = BaseTrack()
+            >>> print(track.track_id)
+            0
+        """
         self.track_id = 0
         self.is_activated = False
         self.state = TrackState.New
@@ -70,36 +89,36 @@ class BaseTrack:
 
     @property
     def end_frame(self):
-        """Return the last frame ID of the track."""
+        """Returns the ID of the most recent frame where the object was tracked."""
         return self.frame_id
 
     @staticmethod
     def next_id():
-        """Increment and return the global track ID counter."""
+        """Increment and return the next unique global track ID for object tracking."""
         BaseTrack._count += 1
         return BaseTrack._count
 
     def activate(self, *args):
-        """Abstract method to activate the track with provided arguments."""
+        """Activates the track with provided arguments, initializing necessary attributes for tracking."""
         raise NotImplementedError
 
     def predict(self):
-        """Abstract method to predict the next state of the track."""
+        """Predicts the next state of the track based on the current state and tracking model."""
         raise NotImplementedError
 
     def update(self, *args, **kwargs):
-        """Abstract method to update the track with new observations."""
+        """Updates the track with new observations and data, modifying its state and attributes accordingly."""
         raise NotImplementedError
 
     def mark_lost(self):
-        """Mark the track as lost."""
+        """Marks the track as lost by updating its state to TrackState.Lost."""
         self.state = TrackState.Lost
 
     def mark_removed(self):
-        """Mark the track as removed."""
+        """Marks the track as removed by setting its state to TrackState.Removed."""
         self.state = TrackState.Removed
 
     @staticmethod
     def reset_id():
-        """Reset the global track ID counter."""
+        """Reset the global track ID counter to its initial value."""
         BaseTrack._count = 0

ultralytics/trackers/bot_sort.py

@@ -1,4 +1,4 @@
-# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 from collections import deque
 
@@ -15,6 +15,9 @@ class BOTrack(STrack):
     """
     An extended version of the STrack class for YOLOv8, 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.
+
     Attributes:
         shared_kalman (KalmanFilterXYWH): A shared Kalman filter for all instances of BOTrack.
         smooth_feat (np.ndarray): Smoothed feature vector.
@@ -34,16 +37,35 @@ class BOTrack(STrack):
         convert_coords(tlwh): Converts tlwh bounding box coordinates to xywh format.
         tlwh_to_xywh(tlwh): Convert bounding box to xywh format `(center x, center y, width, height)`.
 
-    Usage:
-        bo_track = BOTrack(tlwh, score, cls, feat)
-        bo_track.predict()
-        bo_track.update(new_track, frame_id)
+    Examples:
+        Create a BOTrack instance and update its features
+        >>> bo_track = BOTrack(tlwh=[100, 50, 80, 40], score=0.9, cls=1, feat=np.random.rand(128))
+        >>> bo_track.predict()
+        >>> new_track = BOTrack(tlwh=[110, 60, 80, 40], score=0.85, cls=1, feat=np.random.rand(128))
+        >>> bo_track.update(new_track, frame_id=2)
     """
 
     shared_kalman = KalmanFilterXYWH()
 
     def __init__(self, tlwh, score, cls, feat=None, feat_history=50):
-        """Initialize YOLOv8 object with temporal parameters, such as feature history, alpha and current features."""
+        """
+        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).
+            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_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])
+            >>> score = 0.9
+            >>> cls = 1
+            >>> feat = np.random.rand(128)
+            >>> bo_track = BOTrack(tlwh, score, cls, feat)
+        """
         super().__init__(tlwh, score, cls)
 
         self.smooth_feat = None
@@ -54,7 +76,7 @@ class BOTrack(STrack):
         self.alpha = 0.9
 
     def update_features(self, feat):
-        """Update features vector and smooth it using exponential moving average."""
+        """Update the feature vector and apply exponential moving average smoothing."""
         feat /= np.linalg.norm(feat)
         self.curr_feat = feat
         if self.smooth_feat is None:
@@ -65,7 +87,7 @@ class BOTrack(STrack):
         self.smooth_feat /= np.linalg.norm(self.smooth_feat)
 
     def predict(self):
-        """Predicts the mean and covariance using Kalman filter."""
+        """Predicts 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:
             mean_state[6] = 0
@@ -80,14 +102,14 @@ class BOTrack(STrack):
         super().re_activate(new_track, frame_id, new_id)
 
     def update(self, new_track, frame_id):
-        """Update the YOLOv8 instance with new track and frame ID."""
+        """Updates the YOLOv8 instance with new track 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):
-        """Get current position in bounding box format `(top left x, top left y, width, height)`."""
+        """Returns the current bounding box position in `(top left x, top left y, width, height)` format."""
         if self.mean is None:
             return self._tlwh.copy()
         ret = self.mean[:4].copy()
@@ -96,7 +118,7 @@ class BOTrack(STrack):
 
     @staticmethod
     def multi_predict(stracks):
-        """Predicts the mean and covariance of multiple object tracks using shared Kalman filter."""
+        """Predicts the mean and covariance for multiple object tracks using a shared Kalman filter."""
         if len(stracks) <= 0:
             return
         multi_mean = np.asarray([st.mean.copy() for st in stracks])
@@ -111,12 +133,12 @@ class BOTrack(STrack):
             stracks[i].covariance = cov
 
     def convert_coords(self, tlwh):
-        """Converts Top-Left-Width-Height bounding box coordinates to X-Y-Width-Height format."""
+        """Converts tlwh bounding box coordinates to xywh format."""
         return self.tlwh_to_xywh(tlwh)
 
     @staticmethod
     def tlwh_to_xywh(tlwh):
-        """Convert bounding box to format `(center x, center y, width, height)`."""
+        """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
         return ret
@@ -129,9 +151,9 @@ class BOTSORT(BYTETracker):
     Attributes:
         proximity_thresh (float): Threshold for spatial proximity (IoU) between tracks and detections.
         appearance_thresh (float): Threshold for appearance similarity (ReID embeddings) between tracks and detections.
-        encoder (object): Object to handle ReID embeddings, set to None if ReID is not enabled.
+        encoder (Any): Object to handle ReID embeddings, set to None if ReID is not enabled.
         gmc (GMC): An instance of the GMC algorithm for data association.
-        args (object): Parsed command-line arguments containing tracking parameters.
+        args (Any): Parsed command-line arguments containing tracking parameters.
 
     Methods:
         get_kalmanfilter(): Returns an instance of KalmanFilterXYWH for object tracking.
@@ -139,17 +161,29 @@ class BOTSORT(BYTETracker):
         get_dists(tracks, detections): Get distances between tracks and detections using IoU and (optionally) ReID.
         multi_predict(tracks): Predict and track multiple objects with YOLOv8 model.
 
-    Usage:
-        bot_sort = BOTSORT(args, frame_rate)
-        bot_sort.init_track(dets, scores, cls, img)
-        bot_sort.multi_predict(tracks)
+    Examples:
+        Initialize BOTSORT and process detections
+        >>> bot_sort = BOTSORT(args, frame_rate=30)
+        >>> bot_sort.init_track(dets, scores, cls, img)
+        >>> bot_sort.multi_predict(tracks)
 
     Note:
         The class is designed to work with the YOLOv8 object detection model and supports ReID only if enabled via args.
     """
 
     def __init__(self, args, frame_rate=30):
-        """Initialize YOLOv8 object with ReID module and GMC algorithm."""
+        """
+        Initialize YOLOv8 object with ReID module and GMC algorithm.
+
+        Args:
+            args (object): Parsed command-line arguments containing tracking parameters.
+            frame_rate (int): Frame rate of the video being processed.
+
+        Examples:
+            Initialize BOTSORT with command-line arguments and a specified frame rate:
+            >>> args = parse_args()
+            >>> bot_sort = BOTSORT(args, frame_rate=30)
+        """
         super().__init__(args, frame_rate)
         # ReID module
         self.proximity_thresh = args.proximity_thresh
@@ -161,11 +195,11 @@ class BOTSORT(BYTETracker):
         self.gmc = GMC(method=args.gmc_method)
 
     def get_kalmanfilter(self):
-        """Returns an instance of KalmanFilterXYWH for object tracking."""
+        """Returns 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):
-        """Initialize track with detections, scores, and classes."""
+        """Initialize object tracks using detection bounding boxes, scores, class labels, and optional ReID features."""
         if len(dets) == 0:
             return []
         if self.args.with_reid and self.encoder is not None:
@@ -175,13 +209,12 @@ class BOTSORT(BYTETracker):
             return [BOTrack(xyxy, s, c) for (xyxy, s, c) in zip(dets, scores, cls)]  # detections
 
     def get_dists(self, tracks, detections):
-        """Get distances between tracks and detections using IoU and (optionally) ReID embeddings."""
+        """Calculates distances between tracks and detections using IoU and optionally ReID embeddings."""
         dists = matching.iou_distance(tracks, detections)
         dists_mask = dists > self.proximity_thresh
 
-        # TODO: mot20
-        # if not self.args.mot20:
-        dists = matching.fuse_score(dists, detections)
+        if self.args.fuse_score:
+            dists = matching.fuse_score(dists, detections)
 
         if self.args.with_reid and self.encoder is not None:
             emb_dists = matching.embedding_distance(tracks, detections) / 2.0
@@ -191,10 +224,10 @@ class BOTSORT(BYTETracker):
         return dists
 
     def multi_predict(self, tracks):
-        """Predict and track multiple objects with YOLOv8 model."""
+        """Predicts the mean and covariance of multiple object tracks using a shared Kalman filter."""
         BOTrack.multi_predict(tracks)
 
     def reset(self):
-        """Reset tracker."""
+        """Resets the BOTSORT tracker to its initial state, clearing all tracked objects and internal states."""
         super().reset()
         self.gmc.reset_params()