From 6993ca3d718030fa666ff0e9a23329aafd4c2bf2 Mon Sep 17 00:00:00 2001 From: Adonai Vera <45982251+AdonaiVera@users.noreply.github.com> Date: Wed, 31 Jan 2024 15:07:27 -0500 Subject: [PATCH] Reset function to process multiple videos --- .../ultralytics_example_multiple_videos.py | 102 ++++++++++++++++++ supervision/tracker/byte_tracker/basetrack.py | 10 ++ supervision/tracker/byte_tracker/core.py | 13 ++- 3 files changed, 122 insertions(+), 3 deletions(-) create mode 100644 examples/tracking/ultralytics_example_multiple_videos.py diff --git a/examples/tracking/ultralytics_example_multiple_videos.py b/examples/tracking/ultralytics_example_multiple_videos.py new file mode 100644 index 00000000..9d71e5b2 --- /dev/null +++ b/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, + ) \ No newline at end of file diff --git a/supervision/tracker/byte_tracker/basetrack.py b/supervision/tracker/byte_tracker/basetrack.py index c8fadefa..a62c412f 100644 --- a/supervision/tracker/byte_tracker/basetrack.py +++ b/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 + \ No newline at end of file diff --git a/supervision/tracker/byte_tracker/core.py b/supervision/tracker/byte_tracker/core.py index ffe1b1b5..cceb4498 100644 --- a/supervision/tracker/byte_tracker/core.py +++ b/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() 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 \ No newline at end of file