Merge branch 'prediction-smoothing' of https://github.com/roboflow/supervision into prediction-smoothing

supervision/__init__.py

@@ -32,8 +32,8 @@ from supervision.dataset.core import (
 )
 from supervision.detection.annotate import BoxAnnotator
 from supervision.detection.core import Detections
-from supervision.detection.smoother import Smoother
 from supervision.detection.line_counter import LineZone, LineZoneAnnotator
+from supervision.detection.smoother import Smoother
 from supervision.detection.tools.inference_slicer import InferenceSlicer
 from supervision.detection.tools.polygon_zone import PolygonZone, PolygonZoneAnnotator
 from supervision.detection.utils import (

supervision/detection/smoother.py

@@ -1,7 +1,10 @@
 from collections import defaultdict
+
 import numpy as np
+
 from supervision.detection.core import Detections
 
+
 class Smoother:
     """
     A class for smoothing out noise in predictions over time by using a Tracker
@@ -34,11 +37,11 @@ class Smoother:
 
         # Run a tracker to link predictions across frames
         detections = byte_tracker.update_with_detections(detections)
-        
+
         # Record the new frame and get the smoothed predictions
         smoother.add_frame(detections)
         smoothed_detections = smoother.get_smoothed_detections()
-        
+
         # Render
         image_smoothed = box_annotator.annotate(
             scene=image.copy(),
@@ -68,10 +71,7 @@ class Smoother:
         tracking into your inference pipeline.
     """
 
-    def __init__(
-        self,
-        length: int = 5
-    ) -> None:
+    def __init__(self, length: int = 5) -> None:
         """
         Args:
             length (int): The current count of detected objects within the zone
@@ -108,21 +108,21 @@ class Smoother:
                 # initialize a new tracker_id
                 self.tracks[tracker_id] = []
                 self.track_starts[tracker_id] = self.current_frame
-            
+
             self.tracks[tracker_id].append(detections[detection_idx])
             self.track_ends[tracker_id] = self.current_frame
-        
+
         for track_id in self.tracks:
             track = self.tracks[track_id]
             if self.track_ends[track_id] < self.current_frame:
                 # continue tracking for a few frames after the object has left
                 # (to prevent flickering in case it comes back)
                 track.append(None)
-            
+
             if len(track) > self.length:
                 # remove the oldest detection from the track it's too long
                 track.pop(0)
-    
+
     def get_track(self, track_id):
         track = self.tracks[track_id]
         if track is None:
@@ -131,12 +131,12 @@ class Smoother:
         track = [d for d in track if d is not None]
         if len(track) == 0:
             return None
-        
+
         ret = track[0]
         # set to an average of all the detection boxes
         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):
@@ -152,9 +152,10 @@ class Smoother:
             track = self.get_track(track_id)
             if track is not None:
                 tracked_detections.append(track)
-        
+
         return Detections.merge(tracked_detections)
 
+
 class NoneDict(defaultdict):
     """
     Helper class that returns None instead of raising a KeyError
@@ -168,4 +169,4 @@ class NoneDict(defaultdict):
         try:
             return super(NoneDict, self).__getitem__(key)
         except KeyError:
-            return None
+            return None