diff --git a/examples/heatmap_and_track/README.md b/examples/heatmap_and_track/README.md new file mode 100644 index 00000000..70f53249 --- /dev/null +++ b/examples/heatmap_and_track/README.md @@ -0,0 +1,76 @@ +# heatmap and tracking + +## 👋 hello + +This script performs heatmap and tracking analysis using YOLOv8, an object-detection method and +ByteTrack, a simple yet effective online multi-object tracking method. It uses the +supervision package for multiple tasks such as drawing heatmap annotations, tracking objects, etc. + + +## 💻 install + +- clone repository and navigate to example directory + + ```bash + git clone https://github.com/roboflow/supervision.git + cd supervision/examples/heatmap_and_track + ``` + +- setup python environment and activate it [optional] + + ```bash + python3 -m venv venv + source venv/bin/activate + ``` + +- install required dependencies + + ```bash + pip install -r requirements.txt + ``` + +## 🛠️ script arguments + +- `--source_weights_path`: Required. Specifies the path to the weights file for the +YOLO model. This file contains the trained model data necessary for object detection. +- `--source_video_path`: Required. The path to the source video file that will be +analyzed. This is the input video on which traffic flow analysis will be performed. +- `--target_video_path` (optional): The path to save the output video with annotations. +If not specified, the processed video will be displayed in real-time without being +saved. +- `--confidence_threshold` (optional): Sets the confidence threshold for the YOLO model +to filter detections. Default is `0.3`. This determines how confident the model should +be to recognize an object in the video. +- `--iou_threshold` (optional): Specifies the IOU (Intersection Over Union) threshold +for the model. Default is 0.7. This value is used to manage object detection accuracy, +particularly in distinguishing between different objects. +- `heatmap_alpha` (optional): Opacity of the overlay mask, between 0 and 1. +- `radius` (optional): Radius of the heat circle. +- `track_threshold` (optional): Detection confidence threshold for track activation. +- `track_seconds` (optional): Number of seconds to buffer when a track is lost. +- `match_threshold` (optional): Threshold for matching tracks with detections. + +## ⚙️ run + +```bash +python script.py \ +--source_weights_path data/traffic_analysis.pt \ +--source_video_path data/traffic_analysis.mov \ +--confidence_threshold 0.3 \ +--iou_threshold 0.5 \ +--target_video_path data/traffic_analysis_result.mov +``` + +## © license + +This demo integrates two main components, each with its own licensing: + +YOLOv8: The object detection model used in this demo, YOLOv8, is distributed under the +[AGPL-3.0 license](https://github.com/ultralytics/ultralytics/blob/main/LICENSE). You +can find more details about this license here. + +Supervision: The analytics code that powers the zone-based analysis in this demo is +based on the Supervision library, which is licensed under the +[MIT license](https://github.com/roboflow/supervision/blob/develop/LICENSE.md). This +makes the Supervision part of the code fully open source and freely usable in your +projects. diff --git a/examples/heatmap_and_track/requirements.txt b/examples/heatmap_and_track/requirements.txt new file mode 100644 index 00000000..07f88578 Binary files /dev/null and b/examples/heatmap_and_track/requirements.txt differ diff --git a/examples/heatmap_and_track/script.py b/examples/heatmap_and_track/script.py new file mode 100644 index 00000000..2fbd73ee --- /dev/null +++ b/examples/heatmap_and_track/script.py @@ -0,0 +1,176 @@ +from ultralytics import YOLO +import supervision as sv + +import cv2 +import argparse + +def heatmap_and_track( + source_weights_path: str, + source_video_path: str, + target_video_path: str, + confidence_threshold: float = 0.35, + iou_threshold: float = 0.5, + heatmap_alpha: float = 0.5, + radius: int = 25, + track_threshold: float = 0.35, + track_seconds: int = 5, + match_threshold: float = 0.99, +) -> None: + + ### instantiate model + model = YOLO(source_weights_path) + + ### heatmap config + heat_map_annotator = sv.HeatMapAnnotator( + position = sv.Position.BOTTOM_CENTER, + opacity = heatmap_alpha, + radius = radius, + kernel_size = 25, + top_hue = 0, + low_hue = 125, + ) + + ### annotation config + label_annotator = sv.LabelAnnotator(text_position = sv.Position.CENTER) + + ### get the video fps + cap = cv2.VideoCapture(source_video_path) + fps = int(cap.get(cv2.CAP_PROP_FPS)) + cap.release() + + ### tracker config + byte_tracker = sv.ByteTrack( + track_thresh = track_threshold, + track_buffer = track_seconds * fps, + match_thresh = match_threshold, + frame_rate = fps + ) + + ### video config + video_info = sv.VideoInfo.from_video_path(video_path = source_video_path) + frames_generator = sv.get_video_frames_generator(source_path = source_video_path, stride = 1) + + ### Detect, track, annotate, save + with sv.VideoSink(target_path = target_video_path, video_info = video_info) as sink: + + for frame in frames_generator: + result = model( + source = frame, + classes = [0], # only person class + conf = confidence_threshold, + iou = iou_threshold, + # show_conf = True, + # save_txt = True, + # save_conf = True, + # save = True, + device = None, # use None for CPU, 0 for single GPU, or [0,1] for dual GPU + )[0] + + detections = sv.Detections.from_ultralytics(result) # get detections + + detections = byte_tracker.update_with_detections(detections) # update tracker + + ### draw heatmap + annotated_frame = heat_map_annotator.annotate( + scene = frame.copy(), + detections = detections + ) + + ### draw other attributes from `detections` object + labels = [ + f"#{tracker_id}" + for class_id, tracker_id + in zip(detections.class_id, detections.tracker_id) + ] + + label_annotator.annotate( + scene = annotated_frame, + detections = detections, + labels = labels + ) + + sink.write_frame(frame = annotated_frame) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Heatmap and Tracking with Supervision" + ) + parser.add_argument( + "--source_weights_path", + required=True, + help="Path to the source weights file", + type=str, + ) + parser.add_argument( + "--source_video_path", + required=True, + help="Path to the source video file", + type=str, + ) + parser.add_argument( + "--target_video_path", + required=True, + help="Path to the target video file (output)", + type=str, + ) + parser.add_argument( + "--confidence_threshold", + default=0.35, + help="Confidence threshold for the model", + type=float, + ) + parser.add_argument( + "--iou_threshold", + default=0.5, + help="IOU threshold for the model", + type=float, + ) + parser.add_argument( + "--heatmap_alpha", + default=0.5, + help="Opacity of the overlay mask, between 0 and 1", + type=float, + ) + parser.add_argument( + "--radius", + default=25, + help="Radius of the heat circle", + type=float, + ) + parser.add_argument( + "--track_threshold", + default=0.35, + help="Detection confidence threshold for track activation", + type=float, + ) + parser.add_argument( + "--track_seconds", + default=5, + help="Number of seconds to buffer when a track is lost", + type=int, + ) + parser.add_argument( + "--match_threshold", + default=0.99, + help="Threshold for matching tracks with detections", + type=float, + ) + + args = parser.parse_args() + + heatmap_and_track( + source_weights_path=args.source_weights_path, + source_video_path=args.source_video_path, + target_video_path=args.target_video_path, + confidence_threshold=args.confidence_threshold, + iou_threshold=args.iou_threshold, + heatmap_alpha=args.heatmap_alpha, + radius=args.radius, + track_threshold=args.track_threshold, + track_seconds=args.track_seconds, + match_threshold=args.match_threshold, + ) + +### usage +# python script.py --source_weights_path yolov8s.pt --source_video_path input.mp4 --target_video_path results.mp4 \ No newline at end of file