Merge pull request roboflow#436 from AntonioConsiglio/develop

SkalskiP · web-flow · commit 48851b5ae5cd · 2023-10-09T14:18:31.000+02:00
roboflow#418 Make sv.ByteTrack work with segmentation model
diff --git a/supervision/tracker/byte_tracker/core.py b/supervision/tracker/byte_tracker/core.py
@@ -1,4 +1,4 @@
-from typing import List, Tuple
+from typing import List, Optional, Tuple
 
 import numpy as np
 
@@ -11,7 +11,7 @@
 class STrack(BaseTrack):
     shared_kalman = KalmanFilter()
 
-    def __init__(self, tlwh, score, class_ids):
+    def __init__(self, tlwh, score, class_ids, mask: Optional[np.array] = None):
         # wait activate
         self._tlwh = np.asarray(tlwh, dtype=np.float32)
         self.kalman_filter = None
@@ -21,6 +21,7 @@ def __init__(self, tlwh, score, class_ids):
         self.score = score
         self.class_ids = class_ids
         self.tracklet_len = 0
+        self.mask = mask
 
     def predict(self):
         mean_state = self.mean.copy()
@@ -74,6 +75,7 @@ def re_activate(self, new_track, frame_id, new_id=False):
         if new_id:
             self.track_id = self.next_id()
         self.score = new_track.score
+        self.mask = new_track.mask
 
     def update(self, new_track, frame_id):
         """
@@ -95,6 +97,8 @@ def update(self, new_track, frame_id):
 
         self.score = new_track.score
 
+        self.mask = new_track.mask
+
     @property
     def tlwh(self):
         """Get current position in bounding box format `(top left x, top left y,
@@ -231,9 +235,8 @@ def update_with_detections(self, detections: Detections) -> Detections:
             ... )
             ```
         """
-
         tracks = self.update_with_tensors(
-            tensors=detections2boxes(detections=detections)
+            tensors=detections2boxes(detections=detections), masks=detections.mask
         )
         detections = Detections.empty()
         if len(tracks) > 0:
@@ -249,17 +252,22 @@ def update_with_detections(self, detections: Detections) -> Detections:
             detections.confidence = np.array(
                 [t.score for t in tracks], dtype=np.float32
             )
+            detections.mask = np.array([t.mask for t in tracks], dtype=bool)
+
         else:
             detections.tracker_id = np.array([], dtype=int)
 
         return detections
 
-    def update_with_tensors(self, tensors: np.ndarray) -> List[STrack]:
+    def update_with_tensors(
+        self, tensors: np.ndarray, masks: Optional[np.array] = None
+    ) -> List[STrack]:
         """
         Updates the tracker with the provided tensors and returns the updated tracks.
 
         Parameters:
             tensors: The new tensors to update with.
+            masks: The new masks associated to new tensors
 
         Returns:
             List[STrack]: Updated tracks.
@@ -281,6 +289,12 @@ def update_with_tensors(self, tensors: np.ndarray) -> List[STrack]:
         inds_second = np.logical_and(inds_low, inds_high)
         dets_second = bboxes[inds_second]
         dets = bboxes[remain_inds]
+        if masks is not None:
+            masks_keep = masks[remain_inds]
+            masks_second = masks[inds_second]
+        else:
+            masks_keep = np.array([None] * len(remain_inds))
+            masks_second = np.array([None] * len(inds_second))
         scores_keep = scores[remain_inds]
         scores_second = scores[inds_second]
 
@@ -290,8 +304,10 @@ def update_with_tensors(self, tensors: np.ndarray) -> List[STrack]:
         if len(dets) > 0:
             """Detections"""
             detections = [
-                STrack(STrack.tlbr_to_tlwh(tlbr), s, c)
-                for (tlbr, s, c) in zip(dets, scores_keep, class_ids_keep)
+                STrack(STrack.tlbr_to_tlwh(tlbr), s, c, m)
+                for (tlbr, s, c, m) in zip(
+                    dets, scores_keep, class_ids_keep, masks_keep
+                )
             ]
         else:
             detections = []
@@ -331,8 +347,10 @@ def update_with_tensors(self, tensors: np.ndarray) -> List[STrack]:
         if len(dets_second) > 0:
             """Detections"""
             detections_second = [
-                STrack(STrack.tlbr_to_tlwh(tlbr), s, c)
-                for (tlbr, s, c) in zip(dets_second, scores_second, class_ids_second)
+                STrack(STrack.tlbr_to_tlwh(tlbr), s, c, m)
+                for (tlbr, s, c, m) in zip(
+                    dets_second, scores_second, class_ids_second, masks_second
+                )
             ]
         else:
             detections_second = []
