Various tracker integration

BotSORT, ByteTrack, DeepOCSORT, StrongSORT

plugins/track/deep_oc_sort/association.py

@@ -204,7 +204,7 @@ def linear_assignment(cost_matrix):
         import lap
 
         _, x, y = lap.lapjv(cost_matrix, extend_cost=True)
-        return np.array([[y[i], i] for i in x if i >= 0])  #
+        return np.array([[y[i], i] for i in x])  #
     except ImportError:
         from scipy.optimize import linear_sum_assignment
 

tracklab/configs/config.yaml

@@ -4,13 +4,13 @@
 # to create the final config file. This item *must* be
 # the first element in the file.
 defaults:
-  - dataset: soccernet_gs
+  - dataset: dancetrack
   - eval: trackeval
   - engine: offline
   - state: no_save
   - visualization: save_videos
   - modules/bbox_detector: yolov8
-  - modules/pose_bottomup: openpifpaf
+  - modules/pose_topdown: hrnet_posetrack18
   - modules/reid: bpbreid
   - modules/track: oc_sort
   - _self_
@@ -21,7 +21,6 @@ defaults:
 # Use 'pipeline: []' for an empty pipeline
 pipeline:
   - bbox_detector
-  - reid
   - track
 
 # Experiment name

tracklab/configs/engine/video.yaml

@@ -1,6 +1,6 @@
 _target_: tracklab.engine.VideoOnlineTrackingEngine
 
-filename: "/home/vjoosdeterb/Downloads/ilids/AVSS_AB_EVAL_divx.avi"
+filename: "/path/to/video.mp4"
 target_fps: 1
 num_workers: ${num_cores}
 callbacks:

tracklab/configs/modules/bbox_detector/yolox.yaml

@@ -1,4 +1,5 @@
 _target_: tracklab.wrappers.bbox_detector.mmdetection_api.MMDetection
 config_name: yolox_x_8x8_300e_coco
 path_to_checkpoint: ${model_dir}/mmdet/${.config_name}.pth
-min_confidence: 0.4
+min_confidence: 0.4
+batch_size: 2

tracklab/configs/modules/reid/bpbreid.yaml

@@ -3,7 +3,7 @@ _target_: tracklab.wrappers.BPBReId
 defaults:
   - dataset: default
 
-batch_size: 64
+batch_size: 32
 job_id: "${oc.env:SLURM_JOBID,0}" # TODO
 save_path: reid
 use_keypoints_visibility_scores_for_reid: False

tracklab/configs/modules/track/bot_sort.yaml

@@ -0,0 +1,21 @@
+_target_: tracklab.wrappers.BotSORT
+
+cfg:
+  # weights can be downloaded from :
+  # https://github.com/mikel-brostrom/boxmot/tree/master
+  # https://drive.google.com/uc?id=1vduhq5DpN2q1g4fYEZfPI17MJeh9qyrA
+  model_weights: ${model_dir}/reid/osnet_x1_0_market_256x128_amsgrad_ep150_stp60_lr0.0015_b64_fb10_softmax_labelsmooth_flip.pt
+  fp16: false
+
+  min_confidence: 0.4
+
+  hyperparams:
+    appearance_thresh: 0.4818211117541298  # default = 0.25
+    cmc_method: sparseOptFlow  # default = sparseOptFlow
+    frame_rate: 30  # default = 30
+    lambda_: 0.9896143462366406  # default = 0.985
+    match_thresh: 0.22734550911325851  # default = 0.8
+    new_track_thresh: 0.21144301345190655  # default = 0.6
+    proximity_thresh: 0.5945380911899254  # default = 0.5
+    track_buffer: 60  # default = 30
+    track_high_thresh: 0.33824964456239337  # default = 0.45

tracklab/configs/modules/track/byte_track.yaml

@@ -0,0 +1,10 @@
+_target_: tracklab.wrappers.ByteTrack
+
+cfg:
+  min_confidence: 0.4
+
+  hyperparams:
+    track_thresh: 0.6  # tracking confidence threshold, default = 0.45
+    track_buffer: 30   # the frames for keep lost tracks, default = 25
+    match_thresh: 0.8  # matching threshold for tracking, default = 0.8
+    frame_rate: 30     # FPS, default = 30

tracklab/configs/modules/track/deep_oc_sort.yaml

@@ -0,0 +1,26 @@
+_target_: tracklab.wrappers.DeepOCSORT
+
+cfg:
+  # weights can be downloaded from :
+  # https://github.com/mikel-brostrom/boxmot/tree/master
+  # https://drive.google.com/uc?id=1vduhq5DpN2q1g4fYEZfPI17MJeh9qyrA
+  model_weights: ${model_dir}/reid/osnet_x1_0_market_256x128_amsgrad_ep150_stp60_lr0.0015_b64_fb10_softmax_labelsmooth_flip.pt
+  fp16: false
+
+  min_confidence: 0.4
+
+  hyperparams:
+    det_thresh: 0
+    max_age: 50  # default = 30
+    min_hits: 1  # default = 3
+    iou_threshold: 0.22136877277096445  # default = 0.3
+    delta_t: 1  # default = 3
+    asso_func: giou  # default = iou
+    inertia: 0.3941737016672115  # default = 0.2
+    w_association_emb: 0.75  # default = 0.75
+    alpha_fixed_emb: 0.95  # default = 0.95
+    aw_param: 0.5  # default = 0.5
+    embedding_off: false  # default = false
+    cmc_off: false  # default = false
+    aw_off: false  # default = false
+    new_kf_off: false  # default = false

tracklab/configs/modules/track/strong_sort.yaml

@@ -0,0 +1,21 @@
+_target_: tracklab.wrappers.StrongSORT
+
+cfg:
+  # weights can be downloaded from :
+  # https://github.com/mikel-brostrom/boxmot/tree/master
+  # https://drive.google.com/uc?id=1vduhq5DpN2q1g4fYEZfPI17MJeh9qyrA
+  model_weights: ${model_dir}/reid/osnet_x1_0_market_256x128_amsgrad_ep150_stp60_lr0.0015_b64_fb10_softmax_labelsmooth_flip.pt
+  fp16: false
+
+  min_confidence: 0.4
+  ecc: true  # default = true
+
+  hyperparams:
+    ema_alpha: 0.8962157769329083  # default = 0.9
+    max_age: 40  # default = 70
+    max_dist: 0.1594374041012136  # default = 0.2
+    max_iou_dist: 0.5431835667667874  # default = 0.7
+    max_unmatched_preds: 0  # default = 7
+    mc_lambda: 0.995  # default = 0.995
+    n_init: 3  # default = 3
+    nn_budget: 100  # default = 100

tracklab/engine/video.py

@@ -8,7 +8,6 @@
 import torch
 from lightning import Fabric
 
-from tracklab.engine import TrackingEngine
 from tracklab.engine.engine import merge_dataframes
 from tracklab.pipeline import Pipeline
 

tracklab/utils/__init__.py

@@ -90,6 +90,10 @@ def in_bbox_coord(self, bbox_ltwh):
             lambda x: keypoints_in_bbox_coord(x, bbox_ltwh)
         )
 
+    def keypoints_bbox_xyc(self):
+        """Converts from keypoints in image coordinates to keypoints in bbox coordinates"""
+        return self._obj.apply(
+            lambda r: keypoints_in_bbox_coord(r.keypoints_xyc, r.bbox_ltwh), axis=1)
 
 @pd.api.extensions.register_series_accessor("keypoints")
 class KeypointsSeriesAccessor:
@@ -120,3 +124,7 @@ def conf(self):
 
     def in_bbox_coord(self, bbox_ltwh):
         return keypoints_in_bbox_coord(self._obj.keypoints_xyc, bbox_ltwh)
+
+    def keypoints_bbox_xyc(self):
+        """Converts from keypoints in image coordinates to keypoints in bbox coordinates"""
+        return keypoints_in_bbox_coord(self._obj.keypoints_xyc, self._obj.bbox_ltwh)

tracklab/utils/cv2.py

@@ -1,5 +1,10 @@
+import colorsys
+
 import cv2
 from functools import lru_cache
+
+import matplotlib.cm as cm
+
 from .coordinates import *
 
 
@@ -62,7 +67,7 @@ def cv2_load_image(file_path):
             video_reader.set_filename(video_file)
         image = video_reader[int(frame_id)]
     else:
-        image = cv2.imread(str(file_path))
+        image = cv2.imread(file_path)
     image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
     return image
 
@@ -365,3 +370,41 @@ def draw_text(
         lineType=lineType,
     )
     return text_size
+
+
+def scale_lightness(rgb, scale_l=1.4):
+    # convert rgb to hls
+    h, l, s = colorsys.rgb_to_hls(*rgb)
+    # manipulate h, l, s values and return as rgb
+    return colorsys.hls_to_rgb(h, min(1, l * scale_l), s = s)
+
+def colored_body_parts_overlay(img, masks, clip=True, interpolation=cv2.INTER_CUBIC, alpha=0.28, mask_threshold=0.0, weight_scale=1, rgb=False):
+    width, height = img.shape[1], img.shape[0]
+    white_bckg = np.ones_like(img) * 255
+    for i in range(masks.shape[0]):
+        mask = cv2.resize(masks[i], dsize=(width, height), interpolation=interpolation)
+        if clip:
+            mask = np.clip(mask, 0, 1)
+        else:
+            mask = np.interp(mask, (mask.min(), mask.max()), (0, 255)).astype(np.uint8)
+        weight = mask
+        mask_alpha = np.ones_like(weight)
+        mask_alpha[mask < mask_threshold] = 0
+        mask_alpha = np.expand_dims(mask_alpha, 2)
+        weight = np.expand_dims(weight, 2) / weight_scale
+        color_img = np.zeros_like(img)
+        color = scale_lightness(cm.gist_rainbow(i / (len(masks)-1))[0:-1])
+        color_img[:] = np.flip(np.array(color)*255).astype(np.uint8)
+        white_bckg = white_bckg * (1 - mask_alpha * weight) + color_img * mask_alpha * weight
+    heatmap = masks.sum(axis=0).clip(0, 1)
+    heatmap = cv2.resize(heatmap, dsize=(width, height), interpolation=interpolation)
+    mask_alpha = heatmap
+    mask_alpha[heatmap < mask_threshold] = 0
+    mask_alpha = np.repeat(np.expand_dims(mask_alpha, 2), 3, 2)
+    white_bckg = white_bckg.astype(img.dtype)
+    if rgb:
+        white_bckg = cv2.cvtColor(white_bckg, cv2.COLOR_BGR2RGB)
+    masked_img = (
+            img * (1 - mask_alpha * alpha) + white_bckg * mask_alpha * alpha
+    )
+    return masked_img

tracklab/wrappers/datasets/external_video.py

@@ -115,7 +115,7 @@ def __init__(self, dataset_path: str, video_path: str, *args, **kwargs):
             video_metadata,
             image_metadata,
             None,
-            image_metadata
+            image_metadata,
         )
 
         super().__init__(dataset_path,  dict(val=val_set), *args, **kwargs)

tracklab/wrappers/reid/bpbreid_dataset.py

@@ -18,7 +18,6 @@
 from tracklab.utils.coordinates import rescale_keypoints
 
 from tracklab.utils.cv2 import overlay_heatmap
-import tracklab
 
 from torchreid.data import ImageDataset
 from torchreid.utils.imagetools import (

tracklab/wrappers/track/__init__.py

@@ -3,4 +3,4 @@
 from .byte_track_api import ByteTrack
 from .oc_sort_api import OCSORT
 from .deep_oc_sort_api import DeepOCSORT
-from .bpbreid_strong_sort_api import BPBReIDStrongSORT
+from .bpbreid_strong_sort_api import BPBReIDStrongSORT

tracklab/wrappers/track/bot_sort_api.py

@@ -9,6 +9,8 @@
 
 import logging
 
+from tracklab.utils.cv2 import cv2_load_image
+
 log = logging.getLogger(__name__)
 
 
@@ -20,9 +22,10 @@ class BotSORT(ImageLevelModule):
     ]
     output_columns = ["track_id", "track_bbox_ltwh", "track_bbox_conf"]
 
-    def __init__(self, cfg, device, batch_size):
-        super().__init__(cfg, device, batch_size)
+    def __init__(self, cfg, device, **kwargs):
+        super().__init__(batch_size=1)
         self.cfg = cfg
+        self.device = device
         self.reset()
 
     def reset(self):
@@ -35,21 +38,29 @@ def reset(self):
         )
 
     @torch.no_grad()
-    def preprocess(self, detection: pd.Series, metadata: pd.Series):
-        ltrb = detection.bbox.ltrb()
-        conf = detection.bbox.conf()
-        cls = detection.category_id
-        tracklab_id = detection.name
+    def preprocess(self, image, detections: pd.DataFrame, metadata: pd.Series):
+        processed_detections = []
+        if len(detections) == 0:
+            return {"input": []}
+        for det_id, detection in detections.iterrows():
+            ltrb = detection.bbox.ltrb()
+            conf = detection.bbox.conf()
+            cls = detection.category_id
+            tracklab_id = int(detection.name)
+            processed_detections.append(
+                np.array([*ltrb, conf, cls, tracklab_id])
+            )
         return {
-            "input": np.array(
-                [ltrb[0], ltrb[1], ltrb[2], ltrb[3], conf, cls, tracklab_id]
-            ),
+            "input": np.stack(processed_detections)
         }
 
     @torch.no_grad()
-    def process(self, batch, image, detections: pd.DataFrame):
-        inputs = batch["input"]  # Nx7 [l,t,r,b,conf,class,tracklab_id]
+    def process(self, batch, detections: pd.DataFrame, metadatas: pd.DataFrame):
+        if len(detections) == 0:
+            return []
+        inputs = batch["input"][0]  # Nx7 [l,t,r,b,conf,class,tracklab_id]
         inputs = inputs[inputs[:, 4] > self.cfg.min_confidence]
+        image = cv2_load_image(metadatas['file_path'].values[0])
         results = self.model.update(inputs, image)
         results = np.asarray(results)  # N'x8 [l,t,r,b,track_id,class,conf,idx]
         if results.size:

tracklab/wrappers/track/byte_track_api.py

@@ -19,30 +19,38 @@ class ByteTrack(ImageLevelModule):
     ]
     output_columns = ["track_id", "track_bbox_ltwh", "track_bbox_conf"]
 
-    def __init__(self, cfg, device, batch_size):
-        super().__init__(cfg, device, batch_size)
+    def __init__(self, cfg, device, **kwargs):
+        super().__init__(batch_size=1)
         self.cfg = cfg
+        self.device = device
         self.reset()
 
     def reset(self):
         """Reset the tracker state to start tracking in a new video."""
         self.model = byte_tracker.BYTETracker(**self.cfg.hyperparams)
 
     @torch.no_grad()
-    def preprocess(self, detection: pd.Series, metadata: pd.Series):
-        ltrb = detection.bbox.ltrb()
-        conf = detection.bbox.conf()
-        cls = detection.category_id
-        tracklab_id = detection.name
+    def preprocess(self, image, detections: pd.DataFrame, metadata: pd.Series):
+        processed_detections = []
+        if len(detections) == 0:
+            return {"input": []}
+        for det_id, detection in detections.iterrows():
+            ltrb = detection.bbox.ltrb()
+            conf = detection.bbox.conf()
+            cls = detection.category_id
+            tracklab_id = int(detection.name)
+            processed_detections.append(
+                np.array([*ltrb, conf, cls, tracklab_id])
+            )
         return {
-            "input": np.array(
-                [ltrb[0], ltrb[1], ltrb[2], ltrb[3], conf, cls, tracklab_id]
-            ),
+            "input": np.stack(processed_detections)
         }
 
     @torch.no_grad()
-    def process(self, batch, image, detections: pd.DataFrame):
-        inputs = batch["input"]  # Nx7 [l,t,r,b,conf,class,tracklab_id]
+    def process(self, batch, detections: pd.DataFrame, metadatas: pd.DataFrame):
+        if len(detections) == 0:
+            return []
+        inputs = batch["input"][0]  # Nx7 [l,t,r,b,conf,class,tracklab_id]
         inputs = inputs[inputs[:, 4] > self.cfg.min_confidence]
         results = self.model.update(inputs, None)
         results = np.asarray(results)  # N'x8 [l,t,r,b,track_id,class,conf,idx]