Merge pull request #9 from roboflow/feat/reset-tracker-state

feat(trackers): Implement `reset` function

Author: SkalskiP

trackers/core/base.py

@@ -9,8 +9,16 @@ class BaseTracker(ABC):
     def update(self, detections: sv.Detections) -> sv.Detections:
         pass
 
+    @abstractmethod
+    def reset(self) -> None:
+        pass
+
 
 class BaseTrackerWithFeatures(ABC):
     @abstractmethod
     def update(self, detections: sv.Detections, frame: np.ndarray) -> sv.Detections:
         pass
+
+    @abstractmethod
+    def reset(self) -> None:
+        pass

trackers/core/deepsort/tracker.py

@@ -405,3 +405,7 @@ def update(self, detections: sv.Detections, frame: np.ndarray) -> sv.Detections:
         )
 
         return updated_detections
+
+    def reset(self) -> None:
+        self.trackers = []
+        DeepSORTKalmanBoxTracker.count_id = 0

trackers/core/sort/tracker.py

@@ -22,7 +22,7 @@ class SORTTracker(BaseTracker):
         import supervision as sv
         from rfdetr import RFDETRBase
         from rfdetr.util.coco_classes import COCO_CLASSES
-        from trackers.sort_tracker import SORTTracker
+        from trackers import SORTTracker
 
 
         model = RFDETRBase(device="mps")
@@ -219,3 +219,7 @@ def update(self, detections: sv.Detections) -> sv.Detections:
         )
 
         return updated_detections
+
+    def reset(self) -> None:
+        self.trackers = []
+        SORTKalmanBoxTracker.count_id = 0