dgenerate-ultralytics-headless 8.3.253__py3-none-any.whl

ultralytics/solutions/vision_eye.py

@@ -0,0 +1,67 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from typing import Any
+
+from ultralytics.solutions.solutions import BaseSolution, SolutionAnnotator, SolutionResults
+from ultralytics.utils.plotting import colors
+
+
+class VisionEye(BaseSolution):
+    """A class to manage object detection and vision mapping in images or video streams.
+
+    This class extends the BaseSolution class and provides functionality for detecting objects, mapping vision points,
+    and annotating results with bounding boxes and labels.
+
+    Attributes:
+        vision_point (tuple[int, int]): Coordinates (x, y) where vision will view objects and draw tracks.
+
+    Methods:
+        process: Process the input image to detect objects, annotate them, and apply vision mapping.
+
+    Examples:
+        >>> vision_eye = VisionEye()
+        >>> frame = cv2.imread("frame.jpg")
+        >>> results = vision_eye.process(frame)
+        >>> print(f"Total detected instances: {results.total_tracks}")
+    """
+
+    def __init__(self, **kwargs: Any) -> None:
+        """Initialize the VisionEye class for detecting objects and applying vision mapping.
+
+        Args:
+            **kwargs (Any): Keyword arguments passed to the parent class and for configuring vision_point.
+        """
+        super().__init__(**kwargs)
+        # Set the vision point where the system will view objects and draw tracks
+        self.vision_point = self.CFG["vision_point"]
+
+    def process(self, im0) -> SolutionResults:
+        """Perform object detection, vision mapping, and annotation on the input image.
+
+        Args:
+            im0 (np.ndarray): The input image for detection and annotation.
+
+        Returns:
+            (SolutionResults): Object containing the annotated image and tracking statistics.
+                - plot_im: Annotated output image with bounding boxes and vision mapping
+                - total_tracks: Number of tracked objects in the frame
+
+        Examples:
+            >>> vision_eye = VisionEye()
+            >>> frame = cv2.imread("image.jpg")
+            >>> results = vision_eye.process(frame)
+            >>> print(f"Detected {results.total_tracks} objects")
+        """
+        self.extract_tracks(im0)  # Extract tracks (bounding boxes, classes, and masks)
+        annotator = SolutionAnnotator(im0, self.line_width)
+
+        for cls, t_id, box, conf in zip(self.clss, self.track_ids, self.boxes, self.confs):
+            # Annotate the image with bounding boxes, labels, and vision mapping
+            annotator.box_label(box, label=self.adjust_box_label(cls, conf, t_id), color=colors(int(t_id), True))
+            annotator.visioneye(box, self.vision_point)
+
+        plot_im = annotator.result()
+        self.display_output(plot_im)  # Display the annotated output using the base class function
+
+        # Return a SolutionResults object with the annotated image and tracking statistics
+        return SolutionResults(plot_im=plot_im, total_tracks=len(self.track_ids))

ultralytics/trackers/__init__.py

@@ -0,0 +1,7 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from .bot_sort import BOTSORT
+from .byte_tracker import BYTETracker
+from .track import register_tracker
+
+__all__ = "BOTSORT", "BYTETracker", "register_tracker"  # allow simpler import

ultralytics/trackers/basetrack.py

