Merge pull request #1078 from mikel-brostrom/bytetrack-cleanup

Bytetrack cleanup

boxmot/trackers/botsort/bot_sort.py

@@ -184,7 +184,7 @@ def tlwh(self):
         return ret
 
     @property
-    def tlbr(self):
+    def xyxy(self):
         """Convert bounding box to format `(min x, min y, max x, max y)`, i.e.,
         `(top left, bottom right)`.
         """

boxmot/trackers/bytetrack/byte_tracker.py

@@ -5,15 +5,16 @@
 from boxmot.motion.kalman_filters.adapters import ByteTrackKalmanFilterAdapter
 from boxmot.trackers.bytetrack.basetrack import BaseTrack, TrackState
 from boxmot.utils.matching import fuse_score, iou_distance, linear_assignment
-from boxmot.utils.ops import xywh2xyxy, xyxy2xywh
+from boxmot.utils.ops import tlwh2xyah, xywh2xyxy, xyxy2xywh
 
 
 class STrack(BaseTrack):
     shared_kalman = ByteTrackKalmanFilterAdapter()
 
     def __init__(self, det):
         # wait activate
-        self._tlwh = det[0:4]
+        self.xywh = det[0:4]  # (xc, yc, w, h)
+        self.xyah = tlwh2xyah(det[0:4])
         self.score = det[4]
         self.cls = det[5]
         self.det_ind = det[6]
@@ -49,9 +50,7 @@ def activate(self, kalman_filter, frame_id):
         """Start a new tracklet"""
         self.kalman_filter = kalman_filter
         self.track_id = self.next_id()
-        self.mean, self.covariance = self.kalman_filter.initiate(
-            self.tlwh_to_xyah(self._tlwh)
-        )
+        self.mean, self.covariance = self.kalman_filter.initiate(self.xyah)
 
         self.tracklet_len = 0
         self.state = TrackState.Tracked
@@ -63,7 +62,7 @@ def activate(self, kalman_filter, frame_id):
 
     def re_activate(self, new_track, frame_id, new_id=False):
         self.mean, self.covariance = self.kalman_filter.update(
-            self.mean, self.covariance, self.tlwh_to_xyah(new_track.tlwh)
+            self.mean, self.covariance, new_track.xyah
         )
         self.tracklet_len = 0
         self.state = TrackState.Tracked
@@ -87,9 +86,8 @@ def update(self, new_track, frame_id):
         self.tracklet_len += 1
         # self.cls = cls
 
-        new_tlwh = new_track.tlwh
         self.mean, self.covariance = self.kalman_filter.update(
-            self.mean, self.covariance, self.tlwh_to_xyah(new_tlwh)
+            self.mean, self.covariance, new_track.xyah
         )
         self.state = TrackState.Tracked
         self.is_activated = True
@@ -99,59 +97,19 @@ def update(self, new_track, frame_id):
         self.det_ind = new_track.det_ind
 
     @property
-    # @jit(nopython=True)
-    def tlwh(self):
-        """Get current position in bounding box format `(top left x, top left y,
-        width, height)`.
-        """
-        if self.mean is None:
-            return self._tlwh.copy()
-        ret = self.mean[:4].copy()
-        ret[2] *= ret[3]
-        ret[:2] -= ret[2:] / 2
-        return ret
-
-    @property
-    # @jit(nopython=True)
-    def tlbr(self):
+    def xyxy(self):
         """Convert bounding box to format `(min x, min y, max x, max y)`, i.e.,
         `(top left, bottom right)`.
         """
-        ret = self.tlwh.copy()
-        ret[2:] += ret[:2]
-        return ret
-
-    @staticmethod
-    # @jit(nopython=True)
-    def tlwh_to_xyah(tlwh):
-        """Convert bounding box to format `(center x, center y, aspect ratio,
-        height)`, where the aspect ratio is `width / height`.
-        """
-        ret = np.asarray(tlwh).copy()
-        ret[:2] += ret[2:] / 2
-        ret[2] /= ret[3]
-        return ret
-
-    def to_xyah(self):
-        return self.tlwh_to_xyah(self.tlwh)
-
-    @staticmethod
-    # @jit(nopython=True)
-    def tlbr_to_tlwh(tlbr):
-        ret = np.asarray(tlbr).copy()
-        ret[2:] -= ret[:2]
-        return ret
-
-    @staticmethod
-    # @jit(nopython=True)
-    def tlwh_to_tlbr(tlwh):
-        ret = np.asarray(tlwh).copy()
-        ret[2:] += ret[:2]
+        if self.mean is None:
+            ret = self.xywh.copy()  # (xc, yc, w, h)
+        else:
+            ret = self.mean[:4].copy()
+            ret[2] *= ret[3]
+            ret[:2] -= ret[2:] / 2  # (xc, yc, w, h)
+        ret = xywh2xyxy(ret)
         return ret
 
