Reset function to process multiple videos

examples/tracking/ultralytics_example_multiple_videos.py

@@ -0,0 +1,102 @@
+import argparse
+
+from tqdm import tqdm
+from ultralytics import YOLO
+
+import supervision as sv
+
+
+def process_video(
+    source_weights_path: str,
+    source_video_path: str,
+    target_video_path: str,
+    tracker: sv.ByteTrack,
+    confidence_threshold: float = 0.3,
+    iou_threshold: float = 0.7,
+) -> None:
+    model = YOLO(source_weights_path)
+
+    box_annotator = sv.BoundingBoxAnnotator()
+    label_annotator = sv.LabelAnnotator()
+    frame_generator = sv.get_video_frames_generator(source_path=source_video_path)
+    video_info = sv.VideoInfo.from_video_path(video_path=source_video_path)
+
+    with sv.VideoSink(target_path=target_video_path, video_info=video_info) as sink:
+        for frame in tqdm(frame_generator, total=video_info.total_frames):
+            results = model(
+                frame, verbose=False, conf=confidence_threshold, iou=iou_threshold
+            )[0]
+            detections = sv.Detections.from_ultralytics(results)
+            detections = tracker.update_with_detections(detections)
+
+            labels = [
+                    f"#{tracker_id} {model.model.names[class_id]} {confidence:0.2f}"
+                for _, _, confidence, class_id, tracker_id, data
+                in detections
+            ]
+
+            annotated_frame = box_annotator.annotate(
+                scene=frame.copy(), detections=detections
+            )
+
+            annotated_labeled_frame = label_annotator.annotate(
+                scene=annotated_frame, detections=detections, labels=labels
+            )
+
+            sink.write_frame(frame=annotated_labeled_frame)
+    # Reset the tracker after processing the video
+    tracker.reset()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Video Processing with YOLO and ByteTrack"
+    )
+    parser.add_argument(
+        "--source_weights_path",
+        required=True,
+        help="Path to the source weights file",
+        type=str,
+    )
+    parser.add_argument(
+        "--source_video_paths",
+        required=True,
+        help="Paths to the source video files",
+        nargs="+",
+        type=str,
+    )
+    parser.add_argument(
+        "--target_video_paths",
+        required=True,
+        help="Paths to the target video files (output)",
+        nargs="+",
+        type=str,
+    )
+    parser.add_argument(
+        "--confidence_threshold",
+        default=0.3,
+        help="Confidence threshold for the model",
+        type=float,
+    )
+    parser.add_argument(
+        "--iou_threshold", default=0.7, help="IOU threshold for the model", type=float
+    )
+
+    args = parser.parse_args()
+
+    source_video_paths = args.source_video_paths[0].split(',')
+    target_video_paths = args.target_video_paths[0].split(',')
+
+    source_video_paths = [path.strip() for path in source_video_paths]
+    target_video_paths = [path.strip() for path in target_video_paths]
+    tracker = sv.ByteTrack()
+
+    for source_video_path, target_video_path in zip(source_video_paths, target_video_paths):
+        process_video(
+            source_weights_path=args.source_weights_path,
+            source_video_path=source_video_path,
+            target_video_path=target_video_path,
+            confidence_threshold=args.confidence_threshold,
+            iou_threshold=args.iou_threshold,
+            tracker=tracker,
+        )

supervision/tracker/byte_tracker/basetrack.py

@@ -38,6 +38,15 @@ class BaseTrack:
     def next_id() -> int:
         BaseTrack._count += 1
         return BaseTrack._count
+    
+    @staticmethod
+    def reset_counter():
+        BaseTrack._count = 0
+        BaseTrack.track_id = 0
+        BaseTrack.start_frame = 0
+        BaseTrack.frame_id = 0
+        BaseTrack.time_since_update = 0
+
 
     def activate(self, *args):
         raise NotImplementedError
@@ -53,3 +62,4 @@ class BaseTrack:
 
     def mark_removed(self):
         self.state = TrackState.Removed
+        

supervision/tracker/byte_tracker/core.py

@@ -210,7 +210,6 @@ class ByteTrack:
             ```python
             import supervision as sv
             from ultralytics import YOLO
-            import numpy as np
 
             model = YOLO(<MODEL_PATH>)
             tracker = sv.ByteTrack()
@@ -221,7 +220,7 @@ class ByteTrack:
             def callback(frame: np.ndarray, index: int) -> np.ndarray:
                 results = model(frame)[0]
                 detections = sv.Detections.from_ultralytics(results)
-                detections = tracker.update_with_detections(detections)
+                detections = byte_tracker.update_with_detections(detections)
 
                 labels = [f"#{tracker_id}" for tracker_id in detections.tracker_id]
 
@@ -261,6 +260,13 @@ class ByteTrack:
 
         return detections
 
+    def reset(self):
+        self.frame_id = 0
+        self.tracked_tracks: List[STrack] = []
+        self.lost_tracks: List[STrack] = []
+        self.removed_tracks: List[STrack] = []
+        BaseTrack.reset_counter()
+
     def update_with_tensors(self, tensors: np.ndarray) -> List[STrack]:
         """
         Updates the tracker with the provided tensors and returns the updated tracks.
@@ -306,6 +312,7 @@ class ByteTrack:
         """ Add newly detected tracklets to tracked_stracks"""
         unconfirmed = []
         tracked_stracks = []  # type: list[STrack]
+
         for track in self.tracked_tracks:
             if not track.is_activated:
                 unconfirmed.append(track)
@@ -481,4 +488,4 @@ def remove_duplicate_tracks(tracks_a: List, tracks_b: List) -> Tuple[List, List]
         track for index, track in enumerate(tracks_b) if index not in duplicates_b
     ]
 
-    return result_a, result_b
+    return result_a, result_b