Merge pull request #696 from roboflow/prediction-smoothing

Add Prediction Smoothing

CONTRIBUTING.md

@@ -52,7 +52,7 @@ git push -u origin main
 
 ### Pre-commit tool
 
-This project utilizes the [pre-commit](https://pre-commit.com/) tool to maintain code quality and consistency. Before submitting a pull request or making any commits, it is important to run the pre-commit tool to ensure that your changes meet the project's guidelines.
+This project uses the [pre-commit](https://pre-commit.com/) tool to maintain code quality and consistency. Before submitting a pull request or making any commits, it is important to run the pre-commit tool to ensure that your changes meet the project's guidelines.
 
 Furthermore, we have integrated a pre-commit GitHub Action into our workflow. This means that with every pull request opened, the pre-commit checks will be automatically enforced, streamlining the code review process and ensuring that all contributions adhere to our quality standards.
 

docs/annotators.md

@@ -1,5 +1,6 @@
 ---
 comments: true
+status: new
 ---
 
 === "BoundingBox"

docs/detection/core.md

@@ -1,5 +1,6 @@
 ---
 comments: true
+status: new
 ---
 
 ## Detections

docs/detection/tools/smoother.md

@@ -0,0 +1,8 @@
+---
+comments: true
+status: new
+---
+
+## Detection Smoother
+
+:::supervision.detection.tools.smoother.DetectionsSmoother

docs/how_to/evaluate_model.md

@@ -1,5 +0,0 @@
----
-comments: true
----
-
-🚧 Page under construction.

docs/how_to/process_video.md

@@ -1,5 +0,0 @@
----
-comments: true
----
-
-🚧 Page under construction.

docs/index.md

@@ -17,6 +17,38 @@ comments: true
 
 We write your reusable computer vision tools. Whether you need to load your dataset from your hard drive, draw detections on an image or video, or count how many detections are in a zone. You can count on us!
 
+<div class="grid cards" markdown>
+
+-   __Detect and Annotate__
+
+    ---
+
+    Annotate predictions from a range of object detection and segmentation models
+
+    [:octicons-arrow-right-24: Tutorial](how_to/detect_and_annotate)
+
+-   __Track Objects__
+
+    ---
+
+    Discover how to enhance video analysis by implementing seamless object tracking
+
+    [:octicons-arrow-right-24: Tutorial](how_to/track_objects)
+
+-   > __Count Objects Crossing Line__
+
+    ---
+
+    Explore methods to accurately count and analyze objects crossing a predefined line
+
+-   > __Filter Objects in Zone__
+
+    ---
+
+    Master the techniques to selectively filter and focus on objects within a specific zone
+
+</div>
+
 ## 💻 Install
 
 You can install `supervision` with pip in a

mkdocs.yml

@@ -28,10 +28,8 @@ nav:
   - Home: index.md
   - How to:
       - Detect and Annotate: how_to/detect_and_annotate.md
-      - Process Video: how_to/process_video.md
       - Track Objects: how_to/track_objects.md
       - Filter Detections: how_to/filter_detections.md
-      - Evaluate Model: how_to/evaluate_model.md
   - API:
     - Classifications:
         - Core: classification/core.md
@@ -42,6 +40,7 @@ nav:
         - Line Zone: detection/tools/line_zone.md
         - Polygon Zone: detection/tools/polygon_zone.md
         - Inference Slicer: detection/tools/inference_slicer.md
+        - Detection Smoother: detection/tools/smoother.md
     - Annotators: annotators.md
     - Trackers: trackers.md
     - Datasets: datasets.md
@@ -82,6 +81,7 @@ theme:
     code: Roboto Mono
   features:
     - content.code.copy
+    - content.code.annotate
 
 plugins:
   - mkdocstrings

supervision/__init__.py

@@ -37,6 +37,7 @@ from supervision.detection.core import Detections
 from supervision.detection.line_counter import LineZone, LineZoneAnnotator
 from supervision.detection.tools.inference_slicer import InferenceSlicer
 from supervision.detection.tools.polygon_zone import PolygonZone, PolygonZoneAnnotator
+from supervision.detection.tools.smoother import DetectionsSmoother
 from supervision.detection.utils import (
     box_iou_batch,
     calculate_masks_centroids,

supervision/annotators/core.py

@@ -98,7 +98,7 @@ class MaskAnnotator(BaseAnnotator):
 
     !!! warning
 
-        This annotator utilizes the `sv.Detections.mask`.
+        This annotator uses `sv.Detections.mask`.
     """
 
     def __init__(
@@ -181,7 +181,7 @@ class PolygonAnnotator(BaseAnnotator):
 
     !!! warning
 
-        This annotator utilizes the `sv.Detections.mask`.
+        This annotator uses `sv.Detections.mask`.
     """
 
     def __init__(
@@ -349,7 +349,7 @@ class HaloAnnotator(BaseAnnotator):
 
     !!! warning
 
-        This annotator utilizes the `sv.Detections.mask`.
+        This annotator uses `sv.Detections.mask`.
     """
 
     def __init__(
@@ -1020,7 +1020,7 @@ class TraceAnnotator:
 
     !!! warning
 
-        This annotator utilizes the `sv.Detections.tracker_id`. Read
+        This annotator uses the `sv.Detections.tracker_id`. Read
         [here](https://supervision.roboflow.com/trackers/) to learn how to plug
         tracking into your inference pipeline.
     """

supervision/detection/line_counter.py

@@ -15,7 +15,7 @@ class LineZone:
 
     !!! warning
 
-        LineZone utilizes the `tracker_id`. Read
+        LineZone uses the `tracker_id`. Read
         [here](https://supervision.roboflow.com/trackers/) to learn how to plug
         tracking into your inference pipeline.
 

supervision/detection/tools/smoother.py

@@ -0,0 +1,118 @@
+from collections import defaultdict, deque
+from copy import deepcopy
+from typing import Optional
+
+import numpy as np
+
+from supervision.detection.core import Detections
+
+
+class DetectionsSmoother:
+    """
+    A utility class for smoothing detections over multiple frames in video tracking.
+    It maintains a history of detections for each track and provides smoothed
+    predictions based on these histories.
+
+    <video controls>
+        <source
+            src="https://media.roboflow.com/supervision-detection-smoothing.mp4"
+            type="video/mp4">
+    </video>
+
+    !!! warning
+
+        - `DetectionsSmoother` requires the `tracker_id` for each detection. Refer to
+          [Roboflow Trackers](https://supervision.roboflow.com/trackers/) for
+          information on integrating tracking into your inference pipeline.
+        - This class is not compatible with segmentation models.
+
+    Example:
+        ```python
+        import supervision as sv
+
+        from ultralytics import YOLO
+
+        video_info = sv.VideoInfo.from_video_path(video_path=<SOURCE_FILE_PATH>)
+        frame_generator = sv.get_video_frames_generator(source_path=<SOURCE_FILE_PATH>)
+
+        model = YOLO(<MODEL_PATH>)
+        tracker = sv.ByteTrack(frame_rate=video_info.fps)
+        smoother = sv.DetectionsSmoother()
+
+        annotator = sv.BoundingBoxAnnotator()
+
+        with sv.VideoSink(<TARGET_FILE_PATH>, video_info=video_info) as sink:
+            for frame in frame_generator:
+                result = model(frame)[0]
+                detections = sv.Detections.from_ultralytics(result)
+                detections = tracker.update_with_detections(detections)
+                detections = smoother.update_with_detections(detections)
+
+                annotated_frame = bounding_box_annotator.annotate(frame.copy(), detections)
+                sink.write_frame(annotated_frame)
+        ```
+    """  # noqa: E501 // docs
+
+    def __init__(self, length: int = 5) -> None:
+        """
+        Args:
+            length (int): The maximum number of frames to consider for smoothing
+                detections. Defaults to 5.
+        """
+        self.tracks = defaultdict(lambda: deque(maxlen=length))
+
+    def update_with_detections(self, detections: Detections) -> Detections:
+        """
+        Updates the smoother with a new set of detections from a frame.
+
+        Args:
+            detections (Detections): The detections to add to the smoother.
+        """
+
+        if detections.tracker_id is None:
+            print(
+                "Smoothing skipped. DetectionsSmoother requires tracker_id. Refer to "
+                "https://supervision.roboflow.com/trackers for more information."
+            )
+            return detections
+
+        for detection_idx in range(len(detections)):
+            tracker_id = detections.tracker_id[detection_idx]
+            if tracker_id is None:
+                continue
+
+            self.tracks[tracker_id].append(detections[detection_idx])
+
+        for track_id in self.tracks.keys():
+            if track_id not in detections.tracker_id:
+                self.tracks[track_id].append(None)
+
+        for track_id in list(self.tracks.keys()):
+            if all([d is None for d in self.tracks[track_id]]):
+                del self.tracks[track_id]
+
+        return self.get_smoothed_detections()
+
+    def get_track(self, track_id: int) -> Optional[Detections]:
+        track = self.tracks.get(track_id, None)
+        if track is None:
+            return None
+
+        track = [d for d in track if d is not None]
+        if len(track) == 0:
+            return None
+
+        ret = deepcopy(track[0])
+        ret.xyxy = np.mean([d.xyxy for d in track], axis=0)
+        ret.confidence = np.mean([d.confidence for d in track], axis=0)
+
+        return ret
+
+    def get_smoothed_detections(self) -> Detections:
+        tracked_detections = []
+        for track_id in self.tracks:
+            track = self.get_track(track_id)
+            if track is not None:
+                tracked_detections.append(track)
+
+        return Detections.merge(tracked_detections)