@@ -0,0 +1,115 @@
+# 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
+from typing import Any
+
+import numpy as np
+
+
+class TrackState:
+    """Enumeration class representing the possible states of an object being tracked.
+
+    Attributes:
+        New (int): State when the object is newly detected.
+        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
+    Tracked = 1
+    Lost = 2
+    Removed = 3
+
+
+class BaseTrack:
+    """Base class for object tracking, providing foundational attributes and methods.
+
+    Attributes:
+        _count (int): Class-level counter for unique track IDs.
+        track_id (int): Unique identifier for the track.
+        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.
+        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.
+        time_since_update (int): Frames passed since the last update.
+        location (tuple): The location of the object in the context of multi-camera tracking.
+
+    Methods:
+        end_frame: Returns the ID of the last frame where the object was tracked.
+        next_id: Increments and returns the next global track ID.
+        activate: Abstract method to activate the track.
+        predict: Abstract method to predict the next state of the track.
+        update: Abstract method to update the track with new data.
+        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):
+        """Initialize a new track with a unique ID and foundational tracking attributes."""
+        self.track_id = 0
+        self.is_activated = False
+        self.state = TrackState.New
+        self.history = OrderedDict()
+        self.features = []
+        self.curr_feature = None
+        self.score = 0
+        self.start_frame = 0
+        self.frame_id = 0
+        self.time_since_update = 0
+        self.location = (np.inf, np.inf)
+
+    @property
+    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() -> int:
+        """Increment and return the next unique global track ID for object tracking."""
+        BaseTrack._count += 1
+        return BaseTrack._count
+
+    def activate(self, *args: Any) -> None:
+        """Activate the track with provided arguments, initializing necessary attributes for tracking."""
+        raise NotImplementedError
+
+    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: Any, **kwargs: Any) -> None:
+        """Update the track with new observations and data, modifying its state and attributes accordingly."""
+        raise NotImplementedError
+
+    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) -> None:
+        """Mark the track as removed by setting its state to TrackState.Removed."""
+        self.state = TrackState.Removed
+
+    @staticmethod
+    def reset_id() -> None:
+        """Reset the global track ID counter to its initial value."""
+        BaseTrack._count = 0

ultralytics/trackers/bot_sort.py

