TarViS: A Unified Architecture for Target-based Video Segmentation (CVPR'23 Highlight)

Ali Athar, Alexander Hermans, Jonathon Luiten, Deva Ramanan, Bastian Leibe

PDF | Paperswithcode | Cite

Video Instance Segmentation	Video Panoptic Segmentation	Video Object Segmentation	Point Exemplar-guided Tracking

Updates

30.04.2023: Complete code and pre-trained checkpoints uploaded.

Abstract

The general domain of video segmentation is currently fragmented into different tasks spanning multiple benchmarks. Despite rapid progress in the state-of-the-art, current methods are overwhelmingly task-specific and cannot conceptually generalize to other tasks. Inspired by recent approaches with multi-task capability, we propose TarViS: a novel, unified network architecture that can be applied to any task that requires segmenting a set of arbitrarily def ined ‘targets’ in video. Our approach is flexible with respect to how tasks define these targets, since it models the latter as abstract ‘queries’ which are then used to predict pixel-precise target masks. A single TarViS model can be trained jointly on a collection of datasets spanning different tasks, and can hot-swap between tasks during inference without any task-specific retraining. To demonstrate its effectiveness, we apply TarViS to four different tasks, namely Video Instance Segmentation (VIS), Video Panoptic Segmentation (VPS), Video Object Segmentation (VOS) and Point Exemplar-guided Tracking (PET). Our unified, jointly trained model achieves state-of-the-art performance on 5/7 benchmarks spanning these four tasks, and competitive performance on the remaining two

Setup

Building CUDA Kernels

Build the deformable attention CUDA kernels as follows:

cd tarvis/modelling/backbone/temporal_neck/ops
bash make.sh

If you already have the kernels installed from Mask2Former then you can skip this step since they're identical.

Directory Structure

For managing datasets, checkpoints and pretrained backbones, we use a single environment variable $TARVIS_WORKSPACE_DIR which points to a directory that is organized as follows:

├── $TARVIS_WORKSPACE_DIR
│   ├── checkpoints                   <- model weights saved here during training
│   ├── pretrained_backbones          <- ImageNet pretrained Swin backbone weights
│   ├── dataset_images                <- Images/videos for all datasets go here
|   |   ├── training
|   |   |   ├── ade20k               
|   |   |   |   ├── ADE_train_00000001.jpg
|   |   |   |   ├── ...
|   |   |   ├── cityscapes
|   |   |   |   ├── aachen
|   |   |   |   ├── ...
|   |   |   ├── coco
|   |   |   |   ├── 000000000009.jpg
|   |   |   |   ├── ...
|   |   |   ├── mapillary
|   |   |   |   ├── 0035fkbjWljhaftpVM37-g.jpg
|   |   |   |   ├── ...
|   |   |   ├── cityscapes_vps
|   |   |   |   ├── 0001_0001_frankfurt_000000_000279_newImg8bit.png
|   |   |   |   ├── ...
|   |   |   ├── kitti_step
|   |   |   |   ├── 0000
|   |   |   |   ├── ...
|   |   |   ├── vipseg
|   |   |   |   ├── 0_wHveSGjXyDY
|   |   |   |   ├── ...
|   |   |   ├── youtube_vis_2021
|   |   |   |   ├── 3245e049fb
|   |   |   |   ├── ...
|   |   |   ├── ovis
|   |   |   |   ├── 001ca3cb
|   |   |   |   ├── ...
|   |   |   ├── davis 
|   |   |   |   ├── bear
|   |   |   |   ├── ...
|   |   |   ├── burst   
|   |   |   |   ├── YFCC100M
|   |   |   |   ├── ...
|   |   ├── inference
|   |   |   ├── cityscapes_vps_val     <- same directory structure as training
|   |   |   ├── kitti_step_val
|   |   |   ├── vipseg
|   |   |   ├── youtube_vis_2021     
|   |   |   ├── ovis
|   |   |   ├── davis                 
|   |   |   ├── burst
|   |   |   |   ├── val
|   |   |   |   |   ├── YFCC100M
|   |   |   |   |   ├── ...
|   |   |   |   ├── test
|   |   |   |   |   ├── YFCC100M
|   |   |   |   |   ├── ...
|   ├── dataset_annotations            <- Annotations for all datasets go here
|   |   ├── training
|   |   |   ├── ade20k_panoptic    
|   |   |   |   ├── pan_maps
|   |   |   |   |   ├── ADE_train_00000001.png
|   |   |   |   ├── segments.json
|   |   |   ├── cityscapes_panoptic
|   |   |   |   ├── pan_maps
|   |   |   |   |   ├── aachen_000000_000019_gtFine_panoptic.png
|   |   |   |   |   ├── ...
|   |   |   |   ├── segments.json
|   |   |   ├── coco_panoptic
|   |   |   |   ├── pan_maps
|   |   |   |   |   ├── 000000000009.png
|   |   |   |   |   ├── ...
|   |   |   |   ├── segments.json
|   |   |   ├── mapillary_panoptic
|   |   |   |   ├── pan_maps    
|   |   |   |   |   ├── 0035fkbjWljhaftpVM37-g.png
|   |   |   |   |   ├── ...
|   |   |   |   ├── segments.json
|   |   |   ├── cityscapes_vps.json
|   |   |   ├── kitti_step.json
|   |   |   ├── vipseg
|   |   |   |   ├── panoptic_masks
|   |   |   |   |   ├── 0_wHveSGjXyDY
|   |   |   |   |   ├── ...
|   |   |   |   ├── video_info.json
|   |   |   ├── youtube_vis_2021.json
|   |   |   ├── ovis.json
|   |   |   ├── davis_semisupervised.json
|   |   |   ├── burst.json
|   |   ├── inference   
|   |   |   ├── cityscapes_vps
|   |   |   |   ├── im_all_info_val_city_vps.json
|   |   |   ├── vipseg
|   |   |   |   ├── val.json
|   |   |   ├── youtube_vis
|   |   |   |   ├── valid_2021.json
|   |   |   ├── ovis
|   |   |   |   ├── valid.json
|   |   |   ├── davis
|   |   |   |   ├── Annotations
|   |   |   |   ├── ImageSet_val.txt
|   |   |   |   ├── ImageSet_testdev.txt
|   |   |   ├── burst
|   |   |   |   ├── first_frame_annotations_val.json
|   |   |   |   ├── first_frame_annotations_test.json

You do not need to setup all the datasets if you only want to train/infer on a sub-set of them. For training the full model however, you need the complete directory tree above. All paths in-code are provided from tarvis/utils/paths.py, so any changes can be easily made there.

Partially complete workspace directory download

To make the setup easier, you can download a partially complete workspace directory from HERE. It does not include any image files, but it does contain all the JSON format annotation files. For some datasets, e.g. DAVIS, we converted the dataset annotations into a custom format to make it easier to re-use data loading code.

Important: The license agreement for this repository does not apply to the annotations. Please refer to the license agreements for the respective datasets if you wish to use these annotations.

Trained Models

Download links to trained checkpoints are available below. Each zipped directory contains a checkpoint file ending with .pth and a config.yaml file. We also provide checkpoints for the models after the pre-training step on augmented image sequences.

Backbone	Pretrain (augmented images)	Finetune (video)
ResNet-50	URL	URL
Swin-T	URL	URL
Swin-L	URL	URL

Inference

Run the following command with the desired dataset. The checkpoint path can be any of the finetuned models from the the download links above, or models that you train yourself. It is important that config.yaml is also present in the same directory as the checkpoint .pth file.

python tarvis/inference/main.py /path/to/checkpoint.pth --dataset {YOUTUBE_VIS,OVIS,CITYSCAPES_VPS,KITTI_STEP,DAVIS,BURST} --amp --output_dir /path/to/output/directory

Run python tarvis/inference/main.py --help to see additional options.

Training

First run the pretraining on augmented image datasets (COCO, ADE20k, Cityscapes, Mapillary). The following command does this assuming 8 nodes, each containing 4 GPUs:

torchrun --nnodes=8 --nproc_per_node=4 --rdzv_id=22021994 --rdzv_backend=c10d --rdzv_endpoint ${DDP_HOST_ADDRESS} tarvis/training/main.py --model_dir my_first_tarvis_pretraining --cfg pretrain_{resnet50_swin-tiny,swin-large}.yaml --amp  --cfg.MODEL.BACKBONE {ResNet50,SwinTiny,SwinLarge}

The --model_dir argument is interpreted relative to ${TARVIS_WORKSPACE_DIR}/checkpoints. You can also give as absolute path here.

Then run the finetuning on video data (YouTube-VIS, OVIS, KITTI-STEP, Cityscapes-VPS, DAVIS, BURST):

torchrun --nnodes=8 --nproc_per_node=4 --rdzv_id=22021994 --rdzv_backend=c10d --rdzv_endpoint ${DDP_HOST_ADDRESS} tarvis/training/main.py --model_dir my_first_tarvis_finetuning --cfg finetune_{resnet50_swin-tiny,swin-large}.yaml --amp  --finetune_from /path/to/my_first_tarvis_pretraining

For this, the --finetune_from argument should point to the pretraining directory (not the checkpoint file). It will automatically load the latest checkpoint from this directory.

Notes on Training

To run on a single GPU, omit torchrun and its arguments and just run the script normally: python tarvis/training/main.py ...
The training code applies gradient accumulation, so the batch size specified in the config (32 by default) is maintained regardless of how many GPUs are used e.g. if you use 8 GPUs, the gradients for 4 consecutive iterations will be accumulated before running optimizer.step(). The only constraint is that the batch size should be exactly divisible by the number of GPUs.
Training can be resumed from a checkpoint by running the training script as follows:

torchrun --nnodes=8 --nproc_per_node=4 --rdzv_id=22021994 --rdzv_backend=c10d --rdzv_endpoint ${DDP_HOST_ADDRESS} tarvis/training/main.py --restore_session /path/to/latest/checkpoint.pth

The data loading code uses OpenCV which sometimes behaves strangely on compute clusters by spawning too many worker threads. If you experience slow data loading times, try setting OMP_NUM_THREADS=1 and running the training script with --cv2_num_threads=2.
Training metrics are logged using tensorboard by default, but logging with weights and biases is also supported by provided the --wandb_session option.
With the default settings, each GPU should have at least 24GB VRAM to prevent OOM errors. If you want to train on smaller GPUs, consider reducing the input image resolution.
Run python tarvis/training/main.py --help to see additional options.

Cite

@inproceedings{athar2023tarvis,
  title={TarViS: A Unified Architecture for Target-based Video Segmentation},
  author={Athar, Ali and Hermans, Alexander and Luiten, Jonathon and Ramanan, Deva and Leibe, Bastian},
  booktitle={CVPR},
  year={2023}
}

Jinming-Su / TarViS