Mamba (Selective State Space Model) for Multi-object tracking

Language: English | 简体中文

✅Introduction

This repo is the unofficial implementation of the following Mamba-based Multi-object trackers: (This two papers have not released the official code yet)

1. MambaTrack: A Simple Baseline for Multiple Object Tracking with State Space Model ACM MM'25

2. TrackSSM: A General Motion Predictor by State-Space Model arXiv:2409.00487

🗺️RoadMap

• [] Add model of paper: Exploring Learning-based Motion Models inMulti-Object Tracking arXiv:2403.10826

🏃Results & Checkpoints

• Supporting dataset: DanceTrack, MOT17 and VisDrone2019

Demo result of DanceTrack-val:

MambaTrack:

TrackSSM:

• Results of DacneTrack-val and checkpoints:

Models	HOTA	MOTA	IDF1	checkpoint
MambaTrack	32.672	78.392	26.419	Baidu Disk, code: e0mv
TrackSSM	27.536	72.366	20.756	Baidu Disk, code: 2797

I'm trying to achieve better results

📑Dataset Preparation

Training

For training the mamba-based models, all data is converted to the trajectory format by tools/gen_traj_data.py.

MOT17 dataset, run:

python tools/gen_traj_data.py --mot17 --save_name mot17

DanceTrack dataset, run:

python tools/gen_traj_data.py --dancetrack --save_name dancetrack

VisDrone dataset, run:

python tools/gen_traj_data.py --visdrone --save_name visdrone

Remember to modify the DATA_ROOT according to your autual path

After running the code, a json file will be generated in ssm_tracker/traj_anno_data.

Inference

For inference, we preprocess the detection result first.

First things first, organize all the video frames in the subfoler images, and then test (or val), following the format: (same as YOLO format)

DanceTrack
    |
    |____images
            |_____test
                    |_____dancetrack0001
                                |______xxxx.jpg

Then run the yolox detector:

python tools/gen_det_results.py --dataset_name dancetrack --data_root /data/datasets/DanceTrack/images/ --split val --exp_file yolox_exps/custom/yolox_x.py --model_path weights/yolox_dancetrack.pth.tar --generate_meta_data

you can refer to experiments/gen_yolox_det_results.sh

The pretrained YOLOX-X model of DanceTrack can be downloaded from their huggingface repo. The detection results are provided in ./det_results

🔨Environment Setup

Some key package versions are as follows:

cython==3.0.11
cython-bbox==0.1.5
lap==0.5.12
loguru==0.7.3
opencv-python==4.7.0

torch==1.13.1+cu116
torchvision==0.14.1+cu116
yolox==0.1.0
mamba-ssm==2.2.3.post2

For installing yolox, please refer to ByteTrack, and for installing mamba-ssm, please refer to Mamba.

🔍Models & Guidelines

All mamba-like models are under ./ssm_tracker, and all kalman-like models are under ./kalman_tracker.

1. MambaTrack:

The architecture of MambaTrack is as follows:

The corresponding config file is ssm_tracker/cfgs/MambaTrack.yaml

2. TrackSSM:

The architecture of TrackSSM is as follows:

The corresponding config file is ssm_tracker/cfgs/TrackSSM.yaml

3. Training

For training, please first modify the corresponding 'true' and 'false' in the bash file experiments/train_ssm_tracker.sh, and run:

sh experiments/train_ssm_tracker.sh

4. Inference

For testing, please first modify the corresponding 'true' and 'false' in the bash file experiments/inference_ssm_tracker.sh, and modify the following arguments:

1. --det_path: the detection result files path, contains {seq_name}.txt
2. --motion_model_path: the trained checkpoint path
3. --config_file : same as train
4. --data_root: your data root of the dataset, following the yolo format

Then run

sh experiments/train_ssm_tracker.sh