style and doc change as requested

AntonioConsiglio · AntonioConsiglio · commit 8bca45bde4ec · 2023-10-09T13:13:28.000+02:00
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, Tuple, Optional
 
 import numpy as np
 
@@ -11,7 +11,7 @@
 class STrack(BaseTrack):
     shared_kalman = KalmanFilter()
 
-    def __init__(self, tlwh, score, class_ids, mask=None):
+    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
@@ -259,13 +259,14 @@ def update_with_detections(self, detections: Detections) -> Detections:
         return detections
 
     def update_with_tensors(
-        self, tensors: np.ndarray, masks: np.ndarray = None
+        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.
@@ -288,11 +289,11 @@ def update_with_tensors(
         dets_second = bboxes[inds_second]
         dets = bboxes[remain_inds]
         if masks is not None:
-            masks_hs = masks[remain_inds]
-            masks_ls = masks[inds_second]
+            masks_keep = masks[remain_inds]
+            masks_second = masks[inds_second]
         else:
-            masks_hs = np.array([None] * len(remain_inds))
-            masks_ls = np.array([None] * len(inds_second))
+            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]
 
@@ -303,7 +304,7 @@ def update_with_tensors(
             """Detections"""
             detections = [
                 STrack(STrack.tlbr_to_tlwh(tlbr), s, c, m)
-                for (tlbr, s, c, m) in zip(dets, scores_keep, class_ids_keep, masks_hs)
+                for (tlbr, s, c, m) in zip(dets, scores_keep, class_ids_keep, masks_keep)
             ]
         else:
             detections = []
@@ -345,7 +346,7 @@ def update_with_tensors(
             detections_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_ls
+                    dets_second, scores_second, class_ids_second, masks_second
                 )
             ]
         else:
