Merge pull request #1032 from mikel-brostrom/gsi

Gsi

boxmot/__init__.py

@@ -1,5 +1,6 @@
 __version__ = '10.0.17'
 
+from boxmot.postprocessing.gsi import gsi
 from boxmot.tracker_zoo import create_tracker, get_tracker_config
 from boxmot.trackers.botsort.bot_sort import BoTSORT
 from boxmot.trackers.bytetrack.byte_tracker import BYTETracker
@@ -9,4 +10,4 @@
 
 __all__ = ("__version__",
            "StrongSORT", "OCSORT", "BYTETracker", "BoTSORT", "DeepOCSORT",
-           "create_tracker", "get_tracker_config")
+           "create_tracker", "get_tracker_config", "gsi")

boxmot/postprocessing/gsi.py

@@ -0,0 +1,73 @@
+from pathlib import Path
+
+import numpy as np
+from sklearn.gaussian_process import GaussianProcessRegressor as GPR
+from sklearn.gaussian_process.kernels import RBF
+
+from boxmot.utils import logger as LOGGER
+
+
+def linear_interpolation(input_, interval):
+    input_ = input_[np.lexsort([input_[:, 0], input_[:, 1]])]
+    output_ = input_.copy()
+
+    id_pre, f_pre, row_pre = -1, -1, np.zeros((10,))
+    for row in input_:
+        f_curr, id_curr = row[:2].astype(int)
+        if id_curr == id_pre:
+            if f_pre + 1 < f_curr < f_pre + interval:
+                for i, f in enumerate(range(f_pre + 1, f_curr), start=1):
+                    step = (row - row_pre) / (f_curr - f_pre) * i
+                    row_new = row_pre + step
+                    output_ = np.append(output_, row_new[np.newaxis, :], axis=0)
+        else:
+            id_pre = id_curr
+        row_pre = row
+        f_pre = f_curr
+    output_ = output_[np.lexsort([output_[:, 0], output_[:, 1]])]
+    return output_
+
+
+def gaussian_smooth(input_, tau):
+    output_ = list()
+    print('input_', input_)
+    ids = set(input_[:, 1])
+    for i, id_ in enumerate(ids):
+        tracks = input_[input_[:, 1] == id_]
+        print('tracks', tracks)
+        len_scale = np.clip(tau * np.log(tau ** 3 / len(tracks)), tau ** -1, tau ** 2)
+        gpr = GPR(RBF(len_scale, 'fixed'))
+        t = tracks[:, 0].reshape(-1, 1)
+        x = tracks[:, 2].reshape(-1, 1)
+        y = tracks[:, 3].reshape(-1, 1)
+        w = tracks[:, 4].reshape(-1, 1)
+        h = tracks[:, 5].reshape(-1, 1)
+        gpr.fit(t, x)
+        xx = gpr.predict(t)
+        gpr.fit(t, y)
+        yy = gpr.predict(t)
+        gpr.fit(t, w)
+        ww = gpr.predict(t)
+        gpr.fit(t, h)
+        hh = gpr.predict(t)
+        # frame count, id, x, y, w, h, conf, cls, -1 (don't care)
+        output_.extend([
+            [t[j, 0], id_, xx[j], yy[j], ww[j], hh[j], tracks[j, 6], tracks[j, 7], -1] for j in range(len(t))
+        ])
+    return output_
+
+
+def gsi(mot_results_folder=Path('examples/runs/val/exp87/labels'), interval=20, tau=10):
+    tracking_results_files = mot_results_folder.glob('MOT*FRCNN.txt')
+    for p in tracking_results_files:
+        LOGGER.info(f"Applying gaussian smoothed interpolation (GSI) to: {p}")
+        tracking_results = np.loadtxt(p, dtype=int, delimiter=' ')
+        if tracking_results.size != 0:
+            li = linear_interpolation(tracking_results, interval)
+            gsi = gaussian_smooth(li, tau)
+            np.savetxt(p, gsi, fmt='%d %d %d %d %d %d %d %d %d')
+        else:
+            print('No tracking result in {p}. Skipping...')
+
+
+gsi()

examples/evolve.py

@@ -280,8 +280,10 @@ def parse_opt():
                         help='evaluate existing tracker results under mot_callenge/MOTXX-YY/...')
     parser.add_argument('--conf', type=float, default=0.45,
                         help='confidence threshold')
-    parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[1280],
+    parser.add_argument('--imgsz', '--img-size', nargs='+', type=int, default=[1280],
                         help='inference size h,w')
+    parser.add_argument('--gsi', action='store_true',
+                        help='apply gsi to results')
     parser.add_argument('--device', default='',
                         help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
     parser.add_argument('--n-trials', type=int, default=10,

examples/val.py

@@ -171,7 +171,7 @@ def eval(self, opt, seq_paths, save_dir, MOT_results_folder, val_tools_path, gt_
             processes = []
 
             busy_devices = []
-            print(seq_paths)
+
             for i, seq_path in enumerate(seq_paths):
                 # spawn one subprocess per GPU in increasing order.
                 # When max devices are reached start at 0 again
@@ -215,6 +215,11 @@ def eval(self, opt, seq_paths, save_dir, MOT_results_folder, val_tools_path, gt_
 
         print_args(vars(self.opt))
 
+        if opt.gsi:
+            # apply gaussian-smoothed interpolation
+            from boxmot.postprocessing.gsi import gsi
+            gsi(mot_results_folder=save_dir / 'labels')
+
         # run the evaluation on the generated txts
         d = [seq_path.parent.name for seq_path in seq_paths]
         p = subprocess.Popen(
@@ -323,6 +328,8 @@ def parse_opt():
                         help='save results to project/name')
     parser.add_argument('--exist-ok', action='store_true',
                         help='existing project/name ok, do not increment')
+    parser.add_argument('--gsi', action='store_true',
+                        help='apply gsi to results')
     parser.add_argument('--benchmark', type=str, default='MOT17-mini',
                         help='MOT16, MOT17, MOT20')
     parser.add_argument('--split', type=str, default='train',
@@ -331,7 +338,7 @@ def parse_opt():
                         help='evaluate existing results under project/name/labels')
     parser.add_argument('--conf', type=float, default=0.45,
                         help='confidence threshold')
-    parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[1280],
+    parser.add_argument('--imgsz', '--img-size', nargs='+', type=int, default=[1280],
                         help='inference size h,w')
     parser.add_argument('--device', default='',
                         help='cuda device, i.e. 0 or 0,1,2,3 or cpu')

requirements.txt

@@ -11,6 +11,8 @@ opencv-python>=4.6.0
 pandas>=1.1.4           # export matrix
 pre-commit>=3.3.3
 PyYAML>=5.3.1           # read tracker configs
+
+scikit-learn>=1.3.0     # gsi
 tensorboard>=2.13.0
 # base -------------------------------------------------------------------------
 torch>=1.7.0