Awesome

Correlation-Embedded Transformer Tracking: A Single-Branch Framework (SuperSBT)

Fei Xie, Wankou Yang, Chunyu Wang, Lei Chu, Yue Cao, Chao Ma, Wenjun Zeng

:star: SuperSBT is accepted by T-PAMI in Aug. 2024!

:star: The official implementation for the SuperSBT.

:star: SuperSBT is the improved version of the SBT which is published in CVPR 2022 Correlation-Aware Deep Tracking .

[Pretrained Weight] [Models] [Raw Results]

Highlights

Improved One-stream Tracking Framework

SuperSBT adopts a hierarchical architecture with a local modeling layer to enhance shallow-level features. A unified relation modeling is proposed to remove complex handcrafted layer pattern designs. SuperSBT is further improved by masked image modeling pre-training, integrating temporal modeling, and equipping with dedicated prediction heads.

Better Speed-performance Trade-off

Install the environment

Option1: Use the Anaconda (CUDA 10.2)

conda create -n supersbt python=3.8
conda activate supersbt
bash install.sh

Option2: Use the Anaconda (CUDA 11.3)

conda env create -f supersbt_cuda113_env.yaml

Option3: Use the docker file

We provide the full docker file here.

Set project paths

Run the following command to set paths for this project

python tracking/create_default_local_file.py --workspace_dir . --data_dir ./data --save_dir ./output

After running this command, you can also modify paths by editing these two files

lib/train/admin/local.py  # paths about training
lib/test/evaluation/local.py  # paths about testing

Data Preparation

Put the tracking datasets in ./data. It should look like:

${PROJECT_ROOT}
 -- data
     -- lasot
         |-- airplane
         |-- basketball
         |-- bear
         ...
     -- got10k
         |-- test
         |-- train
         |-- val
     -- coco
         |-- annotations
         |-- images
     -- trackingnet
         |-- TRAIN_0
         |-- TRAIN_1
         ...
         |-- TRAIN_11
         |-- TEST

Training

Download pre-trained MAE weight and put it under $PROJECT_ROOT$/pretrained_models

python tracking/train.py --script supersbt --config supersbt_base --save_dir ./output --mode multiple --nproc_per_node 4 --use_wandb 0
python -m torch.distributed.launch --nproc_per_node 8 lib/train/run_training.py --script supersbt --config supersbt_base --save_dir .

Replace --config with the desired model config under experiments/supersbt. We use wandb to record detailed training logs, in case you don't want to use wandb, set --use_wandb 0.

Evaluation

Download the model weights mentioned above

Put the downloaded weights on $PROJECT_ROOT$/output/checkpoints/train/supersbt

Change the corresponding values of lib/test/evaluation/local.py to the actual benchmark saving paths

Some testing examples:

• LaSOT or other off-line evaluated benchmarks (modify --dataset correspondingly)

python tracking/test.py supersbt supersbt_base --dataset lasot --threads 16 --num_gpus 4
python tracking/analysis_results.py # need to modify tracker configs and names

• GOT10K-test

python tracking/test.py supersbt supersbt_base_got --dataset got10k_test --threads 16 --num_gpus 4
python lib/test/utils/transform_got10k.py --tracker_name supersbt --cfg_name supersbt_base_got

• TrackingNet

python tracking/test.py supersbt supersbt_base --dataset trackingnet --threads 16 --num_gpus 4
python lib/test/utils/transform_trackingnet.py --tracker_name supersbt --cfg_name supersbt_base

Test FLOPs, and Speed

Note: The speeds reported in our paper were tested on a single RTX2080Ti GPU.

# Profiling supersbt_base
python tracking/profile_model.py --script supersbt --config supersbt_base
# Profiling vitb_384_mae_ce_32x4_ep300
python tracking/profile_model.py --script supersbt --config supersbt_base

Acknowledgments

• Thanks for the STARK, OStrack, and PyTracking library, which helps us to quickly implement our ideas.
• We use the implementation of the HiViT.

Citation

If our work is useful for your research, please consider citing:

@article{xie_2024_SuperSBT,
  title={Correlation-Embedded Transformer Tracking: A Single-Branch Framework},
  author={Xie, Fei and Yang, Wankou and Wang, Chunyu and Chu, Lei and Cao, Yue and Ma, Chao and Zeng, Wenjun},
  journal={arXiv preprint arXiv:2401.12743},
  year={2024}
}

@InProceedings{xie_2022_SBT,
    author    = {Xie, Fei and Wang, Chunyu and Wang, Guangting and Cao, Yue and Yang, Wankou and Zeng, Wenjun},
    title     = {Correlation-Aware Deep Tracking},
    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    year      = {2022},
}