-    def __repr__(self):
-        return "OT_{}_({}-{})".format(self.track_id, self.start_frame, self.end_frame)
-
 
 class BYTETracker(object):
     def __init__(
@@ -183,13 +141,12 @@ def update(self, dets, _):
         ), "Unsupported 'dets' 2nd dimension lenght, valid lenghts is 6"
 
         dets = np.hstack([dets, np.arange(len(dets)).reshape(-1, 1)])
+        dets = xyxy2xywh(dets)
         self.frame_id += 1
         activated_starcks = []
         refind_stracks = []
         lost_stracks = []
         removed_stracks = []
-
-        dets = xyxy2xywh(dets)
         confs = dets[:, 4]
 
         remain_inds = confs > self.track_thresh
@@ -315,13 +272,8 @@ def update(self, dets, _):
         outputs = []
         for t in output_stracks:
             output = []
-            tlwh = t.tlwh
-            tid = t.track_id
-            tlwh = np.expand_dims(tlwh, axis=0)
-            xyxy = xywh2xyxy(tlwh)
-            xyxy = np.squeeze(xyxy, axis=0)
-            output.extend(xyxy)
-            output.append(tid)
+            output.extend(t.xyxy)
+            output.append(t.track_id)
             output.append(t.score)
             output.append(t.cls)
             output.append(t.det_ind)

boxmot/utils/matching.py

@@ -105,8 +105,8 @@ def iou_distance(atracks, btracks):
         atlbrs = atracks
         btlbrs = btracks
     else:
-        atlbrs = [track.tlbr for track in atracks]
-        btlbrs = [track.tlbr for track in btracks]
+        atlbrs = [track.xyxy for track in atracks]
+        btlbrs = [track.xyxy for track in btracks]
     _ious = ious(atlbrs, btlbrs)
     cost_matrix = 1 - _ious
 

boxmot/utils/ops.py

@@ -95,3 +95,58 @@ def tlwh2xyah(x):
     y[..., 2] = x[..., 2] / x[..., 3]
     y[..., 3] = x[..., 3]
     return y
+
+
+def xyah2tlbr(x):
+    """
+    Convert bounding box coordinates from (xc, yc ,a ,h), where the aspect ratio is `width / height`
+    to (t, l, w, h)`
+    """
+    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
+    y[..., 2] = x[..., 2] * x[..., 3]  # (xc, yc, w, h)
+    y[..., 0] = x[..., 0] - (x[..., 2] / 2)  # xc --> t
+    y[..., 1] = x[..., 1] - (x[..., 3] / 2)  # yc --> l
+    y[..., 2] = x[..., 0] + x[..., 2]
+    y[..., 3] = x[..., 1] + x[..., 3]
+    return y
+
+
+def tlwh2tlbr(x):
+    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
+    y[..., 0] = x[..., 0]
+    y[..., 1] = x[..., 1]
+    y[..., 2] = x[..., 0] + x[..., 2]
+    y[..., 3] = x[..., 1] + x[..., 3]
+    return y
+
+
+def tlbr2tlwh(x):
+    """
+    """
+    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
+    y[..., 0] = x[..., 0]
+    y[..., 1] = x[..., 1]
+    y[..., 2] = x[..., 2] - x[..., 0]
+    y[..., 3] = x[..., 3] + x[..., 1]
+    return y
+
+
+def tlbr2xyah(x):
+    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
+    y[..., 0] = x[..., 0] + (x[..., 2] / 2)  # t --> xc
+    y[..., 1] = x[..., 1] + (x[..., 3] / 2)  # l --> yc
+    y[..., 2] = x[..., 2] / x[..., 3]  # w --> a
+    y[..., 3] = x[..., 3] - x[..., 1]
+    return y
+
+
+def xyah2tlwh(x):
+    """
+    Convert bounding box coordinates from (xc, yc ,a ,h), where the aspect ratio is `width / height`
+    to (t, l, w, h)`
+    """
+    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
+    y[..., 2] = x[..., 2] * x[..., 3]  # (xc, yc, w, h)
+    y[..., 0] = x[..., 0] - (x[..., 2] / 2)  # xc --> t
+    y[..., 1] = x[..., 1] - (x[..., 3] / 2)  # yc --> l
+    return y