@@ -0,0 +1,257 @@
+# 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
+
+from ultralytics.utils.ops import xywh2xyxy
+from ultralytics.utils.plotting import save_one_box
+
+from .basetrack import TrackState
+from .byte_tracker import BYTETracker, STrack
+from .utils import matching
+from .utils.gmc import GMC
+from .utils.kalman_filter import KalmanFilterXYWH
+
+
+class BOTrack(STrack):
+    """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.
+
+    Attributes:
+        shared_kalman (KalmanFilterXYWH): A shared Kalman filter for all instances of BOTrack.
+        smooth_feat (np.ndarray): Smoothed feature vector.
+        curr_feat (np.ndarray): Current feature vector.
+        features (deque): A deque to store feature vectors with a maximum length defined by `feat_history`.
+        alpha (float): Smoothing factor for the exponential moving average of features.
+        mean (np.ndarray): The mean state of the Kalman filter.
+        covariance (np.ndarray): The covariance matrix of the Kalman filter.
+
+    Methods:
+        update_features: Update features vector and smooth it using exponential moving average.
+        predict: Predict the mean and covariance using Kalman filter.
+        re_activate: Reactivate a track with updated features and optionally new ID.
+        update: Update the track with new detection and frame ID.
+        tlwh: Property that gets the current position in tlwh format `(top left x, top left y, width, height)`.
+        multi_predict: Predict the mean and covariance of multiple object tracks using shared Kalman filter.
+        convert_coords: Convert tlwh bounding box coordinates to xywh format.
+        tlwh_to_xywh: Convert bounding box to xywh format `(center x, center y, width, height)`.
+
+    Examples:
+        Create a BOTrack instance and update its features
+        >>> bo_track = BOTrack(xywh=np.array([100, 50, 80, 40, 0]), score=0.9, cls=1, feat=np.random.rand(128))
+        >>> bo_track.predict()
+        >>> new_track = BOTrack(xywh=np.array([110, 60, 80, 40, 0]), score=0.85, cls=1, feat=np.random.rand(128))
+        >>> bo_track.update(new_track, frame_id=2)
+    """
+
+    shared_kalman = KalmanFilterXYWH()
+
+    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:
+            xywh (np.ndarray): Bounding box in `(x, y, w, h, idx)` or `(x, y, w, h, angle, idx)` format, where (x, y) is
+                the center, (w, h) are width and height, and `idx` is the detection index.
+            score (float): Confidence score of the detection.
+            cls (int): Class ID of the detected object.
+            feat (np.ndarray, optional): Feature vector associated with the detection.
+            feat_history (int): Maximum length of the feature history deque.
+        """
+        super().__init__(xywh, score, cls)
+
+        self.smooth_feat = None
+        self.curr_feat = None
+        if feat is not None:
+            self.update_features(feat)
+        self.features = deque([], maxlen=feat_history)
+        self.alpha = 0.9
+
+    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
+        if self.smooth_feat is None:
+            self.smooth_feat = feat
+        else:
+            self.smooth_feat = self.alpha * self.smooth_feat + (1 - self.alpha) * feat
+        self.features.append(feat)
+        self.smooth_feat /= np.linalg.norm(self.smooth_feat)
+
+    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:
+            mean_state[6] = 0
+            mean_state[7] = 0
+
+        self.mean, self.covariance = self.kalman_filter.predict(mean_state, self.covariance)
+
+    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: 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) -> 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()
+        ret = self.mean[:4].copy()
+        ret[:2] -= ret[2:] / 2
+        return ret
+
+    @staticmethod
+    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
+        multi_mean = np.asarray([st.mean.copy() for st in stracks])
+        multi_covariance = np.asarray([st.covariance for st in stracks])
+        for i, st in enumerate(stracks):
+            if st.state != TrackState.Tracked:
+                multi_mean[i][6] = 0
+                multi_mean[i][7] = 0
+        multi_mean, multi_covariance = BOTrack.shared_kalman.multi_predict(multi_mean, multi_covariance)
+        for i, (mean, cov) in enumerate(zip(multi_mean, multi_covariance)):
+            stracks[i].mean = mean
+            stracks[i].covariance = cov
+
+    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: 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
+        return ret
+
+
+class BOTSORT(BYTETracker):
+    """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.
+        appearance_thresh (float): Threshold for appearance similarity (ReID embeddings) between tracks and detections.
+        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 (Any): Parsed command-line arguments containing tracking parameters.
+
+    Methods:
+        get_kalmanfilter: Return an instance of KalmanFilterXYWH for object tracking.
+        init_track: Initialize track with detections, scores, and classes.
+        get_dists: Get distances between tracks and detections using IoU and (optionally) ReID.
+        multi_predict: Predict and track multiple objects with a YOLO model.
+        reset: Reset the BOTSORT tracker to its initial state.
+
+    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)
+
+    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: Any, frame_rate: int = 30):
+        """Initialize BOTSORT object with ReID module and GMC algorithm.
+
+        Args:
+            args (Any): Parsed command-line arguments containing tracking parameters.
+            frame_rate (int): Frame rate of the video being processed.
+        """
+        super().__init__(args, frame_rate)
+        self.gmc = GMC(method=args.gmc_method)
+
+        # ReID module
+        self.proximity_thresh = args.proximity_thresh
+        self.appearance_thresh = args.appearance_thresh
+        self.encoder = (
+            (lambda feats, s: [f.cpu().numpy() for f in feats])  # native features do not require any model
+            if args.with_reid and self.args.model == "auto"
+            else ReID(args.model)
+            if args.with_reid
+            else None
+        )
+
+    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, results, img: np.ndarray | None = None) -> list[BOTrack]:
+        """Initialize object tracks using detection bounding boxes, scores, class labels, and optional ReID features."""
+        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, bboxes)
+            return [BOTrack(xywh, s, c, f) for (xywh, s, c, f) in zip(bboxes, results.conf, results.cls, features_keep)]
+        else:
+            return [BOTrack(xywh, s, c) for (xywh, s, c) in zip(bboxes, results.conf, results.cls)]
+
+    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)
+
+        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
+            emb_dists[emb_dists > (1 - self.appearance_thresh)] = 1.0
+            emb_dists[dists_mask] = 1.0
+            dists = np.minimum(dists, emb_dists)
+        return dists
+
+    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) -> None:
+        """Reset the BOTSORT tracker to its initial state, clearing all tracked objects and internal states."""
+        super().reset()
+        self.gmc.reset_params()
+
+
+class ReID:
+    """YOLO model as 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, save=False)  # init
+
+    def __call__(self, img: np.ndarray, dets: np.ndarray) -> list[np.ndarray]:
+        """Extract embeddings for detected objects."""
+        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]