filter dets, segs and poses based on tracking results

examples/track.py

@@ -108,6 +108,7 @@ def run(args):
         with predictor.profilers[1]:
             preds = predictor.model(im, augment=predictor.args.augment, visualize=predictor.args.visualize)
 
+
         # Postprocess
         with predictor.profilers[2]:
             predictor.results = predictor.postprocess(preds, im, im0s)
@@ -127,21 +128,36 @@ def run(args):
                 dets = predictor.results[i].boxes.data
                 # get tracker predictions
                 predictor.tracker_outputs[i] = predictor.trackers[i].update(dets.cpu().detach(), im0)
+                predictor.tracker_outputs[i] = predictor.tracker_outputs[i][predictor.tracker_outputs[i][:, 5].argsort()[::-1]]
             predictor.results[i].speed = {
                 'preprocess': predictor.profilers[0].dt * 1E3 / n,
                 'inference': predictor.profilers[1].dt * 1E3 / n,
                 'postprocess': predictor.profilers[2].dt * 1E3 / n,
                 'tracking': predictor.profilers[3].dt * 1E3 / n
-
             }
 
+            # filter boxes masks and pose results by tracking results
+            yolo_confs = predictor.results[i].boxes.conf.cpu().numpy()
+            tracker_confs = predictor.tracker_outputs[i][:, 5]
+            mask = np.in1d(yolo_confs, tracker_confs)
+
+            if predictor.results[i].masks is not None:
+                predictor.results[i].masks = predictor.results[i].masks[mask]
+                predictor.results[i].boxes = predictor.results[i].boxes[mask]
+            elif predictor.results[i].keypoints is not None:
+                predictor.results[i].boxes = predictor.results[i].boxes[mask]
+                predictor.results[i].keypoints = predictor.results[i].keypoints[mask]
+
             # overwrite bbox results with tracker predictions
             if predictor.tracker_outputs[i].size != 0:
                 predictor.results[i].boxes = Boxes(
                     # xyxy, (track_id), conf, cls
                     boxes=torch.from_numpy(predictor.tracker_outputs[i]).to(dets.device),
                     orig_shape=im0.shape[:2],  # (height, width)
                 )
+
+                #len_boxes = len(predictor.results[i].boxes)
+                #predictor.results[i].masks = predictor.results[i].masks[0:len_boxes]
 
             # write inference results to a file or directory   
             if predictor.args.verbose or predictor.args.save or predictor.args.save_txt or predictor.args.show: