add yolo v10 and modify pipeline

ultralytics/trackers/track.py

@@ -1,19 +1,20 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
 
 from functools import partial
+from pathlib import Path
 
 import torch
 
 from ultralytics.utils import IterableSimpleNamespace, yaml_load
 from ultralytics.utils.checks import check_yaml
-
 from .bot_sort import BOTSORT
 from .byte_tracker import BYTETracker
 
-TRACKER_MAP = {'bytetrack': BYTETracker, 'botsort': BOTSORT}
+# A mapping of tracker types to corresponding tracker classes
+TRACKER_MAP = {"bytetrack": BYTETracker, "botsort": BOTSORT}
 
 
-def on_predict_start(predictor, persist=False):
+def on_predict_start(predictor: object, persist: bool = False) -> None:
     """
     Initialize trackers for object tracking during prediction.
 
@@ -24,43 +25,65 @@ def on_predict_start(predictor, persist=False):
     Raises:
         AssertionError: If the tracker_type is not 'bytetrack' or 'botsort'.
     """
-    if hasattr(predictor, 'trackers') and persist:
+    if hasattr(predictor, "trackers") and persist:
         return
+
     tracker = check_yaml(predictor.args.tracker)
     cfg = IterableSimpleNamespace(**yaml_load(tracker))
-    assert cfg.tracker_type in ['bytetrack', 'botsort'], \
-        f"Only support 'bytetrack' and 'botsort' for now, but got '{cfg.tracker_type}'"
+
+    if cfg.tracker_type not in ["bytetrack", "botsort"]:
+        raise AssertionError(f"Only 'bytetrack' and 'botsort' are supported for now, but got '{cfg.tracker_type}'")
+
     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):
-    """Postprocess detected boxes and update with object tracking."""
-    bs = predictor.dataset.bs
-    im0s = predictor.batch[1]
-    for i in range(bs):
-        det = predictor.results[i].boxes.cpu().numpy()
+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. Defaults to False.
+    """
+    path, im0s = predictor.batch[:2]
+
+    is_obb = predictor.args.task == "obb"
+    is_stream = predictor.dataset.mode == "stream"
+    for i in range(len(im0s)):
+        tracker = predictor.trackers[i if is_stream else 0]
+        vid_path = predictor.save_dir / Path(path[i]).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 = (predictor.results[i].obb if is_obb else predictor.results[i].boxes).cpu().numpy()
         if len(det) == 0:
             continue
-        tracks = predictor.trackers[i].update(det, im0s[i])
+        tracks = tracker.update(det, im0s[i])
         if len(tracks) == 0:
             continue
         idx = tracks[:, -1].astype(int)
         predictor.results[i] = predictor.results[i][idx]
-        predictor.results[i].update(boxes=torch.as_tensor(tracks[:, :-1]))
+
+        update_args = dict()
+        update_args["obb" if is_obb else "boxes"] = torch.as_tensor(tracks[:, :-1])
+        predictor.results[i].update(**update_args)
 
 
-def register_tracker(model, persist):
+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.
-
     """
-    model.add_callback('on_predict_start', partial(on_predict_start, persist=persist))
-    model.add_callback('on_predict_postprocess_end', on_predict_postprocess_end)
+    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))