Temporally Efficient Vision Transformer for Video Instance Segmentation

Temporally Efficient Vision Transformer for Video Instance Segmentation (CVPR 2022, Oral)

by Shusheng Yang^1,3, Xinggang Wang^{1 📧}, Yu Li⁴, Yuxin Fang¹, Jiemin Fang^1,2, Wenyu Liu¹, Xun Zhao³, Ying Shan³.

¹ School of EIC, HUST, ² AIA, HUST, ³ ARC Lab, Tencent PCG, ⁴ IDEA.

(^📧) corresponding author.

• This repo provides code, models and training/inference recipes for TeViT(Temporally Efficient Vision Transformer for Video Instance Segmentation).
• TeViT is a transformer-based end-to-end video instance segmentation framework. We build our framework upon the query-based instance segmentation methods, i.e., QueryInst.
• We propose a messenger shift mechanism in the transformer backbone, as well as a spatiotemporal query interaction head in the instance heads. These two designs fully utlizes both frame-level and instance-level temporal context information and obtains strong temporal modeling capacity with negligible extra computational cost.

Models and Main Results

• We provide both checkpoints and codalab server submissions on YouTube-VIS-2019 dataset.

• We have conducted multiple runs due to the training instability and checkpoints above are all the best one among multiple runs. The average performances are reported in our paper.
• Besides basic models, we also provide TeViT with ResNet-50 and Swin-L backbone, models are also trained on YouTube-VIS-2019 dataset.
• MST denotes multi-scale traning.

• Due to backbone limitations, TeViT models with ResNet-50 and Swin-L backbone are conducted with STQI Head only (i.e., without our proposed messenger shift mechanism).
• With Swin-L as backbone network, we apply more instance queries (i.e., from 100 to 300) and stronger data augmentation strategies. Both of them can further boost the final performance.

Installation

Prerequisites

• Linux
• Python 3.7+
• CUDA 10.2+
• GCC 5+

Prepare

• Clone the repository locally:

git clone https://github.com/hustvl/TeViT.git

• Create a conda virtual environment and activate it:

conda create --name tevit python=3.7.7
conda activate tevit

• Install YTVOS Version API from youtubevos/cocoapi:

pip install git+https://github.com/youtubevos/cocoapi.git#"egg=pycocotools&subdirectory=PythonAPI

• Install Python requirements

torch==1.9.0
torchvision==0.10.0
mmcv==1.4.8
pip install -r requirements.txt

• Please follow Docs to install MMDetection

python setup.py develop

• Download YouTube-VIS 2019 dataset from here, and organize dataset as follows:

TeViT
├── data
│   ├── youtubevis
│   │   ├── train
│   │   │   ├── 003234408d
│   │   │   ├── ...
│   │   ├── val
│   │   │   ├── ...
│   │   ├── annotations
│   │   │   ├── train.json
│   │   │   ├── valid.json

Inference

python tools/test_vis.py configs/tevit/tevit_msgshift.py $PATH_TO_CHECKPOINT

After inference process, the predicted results is stored in results.json, submit it to the evaluation server to get the final performance.

Training

• Download the COCO pretrained QueryInst with PVT-B1 backbone from here.
• Train TeViT with 8 GPUs:

./tools/dist_train.sh configs/tevit/tevit_msgshift.py 8 --no-validate --cfg-options load_from=$PATH_TO_PRETRAINED_WEIGHT

• Train TeViT with multi-scale data augmentation:

./tools/dist_train.sh configs/tevit/tevit_msgshift_mstrain.py 8 --no-validate --cfg-options load_from=$PATH_TO_PRETRAINED_WEIGHT

• The whole training process will cost about three hours with 8 TESLA V100 GPUs.
• To train TeViT with ResNet-50 or Swin-L backbone, please download the COCO pretrained weights from QueryInst.

Acknowledgement ❤️

This code is mainly based on mmdetection and QueryInst, thanks for their awesome work and great contributions to the computer vision community!

Citation

If you find our paper and code useful in your research, please consider giving a star ⭐ and citation 📝 :

@inproceedings{yang2022tevit,
  title={Temporally Efficient Vision Transformer for Video Instance Segmentation,
  author={Yang, Shusheng and Wang, Xinggang and Li, Yu and Fang, Yuxin and Fang, Jiemin and Liu and Zhao, Xun and Shan, Ying},
  booktitle =   {Proc. IEEE Conf. Computer Vision and Pattern Recognition (CVPR)},
  year      =   {2022}
}