ultralytics 8.3.142__py3-none-any.whl → 8.3.144__py3-none-any.whl

ultralytics/solutions/speed_estimation.py

@@ -12,24 +12,29 @@ class SpeedEstimator(BaseSolution):
     A class to estimate the speed of objects in a real-time video stream based on their tracks.
 
     This class extends the BaseSolution class and provides functionality for estimating object speeds using
-    tracking data in video streams.
+    tracking data in video streams. Speed is calculated based on pixel displacement over time and converted
+    to real-world units using a configurable meters-per-pixel scale factor.
 
     Attributes:
-        spd (Dict[int, float]): Dictionary storing speed data for tracked objects.
-        trk_hist (Dict[int, float]): Dictionary storing the object tracking data.
-        max_hist (int): maximum track history before computing speed
-        meters_per_pixel (float): Real-world meters represented by one pixel (e.g., 0.04 for 4m over 100px).
-        max_speed (int): Maximum allowed object speed; values above this will be capped at 120 km/h.
+        fps (float): Video frame rate for time calculations.
+        frame_count (int): Global frame counter for tracking temporal information.
+        trk_frame_ids (dict): Maps track IDs to their first frame index.
+        spd (dict): Final speed per object in km/h once locked.
+        trk_hist (dict): Maps track IDs to deque of position history.
+        locked_ids (set): Track IDs whose speed has been finalized.
+        max_hist (int): Required frame history before computing speed.
+        meter_per_pixel (float): Real-world meters represented by one pixel for scene scale conversion.
+        max_speed (int): Maximum allowed object speed; values above this will be capped.
 
     Methods:
-        initialize_region: Initializes the speed estimation region.
-        process: Processes input frames to estimate object speeds.
-        store_tracking_history: Stores the tracking history for an object.
-        extract_tracks: Extracts tracks from the current frame.
-        display_output: Displays the output with annotations.
+        process: Process input frames to estimate object speeds based on tracking data.
+        store_tracking_history: Store the tracking history for an object.
+        extract_tracks: Extract tracks from the current frame.
+        display_output: Display the output with annotations.
 
     Examples:
-        >>> estimator = SpeedEstimator()
+        Initialize speed estimator and process a frame
+        >>> estimator = SpeedEstimator(meter_per_pixel=0.04, max_speed=120)
         >>> frame = cv2.imread("frame.jpg")
         >>> results = estimator.process(frame)
         >>> cv2.imshow("Speed Estimation", results.plot_im)
@@ -44,15 +49,15 @@ class SpeedEstimator(BaseSolution):
         """
         super().__init__(**kwargs)
 
-        self.fps = self.CFG["fps"]  # assumed video FPS
-        self.frame_count = 0  # global frame count
+        self.fps = self.CFG["fps"]  # Video frame rate for time calculations
+        self.frame_count = 0  # Global frame counter
         self.trk_frame_ids = {}  # Track ID → first frame index
         self.spd = {}  # Final speed per object (km/h), once locked
         self.trk_hist = {}  # Track ID → deque of (time, position)
         self.locked_ids = set()  # Track IDs whose speed has been finalized
         self.max_hist = self.CFG["max_hist"]  # Required frame history before computing speed
         self.meter_per_pixel = self.CFG["meter_per_pixel"]  # Scene scale, depends on camera details
-        self.max_speed = self.CFG["max_speed"]  # max_speed adjustment
+        self.max_speed = self.CFG["max_speed"]  # Maximum speed adjustment
 
     def process(self, im0):
         """
@@ -65,6 +70,7 @@ class SpeedEstimator(BaseSolution):
             (SolutionResults): Contains processed image `plot_im` and `total_tracks` (number of tracked objects).
 
         Examples:
+            Process a frame for speed estimation
             >>> estimator = SpeedEstimator()
             >>> image = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
             >>> results = estimator.process(image)
