VisTR：基于 Transformer 的端到端视频实例分割

这是 VisTR 论文 的官方实现：

安装

我们提供了通过 conda 安装依赖的说明。 首先，在本地克隆仓库：

git clone https://github.com/Epiphqny/vistr.git

然后，安装 PyTorch 1.6 和 torchvision 0.7：

conda install pytorch==1.6.0 torchvision==0.7.0

安装 pycocotools：

conda install cython scipy
pip install -U 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'
pip install git+https://github.com/youtubevos/cocoapi.git#"egg=pycocotools&subdirectory=PythonAPI"

编译 DCN 模块（需要 GCC>=5.3，CUDA>=10.0）：

cd models/dcn
python setup.py build_ext --inplace

准备工作

从 CodeLab 或 YoutubeVIS 下载并解压 2019 年版 YouTubeVIS 的训练和验证图像及其标注文件。 我们期望目录结构如下：

VisTR
├── data
│   ├── train
│   ├── val
│   ├── annotations
│   │   ├── instances_train_sub.json
│   │   ├── instances_val_sub.json
├── models
...

下载在 COCO 数据集上预训练的 DETR 模型 Google Drive BaiduYun(提取码: alge)，并将其保存到预训练模型路径。

训练

训练该模型至少需要 32GB 显存的 GPU，我们是在 32GB 的 V100 显卡上进行实验的。（由于训练分辨率受显存限制，如果您拥有更大显存的 GPU 并希望进行实验，请与我联系，非常感谢）

要在单节点上使用 8 张 GPU 卡训练基线 VisTR 模型 18 个 epoch，运行以下命令：

python -m torch.distributed.launch --nproc_per_node=8 --use_env main.py --backbone resnet101/50 --ytvos_path /path/to/ytvos --masks --pretrained_weights /path/to/pretrained_path

推理

python inference.py --masks --model_path /path/to/model_weights --save_path /path/to/results.json

模型

我们提供了基线 VisTR 模型，并计划在未来增加更多模型。AP 值是通过将结果 JSON 文件提交到 CodeLab 系统计算得出的，推理时间则仅计算模型本身的推理时间（不包括数据加载和后处理）。

	名称	骨干网络	帧率 (FPS)	掩码 AP	模型	结果 JSON 压缩包	详细 AP
0	VisTR	R50	69.9	36.2	vistr_r50.pth	vistr_r50.zip
1	VisTR	R101	57.7	40.1	vistr_r101.pth	vistr_r101.zip

许可证

VisTR 采用 Apache 2.0 许可证发布。更多信息请参阅 LICENSE 文件。

致谢

我们感谢 DETR 开源项目所做的出色工作，部分代码源自该项目的修改。

引用

如果本项目对您的研究有所帮助，请考虑在您的出版物中引用我们的论文。BibTeX 格式如下：

@inproceedings{wang2020end,
  title={End-to-End Video Instance Segmentation with Transformers},
  author={Wang, Yuqing and Xu, Zhaoliang and Wang, Xinlong and Shen, Chunhua and Cheng, Baoshan and Shen, Hao and Xia, Huaxia},
  booktitle =  {Proc. IEEE Conf. Computer Vision and Pattern Recognition (CVPR)},
  year={2021}
}

VisTR 快速上手指南

VisTR 是一个基于 Transformer 的端到端视频实例分割（Video Instance Segmentation）模型。本指南将帮助开发者快速完成环境搭建、模型准备及推理测试。

1. 环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐 Ubuntu)
• GPU: 至少一张显存 ≥32GB 的 GPU（如 V100 32G），用于训练；推理可根据显存调整分辨率。
• 编译器: GCC ≥ 5.3
• CUDA: ≥ 10.0
• Python 环境: 建议使用 Conda 管理环境

2. 安装步骤

2.1 克隆代码库

git clone https://github.com/Epiphqny/vistr.git
cd vistr

2.2 安装基础依赖

创建并激活 Conda 环境，安装指定版本的 PyTorch 和 torchvision：

conda install pytorch==1.6.0 torchvision==0.7.0

安装 pycocotools 及相关依赖：

conda install cython scipy
pip install -U 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'
pip install git+https://github.com/youtubevos/cocoapi.git#"egg=pycocotools&subdirectory=PythonAPI"

2.3 编译 DCN 模块

进入目录编译可变形卷积（Deformable Convolution）模块：

cd models/dcn
python setup.py build_ext --inplace
cd ../..

3. 数据与模型准备

3.1 数据集准备

下载 YouTubeVIS 2019 版本的训练集和验证集（包含图片和标注文件）。 来源推荐：

• CodeLab 竞赛页面
• YoutubeVIS 官网

解压后，请确保目录结构如下：

VisTR
├── data
│   ├── train
│   ├── val
│   ├── annotations
│   │   ├── instances_train_sub.json
│   │   ├── instances_val_sub.json
├── models
...

3.2 下载预训练权重

下载在 COCO 数据集上预训练的 DETR 模型，并保存至本地路径。

• Google Drive: 下载链接
• 百度网盘: 下载链接 (提取码: alge)

4. 基本使用

4.1 模型推理 (Inference)

使用预训练好的权重对视频进行实例分割推理。以下命令展示了最基础的用法：

python inference.py --masks --model_path /path/to/model_weights --save_path /path/to/results.json

参数说明：

• --model_path: 替换为下载的 .pth 模型文件路径（例如 vistr_r50.pth）。
• --save_path: 指定输出结果 JSON 文件的保存路径。

注：如需重新训练模型，需确保 GPU 显存充足（单卡 32GB 或分布式多卡），训练命令参考原文档 Training 章节。