MeMOTR

MeMOTR: 基于长期记忆增强transformer的多目标跟踪（ICCV 2023）的官方实现。

作者：高若鹏，王利民。

MeMOTR是一个完全端到端的基于Transformer的记忆增强多目标跟踪器。我们利用长期记忆注入和定制的记忆注意力层，显著提高了关联性能。

新闻 :fire:

• 2024.05.09：我们发布了MOTIP，一种将多目标跟踪任务视为ID预测问题的新视角 :telescope:。

• 2024.02.21：我们在arxiv版本（补充部分）中添加了SportsMOT的结果。如果您在SportsMOT比较中引用我们的跟踪器，我们将不胜感激 :chart_with_upwards_trend:。

• 2023.12.24：我们发布了BDD100K上的代码、脚本和检查点 :car:。

• 2023.12.13：我们实现了一个jupyter notebook，可以在您自己的视频上运行我们的模型 :movie_camera:。

• 2023.11.07：我们发布了SportsMOT上的脚本和检查点 :basketball:。

• 2023.08.24：我们发布了DanceTrack上的脚本和检查点 :dancer:。

• 2023.08.09：我们发布了主要代码。更多配置、脚本和检查点将很快发布 :soon:。

安装

conda create -n MeMOTR python=3.10  # 创建虚拟环境
# 我记得我使用了Python 3.10的一些新特性，但我不确定具体是哪些。
conda activate MeMOTR               # 激活环境
conda install pytorch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 pytorch-cuda=11.7 -c pytorch -c nvidia
# 我们的代码主要在PyTorch 1.13.1上运行，
# 但它应该也兼容早期的PyTorch版本（如1.12.1）。
# 然而，太早的pytorch版本可能会导致一些需要修复的问题，因为我们使用了pytorch的一些新提出的特性（例如，ResNet50_Weights）。
conda install matplotlib pyyaml scipy tqdm tensorboard
pip install opencv-python

您还需要编译Deformable Attention CUDA算子：

# 来自 https://github.com/fundamentalvision/Deformable-DETR
cd ./models/ops/
sh make.sh
# 如果需要，您可以测试这些算子：
python test.py

数据

您应该将解压后的MOT17和CrowdHuman数据集分别放入DATADIR/MOT17/images/和DATADIR/CrowdHuman/images/。然后通过运行相应的脚本生成真实标注文件：./data/gen_mot17_gts.py和./data/gen_crowdhuman_gts.py。

最后，您应该得到以下数据集结构：

DATADIR/
  ├── DanceTrack/
  │ ├── train/
  │ ├── val/
  │ ├── test/
  │ ├── train_seqmap.txt
  │ ├── val_seqmap.txt
  │ └── test_seqmap.txt
  ├── SportsMOT/
  │ ├── train/
  │ ├── val/
  │ ├── test/
  │ ├── train_seqmap.txt
  │ ├── val_seqmap.txt
  │ └── test_seqmap.txt
  ├── MOT17/
  │ ├── images/
  │ │ ├── train/     # 从MOT17解压
  │ │ └── test/      # 从MOT17解压
  │ └── gts/
  │   └── train/     # 由./data/gen_mot17_gts.py生成
  └── CrowdHuman/
    ├── images/
    │ ├── train/     # 从CrowdHuman解压
    │ └── val/       # 从CrowdHuman解压
    └── gts/
      ├── train/     # 由./data/gen_crowdhuman_gts.py生成
      └── val/       # 由./data/gen_crowdhuman_gts.py生成

预训练

我们使用官方的DAB-Deformable-DETR（使用R50骨干网络）在COCO数据集上预训练的权重初始化我们的模型，您也可以在这里下载我们使用的检查点。然后将检查点放在项目根目录下。

DanceTrack脚本

训练

使用8个GPU在DanceTrack上训练MeMOTR（建议使用>=32 GB内存的GPU，如V100-32GB或其他）：

python -m torch.distributed.run --nproc_per_node=8 main.py --use-distributed --config-path ./configs/train_dancetrack.yaml --outputs-dir ./outputs/memotr_dancetrack/ --batch-size 1 --data-root <您的数据目录路径>

如果您的GPU内存小于32 GB，我们还实现了一个内存优化版本（通过运行选项--use-checkpoint），如论文中讨论的那样，我们使用梯度检查点来减少分配的GPU内存。以下训练脚本只需要约10 GB的GPU内存：

python -m torch.distributed.run --nproc_per_node=8 main.py --use-distributed --config-path ./configs/train_dancetrack.yaml --outputs-dir ./outputs/memotr_dancetrack/ --batch-size 1 --data-root <您的数据目录路径> --use-checkpoint

提交和评估

您可以使用此脚本在DanceTrack验证集上评估训练好的模型：

python main.py --mode eval --data-root <您的数据目录路径> --eval-mode specific --eval-model <检查点文件名> --eval-dir ./outputs/memotr_dancetrack/ --eval-threads <您的GPU数量>

对于提交，您可以使用以下脚本：

python -m torch.distributed.run --nproc_per_node=8 main.py --mode submit --submit-dir ./outputs/memotr_dancetrack/ --submit-model <检查点文件名> --use-distributed --data-root <您的数据目录路径>

此外，如果您只想直接使用我们训练好的检查点进行评估或提交，您可以在这里获取我们在论文中使用的检查点。然后将此检查点放入./outputs/memotr_dancetrack/并运行上述脚本。

MOT17脚本

提交

对于提交，您可以使用以下脚本：

python -m torch.distributed.run --nproc_per_node=8 main.py --mode submit --config-path ./outputs/memotr_mot17/train/config.yaml --submit-dir ./outputs/memotr_mot17/ --submit-model <检查点文件名> --use-distributed --data-root <你的数据目录路径>

另外，你可以直接在这里下载我们训练好的检查点。然后将其放入 ./outputs/memotr_mot17/ 并运行上述脚本进行提交，以获取MOT17测试集的提交文件。

SportsMOT和其他数据集的脚本

你可以替换DanceTrack脚本中的--config-path。例如，将./configs/train_dancetrack.yaml改为./configs/train_sportsmot.yaml以在SportsMOT上进行训练。

结果

DanceTrack测试集上的多目标跟踪

SportsMOT测试集上的多目标跟踪

对于所有实验，我们不使用额外数据（如CrowdHuman）进行训练。

MOT17测试集上的多目标跟踪

BDD100K验证集上的多类别多目标跟踪

联系方式

• 高若鹏: ruopenggao@gmail.com

引用

@InProceedings{MeMOTR,
    author    = {Gao, Ruopeng and Wang, Limin},
    title     = {{MeMOTR}: Long-Term Memory-Augmented Transformer for Multi-Object Tracking},
    booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
    month     = {October},
    year      = {2023},
    pages     = {9901-9910}
}

星标

致谢

• Deformable DETR
• DAB DETR
• MOTR
• TrackEval
• CV-Framework
• Star History