@@ -89,15 +95,15 @@ class SpeedEstimator(BaseSolution):
                     p0, p1 = trk_hist[0], trk_hist[-1]  # First and last points of track
                     dt = (self.frame_count - self.trk_frame_ids[track_id]) / self.fps  # Time in seconds
                     if dt > 0:
-                        dx, dy = p1[0] - p0[0], p1[1] - p0[1]  # pixel displacement
-                        pixel_distance = sqrt(dx * dx + dy * dy)  # get pixel distance
-                        meters = pixel_distance * self.meter_per_pixel  # convert to meters
+                        dx, dy = p1[0] - p0[0], p1[1] - p0[1]  # Pixel displacement
+                        pixel_distance = sqrt(dx * dx + dy * dy)  # Calculate pixel distance
+                        meters = pixel_distance * self.meter_per_pixel  # Convert to meters
                         self.spd[track_id] = int(
                             min((meters / dt) * 3.6, self.max_speed)
-                        )  # convert to km/h and store final speed
-                        self.locked_ids.add(track_id)  # prevent further updates
-                        self.trk_hist.pop(track_id, None)  # free memory
-                        self.trk_frame_ids.pop(track_id, None)  # optional: remove frame start too
+                        )  # Convert to km/h and store final speed
+                        self.locked_ids.add(track_id)  # Prevent further updates
+                        self.trk_hist.pop(track_id, None)  # Free memory
+                        self.trk_frame_ids.pop(track_id, None)  # Remove frame start reference
 
             if track_id in self.spd:
                 speed_label = f"{self.spd[track_id]} km/h"

ultralytics/solutions/streamlit_inference.py

@@ -1,7 +1,7 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 import io
-from typing import Any
+from typing import Any, List
 
 import cv2
 
@@ -24,7 +24,7 @@ class Inference:
         model_path (str): Path to the loaded model.
         model (YOLO): The YOLO model instance.
         source (str): Selected video source (webcam or video file).
-        enable_trk (str): Enable tracking option ("Yes" or "No").
+        enable_trk (bool): Enable tracking option.
         conf (float): Confidence threshold for detection.
         iou (float): IoU threshold for non-maximum suppression.
         org_frame (Any): Container for the original frame to be displayed.
@@ -33,14 +33,19 @@ class Inference:
         selected_ind (List[int]): List of selected class indices for detection.
 
     Methods:
-        web_ui: Sets up the Streamlit web interface with custom HTML elements.
-        sidebar: Configures the Streamlit sidebar for model and inference settings.
-        source_upload: Handles video file uploads through the Streamlit interface.
-        configure: Configures the model and loads selected classes for inference.
-        inference: Performs real-time object detection inference.
+        web_ui: Set up the Streamlit web interface with custom HTML elements.
+        sidebar: Configure the Streamlit sidebar for model and inference settings.
+        source_upload: Handle video file uploads through the Streamlit interface.
+        configure: Configure the model and load selected classes for inference.
+        inference: Perform real-time object detection inference.
 
     Examples:
-        >>> inf = Inference(model="path/to/model.pt")  # Model is an optional argument
+        Create an Inference instance with a custom model
+        >>> inf = Inference(model="path/to/model.pt")
+        >>> inf.inference()
+
+        Create an Inference instance with default settings
+        >>> inf = Inference()
         >>> inf.inference()
     """
 
@@ -62,7 +67,7 @@ class Inference:
         self.org_frame = None  # Container for the original frame display
         self.ann_frame = None  # Container for the annotated frame display
         self.vid_file_name = None  # Video file name or webcam index
-        self.selected_ind = []  # List of selected class indices for detection
+        self.selected_ind: List[int] = []  # List of selected class indices for detection
         self.model = None  # YOLO model instance
 
         self.temp_dict = {"model": None, **kwargs}
@@ -73,7 +78,7 @@ class Inference:
         LOGGER.info(f"Ultralytics Solutions: ✅ {self.temp_dict}")
 
     def web_ui(self):
-        """Sets up the Streamlit web interface with custom HTML elements."""
+        """Set up the Streamlit web interface with custom HTML elements."""
         menu_style_cfg = """<style>MainMenu {visibility: hidden;}</style>"""  # Hide main menu style
 
         # Main title of streamlit application
@@ -102,7 +107,7 @@ class Inference:
             "Video",
             ("webcam", "video"),
         )  # Add source selection dropdown
-        self.enable_trk = self.st.sidebar.radio("Enable Tracking", ("Yes", "No"))  # Enable object tracking
+        self.enable_trk = self.st.sidebar.radio("Enable Tracking", ("Yes", "No")) == "Yes"  # Enable object tracking
         self.conf = float(
             self.st.sidebar.slider("Confidence Threshold", 0.0, 1.0, self.conf, 0.01)
         )  # Slider for confidence
@@ -166,7 +171,7 @@ class Inference:
                     break
 
                 # Process frame with model
-                if self.enable_trk == "Yes":
+                if self.enable_trk:
                     results = self.model.track(
                         frame, conf=self.conf, iou=self.iou, classes=self.selected_ind, persist=True
                     )

ultralytics/solutions/trackzone.py

@@ -23,9 +23,9 @@ class TrackZone(BaseSolution):
         clss (List[int]): Class indices of tracked objects.
 
     Methods:
-        process: Processes each frame of the video, applying region-based tracking.
-        extract_tracks: Extracts tracking information from the input frame.
-        display_output: Displays the processed output.
+        process: Process each frame of the video, applying region-based tracking.
+        extract_tracks: Extract tracking information from the input frame.
+        display_output: Display the processed output.
 
     Examples:
         >>> tracker = TrackZone()
@@ -82,7 +82,7 @@ class TrackZone(BaseSolution):
             )
 
         plot_im = annotator.result()
-        self.display_output(plot_im)  # display output with base class function
+        self.display_output(plot_im)  # Display output with base class function
 
         # Return a SolutionResults
         return SolutionResults(plot_im=plot_im, total_tracks=len(self.track_ids))

ultralytics/trackers/basetrack.py

@@ -2,6 +2,7 @@
 """Module defines the base classes and structures for object tracking in YOLO."""
 
 from collections import OrderedDict
+from typing import Any
 
 import numpy as np
 
@@ -66,15 +67,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 +81,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,7 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 from collections import deque
+from typing import Any, List, Optional
 
 import numpy as np
 import torch
@@ -51,7 +52,9 @@ class BOTrack(STrack):
 
     shared_kalman = KalmanFilterXYWH()
 
-    def __init__(self, tlwh, score, cls, feat=None, feat_history=50):
+    def __init__(
+        self, tlwh: np.ndarray, score: float, cls: int, feat: Optional[np.ndarray] = None, feat_history: int = 50
+    ):
         """
         Initialize a BOTrack object with temporal parameters, such as feature history, alpha, and current features.
 
@@ -59,7 +62,7 @@ class BOTrack(STrack):
             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 (np.ndarray, optional): Feature vector associated with the detection.
             feat_history (int): Maximum length of the feature history deque.
 
         Examples:
@@ -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
@@ -176,12 +179,12 @@ class BOTSORT(BYTETracker):
         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):
+    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,11 +206,13 @@ 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, dets: np.ndarray, scores: np.ndarray, cls: np.ndarray, img: Optional[np.ndarray] = None
+    ) -> List[BOTrack]:
         """Initialize object tracks using detection bounding boxes, scores, class labels, and optional ReID features."""
         if len(dets) == 0:
             return []
@@ -217,7 +222,7 @@ class BOTSORT(BYTETracker):
         else:
             return [BOTrack(xyxy, s, c) for (xyxy, s, c) in zip(dets, scores, cls)]  # detections
 
-    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 +237,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,14 +250,19 @@ 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
 
-    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]))])
         if len(feats) != dets.shape[0] and feats[0].shape[0] == dets.shape[0]: