add tracktrack, improassoc, re arrange all codes

Author: JackWoo0831

README.md

@@ -230,6 +230,10 @@ In addition, you can also specify
 
 - Hybrid SORT: `python tracker/track.py --dataset visdrone_part --detector yolo_ultra --tracker hybridsort --kalman_format hybridsort --detector_model_path weights/yolov8l_VisDrone_35epochs_20230509.pt --save_images`
 
+- ImproAssoc: `python tracker/track.py --dataset visdrone_part --detector yolo_ultra --tracker improassoc --kalman_format bot --detector_model_path weights/yolov8l_VisDrone_35epochs_20230509.pt --save_images`
+
+- TrackTrack: `python tracker/track.py --dataset visdrone_part --detector yolo_ultra --tracker tracktrack --kalman_format bot --detector_model_path weights/yolov8l_VisDrone_35epochs_20230509.pt --save_images --nms_thresh 0.95 --reid`
+
 > **Important notes for UCMC Track:**
 > 
 > 1. Camera parameters. The UCMC Track need the intrinsic and extrinsic parameter of camera. Please organize like the format of `tracker/cam_param_files/uavdt/M0101.txt`. One video sequence corresponds to one txt file. If you do not have the labelled parameters, you can refer to the estimating toolbox in original repo ([https://github.com/corfyi/UCMCTrack](https://github.com/corfyi/UCMCTrack)).

requirements.txt

@@ -1,8 +1,9 @@
 numpy
-cython-bbox==0.1.3
+Cython
 loguru
 motmetrics==1.4.0
 ninja
+opencv-python
 
 pandas
 Pillow
@@ -11,6 +12,7 @@ PyYAML
 
 scikit-learn
 scipy
+filterpy
 seaborn
 
 thop
@@ -22,3 +24,5 @@ tqdm
 wandb
 
 gdown
+
+ultralytics

tracker/config_files/visdrone_part.yaml

@@ -1,6 +1,6 @@
 # Config file of VisDrone dataset
 
-DATASET_ROOT: '/data/wujiapeng/datasets/VisDrone2019/VisDrone2019'
+DATASET_ROOT: '/root/datasets/visdrone/'
 SPLIT: val
 CATEGORY_NAMES: 
   - 'pedestrain'

tracker/track.py

@@ -30,6 +30,8 @@
 from trackers.sparse_tracker import SparseTracker
 from trackers.ucmc_tracker import UCMCTracker
 from trackers.hybridsort_tracker import HybridSortTracker
+from trackers.tracktrack_tracker import TrackTrackTracker
+from trackers.improassoc_tracker import ImproAssocTracker
 
 # YOLOX modules
 try:
@@ -74,7 +76,9 @@
     'strongsort': StrongSortTracker, 
     'sparsetrack': SparseTracker, 
     'ucmctrack': UCMCTracker, 
-    'hybridsort': HybridSortTracker
+    'hybridsort': HybridSortTracker, 
+    'tracktrack': TrackTrackTracker, 
+    'improassoc': ImproAssocTracker
 }
 
 def get_args():
@@ -90,13 +94,14 @@ def get_args():
 
     parser.add_argument('--kalman_format', type=str, default='default', help='use what kind of Kalman, sort, deepsort, byte, etc.')
     parser.add_argument('--img_size', type=int, default=1280, help='image size, [h, w]')
-
-    parser.add_argument('--conf_thresh', type=float, default=0.2, help='filter detections')
+    
+    # thresholds
+    parser.add_argument('--conf_thresh', type=float, default=0.2, help='filter detections, serve as high conf thresh in two-stage association')
     parser.add_argument('--conf_thresh_low', type=float, default=0.1, help='filter low conf detections, used in two-stage association')
-    parser.add_argument('--nms_thresh', type=float, default=0.7, help='thresh for NMS')
-    parser.add_argument('--iou_thresh', type=float, default=0.5, help='IOU thresh to filter tracks')
+    parser.add_argument('--init_thresh', type=float, default=0.3, help='filter new detections, larger than this thresh consider as new tracklet')
+    parser.add_argument('--nms_thresh', type=float, default=0.45, help='thresh for NMS')
 
-    parser.add_argument('--device', type=str, default='6', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
+    parser.add_argument('--device', type=str, default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
 
     """yolox"""
     parser.add_argument('--yolox_exp_file', type=str, default='./tracker/yolox_utils/yolox_m.py')
@@ -121,9 +126,15 @@ def get_args():
 
     parser.add_argument('--track_eval', type=bool, default=True, help='Use TrackEval to evaluate')
 
+    parser.add_argument('--cmc_method', type=str, default='orb', help='feature discriptor in camera motion compensation')
+
     """camera parameter"""
     parser.add_argument('--camera_parameter_folder', type=str, default='./tracker/cam_param_files', help='folder path of camera parameter files')
 
+    
+    """tensorrt options"""
+    parser.add_argument('--tensorrt', action='store_true', help='use tensorrt engine to detect and reid')
+
     return parser.parse_args()
 
 def main(args, dataset_cfgs):
@@ -291,7 +302,7 @@ def main(args, dataset_cfgs):
             timer.toc()
 
             if args.save_images:
-                plot_img(img=ori_img, frame_id=frame_idx, results=[cur_tlwh, cur_id, cur_cls], 
+                plot_img(img=ori_img, frame_id=frame_idx + 1, results=[cur_tlwh, cur_id, cur_cls], 
                          save_dir=os.path.join(save_dir, 'vis_results'))
 
         save_results(save_dir=save_dir, 

tracker/track_demo.py

@@ -29,6 +29,8 @@
 from trackers.sparse_tracker import SparseTracker
 from trackers.ucmc_tracker import UCMCTracker
 from trackers.hybridsort_tracker import HybridSortTracker
+from trackers.tracktrack_tracker import TrackTrackTracker
+from trackers.improassoc_tracker import ImproAssocTracker
 
 # YOLOX modules
 try:
@@ -73,7 +75,9 @@
     'strongsort': StrongSortTracker, 
     'sparsetrack': SparseTracker, 
     'ucmctrack': UCMCTracker, 
-    'hybridsort': HybridSortTracker
+    'hybridsort': HybridSortTracker, 
+    'tracktrack': TrackTrackTracker, 
+    'improassoc': ImproAssocTracker
 }
 
 def get_args():
@@ -93,8 +97,8 @@ def get_args():
 
     parser.add_argument('--conf_thresh', type=float, default=0.2, help='filter tracks')
     parser.add_argument('--conf_thresh_low', type=float, default=0.1, help='filter low conf detections, used in two-stage association')
+    parser.add_argument('--init_thresh', type=float, default=0.3, help='filter new detections, larger than this thresh consider as new tracklet')
     parser.add_argument('--nms_thresh', type=float, default=0.7, help='thresh for NMS')
-    parser.add_argument('--iou_thresh', type=float, default=0.5, help='IOU thresh to filter tracks')
 
     parser.add_argument('--device', type=str, default='6', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
 
@@ -122,6 +126,8 @@ def get_args():
 
     parser.add_argument('--track_eval', type=bool, default=True, help='Use TrackEval to evaluate')
 
+    parser.add_argument('--cmc_method', type=str, default='orb', help='feature discriptor in camera motion compensation')
+
     """camera parameter"""
     parser.add_argument('--camera_parameter_folder', type=str, default='./tracker/cam_param_files', help='folder path of camera parameter files')
     return parser.parse_args()
@@ -265,7 +271,7 @@ def main(args):
         results.append((frame_idx + 1, cur_id, cur_tlwh, cur_cls, cur_score))
 
         if args.save_images:
-            plot_img(img=ori_img, frame_id=frame_idx, results=[cur_tlwh, cur_id, cur_cls], 
+            plot_img(img=ori_img, frame_id=frame_idx + 1, results=[cur_tlwh, cur_id, cur_cls], 
                         save_dir=os.path.join(save_dir, save_obj_name, 'vis_results'))
 
     save_results(save_dir=save_dir, 

tracker/trackers/basetrack.py

@@ -1,3 +1,7 @@
+"""
+Base class of Tracklets, used in tracklet.py
+"""
+
 import numpy as np
 from collections import OrderedDict
 

tracker/trackers/basetracker.py

@@ -0,0 +1,99 @@
+"""
+Base class for every tracker, embedded some general codes
+like init, get_features and tracklet merge
+for code clearity
+"""
+
+import numpy as np 
+import torch 
+from .reid_models.engine import crop_and_resize
+from .matching import iou_distance
+
+class BaseTracker(object):
+    def __init__(self, args, frame_rate=30):
+
+        self.tracked_tracklets = []  # list[Tracklet]
+        self.lost_tracklets = []  # list[Tracklet]
+        self.removed_tracklets = []  # list[Tracklet]
+
+        self.frame_id = 0
+        self.args = args
+
+        self.init_thresh = args.init_thresh
+        self.buffer_size = int(frame_rate / 30.0 * args.track_buffer)
+        self.max_time_lost = self.buffer_size
+
+        self.motion = args.kalman_format
+
+    def update(self, output_results, img, ori_img):
+        raise NotImplementedError
+
+    @torch.no_grad()
+    def get_feature(self, tlwhs, ori_img):
+        """
+        get apperance feature of an object
+        tlwhs: shape (num_of_objects, 4)
+        ori_img: original image, np.ndarray, shape(H, W, C)
+        """
+
+        if tlwhs.size == 0:
+            return np.empty((0, 512))
+
+        crop_bboxes = crop_and_resize(tlwhs, ori_img, input_format='tlwh', sz=(64, 128))
+        features = self.reid_model(crop_bboxes).cpu().numpy()
+
+        return features
+    
+    
+    def merge_tracklets(self, activated_tracklets, refind_tracklets, lost_tracklets, removed_tracklets):
+        """
+        update tracklets with current association results
+        """
+        self.tracked_tracklets = BaseTracker.joint_tracklets(self.tracked_tracklets, activated_tracklets)
+        self.tracked_tracklets = BaseTracker.joint_tracklets(self.tracked_tracklets, refind_tracklets)
+        self.lost_tracklets = BaseTracker.sub_tracklets(self.lost_tracklets, self.tracked_tracklets)
+        self.lost_tracklets.extend(lost_tracklets)
+        self.lost_tracklets = BaseTracker.sub_tracklets(self.lost_tracklets, self.removed_tracklets)
+        self.removed_tracklets.extend(removed_tracklets)
+        self.tracked_tracklets, self.lost_tracklets = BaseTracker.remove_duplicate_tracklets(self.tracked_tracklets, self.lost_tracklets)
+
+    @staticmethod
+    def joint_tracklets(tlista, tlistb):
+        exists = {}
+        res = []
+        for t in tlista:
+            exists[t.track_id] = 1
+            res.append(t)
+        for t in tlistb:
+            tid = t.track_id
+            if not exists.get(tid, 0):
+                exists[tid] = 1
+                res.append(t)
+        return res
+
+    @staticmethod
+    def sub_tracklets(tlista, tlistb):
+        tracklets = {}
+        for t in tlista:
+            tracklets[t.track_id] = t
+        for t in tlistb:
+            tid = t.track_id
+            if tracklets.get(tid, 0):
+                del tracklets[tid]
+        return list(tracklets.values())
+
+    @staticmethod
+    def remove_duplicate_tracklets(trackletsa, trackletsb):
+        pdist = iou_distance(trackletsa, trackletsb)
+        pairs = np.where(pdist < 0.15)
+        dupa, dupb = list(), list()
+        for p, q in zip(*pairs):
+            timep = trackletsa[p].frame_id - trackletsa[p].start_frame
+            timeq = trackletsb[q].frame_id - trackletsb[q].start_frame
+            if timep > timeq:
+                dupb.append(q)
+            else:
+                dupa.append(p)
+        resa = [t for i, t in enumerate(trackletsa) if not i in dupa]
+        resb = [t for i, t in enumerate(trackletsb) if not i in dupb]
+        return resa, resb