supervisely 6.73.418__py3-none-any.whl → 6.73.419__py3-none-any.whl

supervisely/nn/tracker/botsort/tracker/matching.py

@@ -0,0 +1,202 @@
+import numpy as np
+from supervisely import logger
+
+try:
+    import scipy
+    import lap
+    from scipy.spatial.distance import cdist
+
+    from cython_bbox import bbox_overlaps as bbox_ious
+except ImportError:
+    logger.warning("scipy, lap, and cython_bbox are not installed, some tracker functionalities may not work.")
+    
+from . import kalman_filter
+
+def merge_matches(m1, m2, shape):
+    O,P,Q = shape
+    m1 = np.asarray(m1)
+    m2 = np.asarray(m2)
+
+    M1 = scipy.sparse.coo_matrix((np.ones(len(m1)), (m1[:, 0], m1[:, 1])), shape=(O, P))
+    M2 = scipy.sparse.coo_matrix((np.ones(len(m2)), (m2[:, 0], m2[:, 1])), shape=(P, Q))
+
+    mask = M1*M2
+    match = mask.nonzero()
+    match = list(zip(match[0], match[1]))
+    unmatched_O = tuple(set(range(O)) - set([i for i, j in match]))
+    unmatched_Q = tuple(set(range(Q)) - set([j for i, j in match]))
+
+    return match, unmatched_O, unmatched_Q
+
+
+def _indices_to_matches(cost_matrix, indices, thresh):
+    matched_cost = cost_matrix[tuple(zip(*indices))]
+    matched_mask = (matched_cost <= thresh)
+
+    matches = indices[matched_mask]
+    unmatched_a = tuple(set(range(cost_matrix.shape[0])) - set(matches[:, 0]))
+    unmatched_b = tuple(set(range(cost_matrix.shape[1])) - set(matches[:, 1]))
+
+    return matches, unmatched_a, unmatched_b
+
+
+def linear_assignment(cost_matrix, thresh):
+    if cost_matrix.size == 0:
+        return np.empty((0, 2), dtype=int), tuple(range(cost_matrix.shape[0])), tuple(range(cost_matrix.shape[1]))
+    matches, unmatched_a, unmatched_b = [], [], []
+    cost, x, y = lap.lapjv(cost_matrix, extend_cost=True, cost_limit=thresh)
+    for ix, mx in enumerate(x):
+        if mx >= 0:
+            matches.append([ix, mx])
+    unmatched_a = np.where(x < 0)[0]
+    unmatched_b = np.where(y < 0)[0]
+    matches = np.asarray(matches)
+    return matches, unmatched_a, unmatched_b
+
+
+def ious(atlbrs, btlbrs):
+    """
+    Compute cost based on IoU
+    :type atlbrs: list[tlbr] | np.ndarray
+    :type atlbrs: list[tlbr] | np.ndarray
+
+    :rtype ious np.ndarray
+    """
+    ious = np.zeros((len(atlbrs), len(btlbrs)), dtype=float)
+    if ious.size == 0:
+        return ious
+
+    ious = bbox_ious(
+        np.ascontiguousarray(atlbrs, dtype=float),
+        np.ascontiguousarray(btlbrs, dtype=float)
+    )
+
+    return ious
+
+
+def tlbr_expand(tlbr, scale=1.2):
+    w = tlbr[2] - tlbr[0]
+    h = tlbr[3] - tlbr[1]
+
+    half_scale = 0.5 * scale
+
+    tlbr[0] -= half_scale * w
+    tlbr[1] -= half_scale * h
+    tlbr[2] += half_scale * w
+    tlbr[3] += half_scale * h
+
+    return tlbr
+
+
+def iou_distance(atracks, btracks):
+    """
+    Compute cost based on IoU
+    :type atracks: list[STrack]
+    :type btracks: list[STrack]
+
+    :rtype cost_matrix np.ndarray
+    """
+
+    if (len(atracks)>0 and isinstance(atracks[0], np.ndarray)) or (len(btracks) > 0 and isinstance(btracks[0], np.ndarray)):
+        atlbrs = atracks
+        btlbrs = btracks
+    else:
+        atlbrs = [track.tlbr for track in atracks]
+        btlbrs = [track.tlbr for track in btracks]
+    _ious = ious(atlbrs, btlbrs)
+    cost_matrix = 1 - _ious
+
+    return cost_matrix
+
+
+def v_iou_distance(atracks, btracks):
+    """
+    Compute cost based on IoU
+    :type atracks: list[STrack]
+    :type btracks: list[STrack]
+
+    :rtype cost_matrix np.ndarray
+    """
+
+    if (len(atracks)>0 and isinstance(atracks[0], np.ndarray)) or (len(btracks) > 0 and isinstance(btracks[0], np.ndarray)):
+        atlbrs = atracks
+        btlbrs = btracks
+    else:
+        atlbrs = [track.tlwh_to_tlbr(track.pred_bbox) for track in atracks]
+        btlbrs = [track.tlwh_to_tlbr(track.pred_bbox) for track in btracks]
+    _ious = ious(atlbrs, btlbrs)
+    cost_matrix = 1 - _ious
+
+    return cost_matrix
+
+
+def embedding_distance(tracks, detections, metric='cosine'):
+    """
+    :param tracks: list[STrack]
+    :param detections: list[BaseTrack]
+    :param metric:
+    :return: cost_matrix np.ndarray
+    """
+
+    cost_matrix = np.zeros((len(tracks), len(detections)), dtype=float)
+    if cost_matrix.size == 0:
+        return cost_matrix
+    det_features = np.asarray([track.curr_feat for track in detections], dtype=float)
+    track_features = np.asarray([track.smooth_feat for track in tracks], dtype=float)
+
+    cost_matrix = np.maximum(0.0, cdist(track_features, det_features, metric))  # / 2.0  # Nomalized features
+    return cost_matrix
+
+
+def gate_cost_matrix(kf, cost_matrix, tracks, detections, only_position=False):
+    if cost_matrix.size == 0:
+        return cost_matrix
+    gating_dim = 2 if only_position else 4
+    gating_threshold = kalman_filter.chi2inv95[gating_dim]
+    # measurements = np.asarray([det.to_xyah() for det in detections])
+    measurements = np.asarray([det.to_xywh() for det in detections])
+    for row, track in enumerate(tracks):
+        gating_distance = kf.gating_distance(
+            track.mean, track.covariance, measurements, only_position)
+        cost_matrix[row, gating_distance > gating_threshold] = np.inf
+    return cost_matrix
+
+
+def fuse_motion(kf, cost_matrix, tracks, detections, only_position=False, lambda_=0.98):
+    if cost_matrix.size == 0:
+        return cost_matrix
+    gating_dim = 2 if only_position else 4
+    gating_threshold = kalman_filter.chi2inv95[gating_dim]
+    # measurements = np.asarray([det.to_xyah() for det in detections])
+    measurements = np.asarray([det.to_xywh() for det in detections])
+    for row, track in enumerate(tracks):
+        gating_distance = kf.gating_distance(
+            track.mean, track.covariance, measurements, only_position, metric='maha')
+        cost_matrix[row, gating_distance > gating_threshold] = np.inf
+        cost_matrix[row] = lambda_ * cost_matrix[row] + (1 - lambda_) * gating_distance
+    return cost_matrix
+
+
+def fuse_iou(cost_matrix, tracks, detections):
+    if cost_matrix.size == 0:
+        return cost_matrix
+    reid_sim = 1 - cost_matrix
+    iou_dist = iou_distance(tracks, detections)
+    iou_sim = 1 - iou_dist
+    fuse_sim = reid_sim * (1 + iou_sim) / 2
+    det_scores = np.array([det.score for det in detections])
+    det_scores = np.expand_dims(det_scores, axis=0).repeat(cost_matrix.shape[0], axis=0)
+    #fuse_sim = fuse_sim * (1 + det_scores) / 2
+    fuse_cost = 1 - fuse_sim
+    return fuse_cost
+
+
+def fuse_score(cost_matrix, detections):
+    if cost_matrix.size == 0:
+        return cost_matrix
+    iou_sim = 1 - cost_matrix
+    det_scores = np.array([det.score for det in detections])
+    det_scores = np.expand_dims(det_scores, axis=0).repeat(cost_matrix.shape[0], axis=0)
+    fuse_sim = iou_sim * det_scores
+    fuse_cost = 1 - fuse_sim
+    return fuse_cost

supervisely/nn/tracker/{bot_sort/bot_sort.py → botsort/tracker/mc_bot_sort.py}

@@ -1,16 +1,16 @@
-from collections import deque
-
 import numpy as np
-
-from supervisely.nn.tracker.utils.gmc import GMC
-from supervisely.nn.tracker.utils.kalman_filter import KalmanFilterXYWH as KalmanFilter
+from collections import deque
 
 from . import matching
+from .gmc import GMC
 from .basetrack import BaseTrack, TrackState
-from .fast_reid_interface import FastReIDInterface
+from .kalman_filter import KalmanFilter
+
+from supervisely.nn.tracker.botsort.osnet_reid.osnet_reid_interface import OsnetReIDInterface
 
 
 class STrack(BaseTrack):
+
     shared_kalman = KalmanFilter()
 
     def __init__(self, tlwh, score, cls, feat=None, feat_history=50):
@@ -81,9 +81,7 @@ class STrack(BaseTrack):
                 if st.state != TrackState.Tracked:
                     multi_mean[i][6] = 0
                     multi_mean[i][7] = 0
-            multi_mean, multi_covariance = STrack.shared_kalman.multi_predict(
-                multi_mean, multi_covariance
-            )
+            multi_mean, multi_covariance = STrack.shared_kalman.multi_predict(multi_mean, multi_covariance)
             for i, (mean, cov) in enumerate(zip(multi_mean, multi_covariance)):
                 stracks[i].mean = mean
                 stracks[i].covariance = cov
@@ -122,9 +120,7 @@ class STrack(BaseTrack):
 
     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_xywh(new_track.tlwh)
-        )
+        self.mean, self.covariance = self.kalman_filter.update(self.mean, self.covariance, self.tlwh_to_xywh(new_track.tlwh))
         if new_track.curr_feat is not None:
             self.update_features(new_track.curr_feat)
         self.tracklet_len = 0
@@ -150,9 +146,7 @@ class STrack(BaseTrack):
 
         new_tlwh = new_track.tlwh
 
-        self.mean, self.covariance = self.kalman_filter.update(
-            self.mean, self.covariance, self.tlwh_to_xywh(new_tlwh)
-        )
+        self.mean, self.covariance = self.kalman_filter.update(self.mean, self.covariance, self.tlwh_to_xywh(new_tlwh))
 
         if new_track.curr_feat is not None:
             self.update_features(new_track.curr_feat)
@@ -166,7 +160,7 @@ class STrack(BaseTrack):
     @property
     def tlwh(self):
         """Get current position in bounding box format `(top left x, top left y,
-        width, height)`.
+                width, height)`.
         """
         if self.mean is None:
             return self._tlwh.copy()
@@ -227,7 +221,7 @@ class STrack(BaseTrack):
         return ret
 
     def __repr__(self):
-        return "OT_{}_({}-{})".format(self.track_id, self.start_frame, self.end_frame)
+        return 'OT_{}_({}-{})'.format(self.track_id, self.start_frame, self.end_frame)
 
 
 class BoTSORT(object):
@@ -254,11 +248,9 @@ class BoTSORT(object):
         self.appearance_thresh = args.appearance_thresh
 
         if args.with_reid:
-            self.encoder = FastReIDInterface(
-                args.fast_reid_config, args.fast_reid_weights, args.device
-            )
+            self.encoder = OsnetReIDInterface(args.reid_weights, args.device, args.fp16)
 
-        self.gmc = GMC(method=args.cmc_method, gmc_file=args.gmc_config)
+        self.gmc = GMC(method=args.cmc_method, verbose=[args.name, args.ablation])
 
     def update(self, output_results, img):
         self.frame_id += 1
@@ -267,6 +259,14 @@ class BoTSORT(object):
         lost_stracks = []
         removed_stracks = []
 
+        # mapping detection to track - INITIALIZE det_id immediately
+        detection_track_map = []
+        for i in range(len(output_results)):
+            detection_track_map.append({
+                "det_id": i,  # Detection index - fill immediately
+                "track_id": None  # track_id remains None until track assignment
+            })
+
         if len(output_results):
             bboxes = output_results[:, :4]
             scores = output_results[:, 4]
@@ -279,6 +279,7 @@ class BoTSORT(object):
             scores = scores[lowest_inds]
             classes = classes[lowest_inds]
             features = output_results[lowest_inds]
+            low_inds_map = lowest_inds.nonzero()[0]  # Mapping to original indices
 
             # Find high threshold detections
             remain_inds = scores > self.args.track_high_thresh
@@ -286,6 +287,7 @@ class BoTSORT(object):
             scores_keep = scores[remain_inds]
             classes_keep = classes[remain_inds]
             features_keep = features[remain_inds]
+            high_inds_map = low_inds_map[remain_inds]  # Mapping to original indices
         else:
             bboxes = []
             scores = []
@@ -293,86 +295,69 @@ class BoTSORT(object):
             dets = []
             scores_keep = []
             classes_keep = []
+            high_inds_map = []
 
-        """Extract embeddings """
-        if self.args.with_reid:
+        '''Extract embeddings '''
+        if self.args.with_reid and len(dets) > 0:
             features_keep = self.encoder.inference(img, dets)
 
         if len(dets) > 0:
-            """Detections"""
             if self.args.with_reid:
-                detections = [
-                    STrack(STrack.tlbr_to_tlwh(tlbr), s, c, f)
-                    for (tlbr, s, c, f) in zip(dets, scores_keep, classes_keep, features_keep)
-                ]
+                detections = [STrack(STrack.tlbr_to_tlwh(tlbr), s, c, f) for
+                            (tlbr, s, c, f) in zip(dets, scores_keep, classes_keep, features_keep)]
             else:
-                detections = [
-                    STrack(STrack.tlbr_to_tlwh(tlbr), s, c)
-                    for (tlbr, s, c) in zip(dets, scores_keep, classes_keep)
-                ]
+                detections = [STrack(STrack.tlbr_to_tlwh(tlbr), s, c) for
+                            (tlbr, s, c) in zip(dets, scores_keep, classes_keep)]
         else:
             detections = []
 
-        """ Add newly detected tracklets to tracked_stracks"""
+        ''' Add newly detected tracklets to tracked_stracks'''
         unconfirmed = []
-        tracked_stracks = []  # type: list[STrack]
+        tracked_stracks = []
         for track in self.tracked_stracks:
             if not track.is_activated:
                 unconfirmed.append(track)
             else:
                 tracked_stracks.append(track)
 
-        """ Step 2: First association, with high score detection boxes"""
+        ''' Step 2: First association, with high score detection boxes'''
         strack_pool = joint_stracks(tracked_stracks, self.lost_stracks)
 
-        # Predict the current location with KF
         STrack.multi_predict(strack_pool)
-
-        # Fix camera motion
         warp = self.gmc.apply(img, dets)
         STrack.multi_gmc(strack_pool, warp)
         STrack.multi_gmc(unconfirmed, warp)
 
-        # Associate with high score detection boxes
         ious_dists = matching.iou_distance(strack_pool, detections)
-        ious_dists_mask = ious_dists > self.proximity_thresh
+        ious_dists_mask = (ious_dists > self.proximity_thresh)
 
         if not self.args.mot20:
             ious_dists = matching.fuse_score(ious_dists, detections)
 
         if self.args.with_reid:
             emb_dists = matching.embedding_distance(strack_pool, detections) / 2.0
-            raw_emb_dists = emb_dists.copy()
             emb_dists[emb_dists > self.appearance_thresh] = 1.0
             emb_dists[ious_dists_mask] = 1.0
             dists = np.minimum(ious_dists, emb_dists)
-
-            # Popular ReID method (JDE / FairMOT)
-            # raw_emb_dists = matching.embedding_distance(strack_pool, detections)
-            # dists = matching.fuse_motion(self.kalman_filter, raw_emb_dists, strack_pool, detections)
-            # emb_dists = dists
-
-            # IoU making ReID
-            # dists = matching.embedding_distance(strack_pool, detections)
-            # dists[ious_dists_mask] = 1.0
         else:
             dists = ious_dists
 
-        matches, u_track, u_detection = matching.linear_assignment(
-            dists, thresh=self.args.match_thresh
-        )
+        matches, u_track, u_detection = matching.linear_assignment(dists, thresh=self.args.match_thresh)
 
         for itracked, idet in matches:
             track = strack_pool[itracked]
             det = detections[idet]
+            orig_idx = int(high_inds_map[idet])  # Get original detection index
+            # Fill track_id for matched detection
+            detection_track_map[orig_idx]["track_id"] = int(track.track_id)
             if track.state == TrackState.Tracked:
-                track.update(detections[idet], self.frame_id)
+                track.update(det, self.frame_id)
                 activated_starcks.append(track)
             else:
                 track.re_activate(det, self.frame_id, new_id=False)
                 refind_stracks.append(track)
 
-        """ Step 3: Second association, with low score detection boxes"""
+        ''' Step 3: Second association, with low score detection boxes'''
         if len(scores):
             inds_high = scores < self.args.track_high_thresh
             inds_low = scores > self.args.track_low_thresh
@@ -380,34 +365,32 @@ class BoTSORT(object):
             dets_second = bboxes[inds_second]
             scores_second = scores[inds_second]
             classes_second = classes[inds_second]
