Skip to content

Commit

Permalink
Various tracker integration
Browse files Browse the repository at this point in the history
BotSORT, ByteTrack, DeepOCSORT, StrongSORT
  • Loading branch information
bstandaert committed Jul 5, 2024
1 parent ccd6a0f commit 6acda53
Show file tree
Hide file tree
Showing 19 changed files with 228 additions and 60 deletions.
2 changes: 1 addition & 1 deletion plugins/track/deep_oc_sort/association.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
5 changes: 2 additions & 3 deletions tracklab/configs/config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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_
Expand All @@ -21,7 +21,6 @@ defaults:
# Use 'pipeline: []' for an empty pipeline
pipeline:
- bbox_detector
- reid
- track

# Experiment name
Expand Down
2 changes: 1 addition & 1 deletion tracklab/configs/engine/video.yaml
Original file line number Diff line number Diff line change
@@ -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:
Expand Down
3 changes: 2 additions & 1 deletion tracklab/configs/modules/bbox_detector/yolox.yaml
Original file line number Diff line number Diff line change
@@ -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
2 changes: 1 addition & 1 deletion tracklab/configs/modules/reid/bpbreid.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
21 changes: 21 additions & 0 deletions tracklab/configs/modules/track/bot_sort.yaml
Original file line number Diff line number Diff line change
@@ -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
10 changes: 10 additions & 0 deletions tracklab/configs/modules/track/byte_track.yaml
Original file line number Diff line number Diff line change
@@ -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
26 changes: 26 additions & 0 deletions tracklab/configs/modules/track/deep_oc_sort.yaml
Original file line number Diff line number Diff line change
@@ -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
21 changes: 21 additions & 0 deletions tracklab/configs/modules/track/strong_sort.yaml
Original file line number Diff line number Diff line change
@@ -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
1 change: 0 additions & 1 deletion tracklab/engine/video.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
8 changes: 8 additions & 0 deletions tracklab/utils/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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)
45 changes: 44 additions & 1 deletion tracklab/utils/cv2.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,10 @@
import colorsys

import cv2
from functools import lru_cache

import matplotlib.cm as cm

from .coordinates import *


Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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
2 changes: 1 addition & 1 deletion tracklab/wrappers/datasets/external_video.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
1 change: 0 additions & 1 deletion tracklab/wrappers/reid/bpbreid_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down
2 changes: 1 addition & 1 deletion tracklab/wrappers/track/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
35 changes: 23 additions & 12 deletions tracklab/wrappers/track/bot_sort_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@

import logging

from tracklab.utils.cv2 import cv2_load_image

log = logging.getLogger(__name__)


Expand All @@ -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):
Expand All @@ -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:
Expand Down
32 changes: 20 additions & 12 deletions tracklab/wrappers/track/byte_track_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
Loading

0 comments on commit 6acda53

Please sign in to comment.