init commit

ultralytics/trackers/track.py

@@ -0,0 +1,119 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from functools import partial
+from pathlib import Path
+
+import torch
+
+from ultralytics.utils import YAML, IterableSimpleNamespace
+from ultralytics.utils.checks import check_yaml
+
+from .bot_sort import BOTSORT
+from .byte_tracker import BYTETracker
+
+# A mapping of tracker types to corresponding tracker classes
+TRACKER_MAP = {"bytetrack": BYTETracker, "botsort": BOTSORT}
+
+
+def on_predict_start(predictor: object, persist: bool = False) -> None:
+    """
+    Initialize trackers for object tracking during prediction.
+
+    Args:
+        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
+        >>> predictor = SomePredictorClass()
+        >>> on_predict_start(predictor, persist=True)
+    """
+    if predictor.args.task == "classify":
+        raise ValueError("❌ Classification doesn't support 'mode=track'")
+
+    if hasattr(predictor, "trackers") and persist:
+        return
+
+    tracker = check_yaml(predictor.args.tracker)
+    cfg = IterableSimpleNamespace(**YAML.load(tracker))
+
+    if cfg.tracker_type not in {"bytetrack", "botsort"}:
+        raise AssertionError(f"Only 'bytetrack' and 'botsort' are supported for now, but got '{cfg.tracker_type}'")
+
+    predictor._feats = None  # reset in case used earlier
+    if hasattr(predictor, "_hook"):
+        predictor._hook.remove()
+    if cfg.tracker_type == "botsort" and cfg.with_reid and cfg.model == "auto":
+        from ultralytics.nn.modules.head import Detect
+
+        if not (
+            isinstance(predictor.model.model, torch.nn.Module)
+            and isinstance(predictor.model.model.model[-1], Detect)
+            and not predictor.model.model.model[-1].end2end
+        ):
+            cfg.model = "yolo11n-cls.pt"
+        else:
+            # Register hook to extract input of Detect layer
+            def pre_hook(module, input):
+                predictor._feats = list(input[0])  # unroll to new list to avoid mutation in forward
+
+            predictor._hook = predictor.model.model.model[-1].register_forward_pre_hook(pre_hook)
+
+    trackers = []
+    for _ in range(predictor.dataset.bs):
+        tracker = TRACKER_MAP[cfg.tracker_type](args=cfg, frame_rate=30)
+        trackers.append(tracker)
+        if predictor.dataset.mode != "stream":  # only need one tracker for other modes
+            break
+    predictor.trackers = trackers
+    predictor.vid_path = [None] * predictor.dataset.bs  # for determining when to reset tracker on new video
+
+
+def on_predict_postprocess_end(predictor: object, persist: bool = False) -> None:
+    """
+    Postprocess detected boxes and update with object tracking.
+
+    Args:
+        predictor (object): The predictor object containing the predictions.
+        persist (bool, optional): Whether to persist the trackers if they already exist.
+
+    Examples:
+        Postprocess predictions and update with tracking
+        >>> predictor = YourPredictorClass()
+        >>> on_predict_postprocess_end(predictor, persist=True)
+    """
+    is_obb = predictor.args.task == "obb"
+    is_stream = predictor.dataset.mode == "stream"
+    for i, result in enumerate(predictor.results):
+        tracker = predictor.trackers[i if is_stream else 0]
+        vid_path = predictor.save_dir / Path(result.path).name
+        if not persist and predictor.vid_path[i if is_stream else 0] != vid_path:
+            tracker.reset()
+            predictor.vid_path[i if is_stream else 0] = vid_path
+
+        det = (result.obb if is_obb else result.boxes).cpu().numpy()
+        tracks = tracker.update(det, result.orig_img, getattr(result, "feats", None))
+        if len(tracks) == 0:
+            continue
+        idx = tracks[:, -1].astype(int)
+        predictor.results[i] = result[idx]
+
+        update_args = {"obb" if is_obb else "boxes": torch.as_tensor(tracks[:, :-1])}
+        predictor.results[i].update(**update_args)
+
+
+def register_tracker(model: object, persist: bool) -> None:
+    """
+    Register tracking callbacks to the model for object tracking during prediction.
+
+    Args:
+        model (object): The model object to register tracking callbacks for.
+        persist (bool): Whether to persist the trackers if they already exist.
+
+    Examples:
+        Register tracking callbacks to a YOLO model
+        >>> model = YOLOModel()
+        >>> register_tracker(model, persist=True)
+    """
+    model.add_callback("on_predict_start", partial(on_predict_start, persist=persist))
+    model.add_callback("on_predict_postprocess_end", partial(on_predict_postprocess_end, persist=persist))