+            second_inds_map = low_inds_map[inds_second]
         else:
             dets_second = []
             scores_second = []
             classes_second = []
+            second_inds_map = []
 
-        # association the untrack to the low score detections
         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, classes_second)
-            ]
+            detections_second = [STrack(STrack.tlbr_to_tlwh(tlbr), s, c) for
+                                (tlbr, s, c) in zip(dets_second, scores_second, classes_second)]
         else:
             detections_second = []
 
-        r_tracked_stracks = [
-            strack_pool[i] for i in u_track if strack_pool[i].state == TrackState.Tracked
-        ]
+        r_tracked_stracks = [strack_pool[i] for i in u_track if strack_pool[i].state == TrackState.Tracked]
         dists = matching.iou_distance(r_tracked_stracks, detections_second)
         matches, u_track, u_detection_second = matching.linear_assignment(dists, thresh=0.5)
         for itracked, idet in matches:
             track = r_tracked_stracks[itracked]
-            det = detections_second[idet]
+            orig_idx = int(second_inds_map[idet])
+            # Fill track_id for matched second-level detection
+            detection_track_map[orig_idx]["track_id"] = int(track.track_id)
             if track.state == TrackState.Tracked:
-                track.update(det, self.frame_id)
+                track.update(detections_second[idet], self.frame_id)
                 activated_starcks.append(track)
             else:
-                track.re_activate(det, self.frame_id, new_id=False)
+                track.re_activate(detections_second[idet], self.frame_id, new_id=False)
                 refind_stracks.append(track)
 
         for it in u_track:
@@ -416,14 +399,20 @@ class BoTSORT(object):
                 track.mark_lost()
                 lost_stracks.append(track)
 
-        """Deal with unconfirmed tracks, usually tracks with only one beginning frame"""
-        detections = [detections[i] for i in u_detection]
-        dists = matching.iou_distance(unconfirmed, detections)
+        '''Deal with unconfirmed tracks'''
+        detections_unconf = [detections[i] for i in u_detection]
+        dists = matching.iou_distance(unconfirmed, detections_unconf)
         if not self.args.mot20:
-            dists = matching.fuse_score(dists, detections)
-        matches, u_unconfirmed, u_detection = matching.linear_assignment(dists, thresh=0.7)
+            dists = matching.fuse_score(dists, detections_unconf)
+        matches, u_unconfirmed, u_detection_unconf = matching.linear_assignment(dists, thresh=0.7)
         for itracked, idet in matches:
-            unconfirmed[itracked].update(detections[idet], self.frame_id)
+            track = unconfirmed[itracked]
+            # Get original detection index from u_detection array
+            detection_idx = u_detection[idet]
+            orig_idx = int(high_inds_map[detection_idx])
+            # Fill track_id for matched unconfirmed detection
+            detection_track_map[orig_idx]["track_id"] = int(track.track_id)
+            unconfirmed[itracked].update(detections_unconf[idet], self.frame_id)
             activated_starcks.append(unconfirmed[itracked])
         for it in u_unconfirmed:
             track = unconfirmed[it]
@@ -433,9 +422,11 @@ class BoTSORT(object):
         """ Step 4: Init new stracks"""
         for inew in u_detection:
             track = detections[inew]
+            orig_idx = int(high_inds_map[inew])
             if track.score < self.new_track_thresh:
                 continue
-
+            # Fill track_id for new detection
+            detection_track_map[orig_idx]["track_id"] = int(track.track_id)
             track.activate(self.kalman_filter, self.frame_id)
             activated_starcks.append(track)
 
@@ -453,15 +444,11 @@ class BoTSORT(object):
         self.lost_stracks.extend(lost_stracks)
         self.lost_stracks = sub_stracks(self.lost_stracks, self.removed_stracks)
         self.removed_stracks.extend(removed_stracks)
-        self.tracked_stracks, self.lost_stracks = remove_duplicate_stracks(
-            self.tracked_stracks, self.lost_stracks
-        )
-
-        # output_stracks = [track for track in self.tracked_stracks if track.is_activated]
-        output_stracks = [track for track in self.tracked_stracks]
+        self.tracked_stracks, self.lost_stracks = remove_duplicate_stracks(self.tracked_stracks, self.lost_stracks)
 
-        return output_stracks
+        output_stracks = [track for track in self.tracked_stracks if track.is_activated]
 
+        return output_stracks, detection_track_map
 
 def joint_stracks(tlista, tlistb):
     exists = {}