Merge pull request #680 from HinePo/develop - Adding heatmap example

Adding heatmap example

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
 

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.

examples/heatmap_and_track/requirements.txt

@@ -0,0 +1,2 @@
+supervision[assets]
+ultralytics

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,
+    )