diff --git a/examples/README.md b/examples/README.md index 20bc3d7d..f3f04c71 100644 --- a/examples/README.md +++ b/examples/README.md @@ -7,6 +7,7 @@ interfaces with diverse applications. - [tracking](./tracking) by [@SkalskiP](https://github.com/SkalskiP) - [count people in zone](./count_people_in_zone) by [@SkalskiP](https://github.com/SkalskiP) - [traffic analysis](./traffic_analysis) by [@SkalskiP](https://github.com/SkalskiP) +- [heatmap and track](./heatmap_and_track/) by [@HinePo](https://github.com/HinePo) ## contributing diff --git a/examples/heatmap_and_track/README.md b/examples/heatmap_and_track/README.md new file mode 100644 index 00000000..c2339956 --- /dev/null +++ b/examples/heatmap_and_track/README.md @@ -0,0 +1,75 @@ +# 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` (optional): The path to the source video file that will be +analyzed. This is the input video on which crowd analysis will be performed. +If not specified default is `people-walking.mp4` from supervision assets +- `--target_video_path` (optional): The path to save the output.mp4 video with annotations. +- `--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 weight.pt \ +--source_video_path input_video.mp4 \ +--confidence_threshold 0.3 \ +--iou_threshold 0.5 \ +--target_video_path output_video.mp4 +``` + +## © 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..ffc73e1f --- /dev/null +++ b/examples/heatmap_and_track/requirements.txt @@ -0,0 +1,2 @@ +supervision[assets] +ultralytics diff --git a/examples/heatmap_and_track/script.py b/examples/heatmap_and_track/script.py new file mode 100644 index 00000000..a4eb80b9 --- /dev/null +++ b/examples/heatmap_and_track/script.py @@ -0,0 +1,181 @@ +import argparse + +import cv2 +from ultralytics import YOLO + +import supervision as sv +from supervision.assets import VideoAssets, download_assets + + +def download_video() -> str: + download_assets(VideoAssets.PEOPLE_WALKING) + return VideoAssets.PEOPLE_WALKING.value + + +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 = CPU, 0 = single GPU, or [0,1] = 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", + default=download_video(), + help="Path to the source video file", + type=str, + ) + parser.add_argument( + "--target_video_path", + default="output.mp4", + 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, + )