/TSAM

[CVPR2021] Official Implementation for Progressive Temporal Feature Alignment Network for Video Inpainting

Primary LanguagePython

Progressive Temporal Feature Alignment Network for Video Inpainting

This work is accepted in CVPR2021 as Poster. It proposed a new video inpainting approach that combines temporal convolution as well as optical flow approach.

Noted: This code is currently a beta version. Not gurantee to be fully correct.

Update

Optical Flow Davis | Optical Flow FVI | Mask Davis | Mask FVI | Checkpoint

Installation

torch==1.7.0
torchvision==0.8.1

Dataset

For FVI dataset, please refer to https://github.com/amjltc295/Free-Form-Video-Inpainting. For DAVIS dataset, please refer to https://davischallenge.org/.

File Structure

TSAM
└── data
    ├── checkpoints
    ├── model_weights
    ├── results
    ├── FVI
    ├── DAVIS    
    └── runs
└── code
    └── master
        └── TSAM
            └── ...

Prepare pretrained weights for training

Pretrained weights: download all the pretrained weights and put it under TSAM/data/model_weights

Model Name
TSM_imagenet_resent50_gated.pth weight
TSM_imagenet_resent50.pth weight

Training

FVI TSM moving object/curve masks:

CUDA_VISIBLE_DEVICES=0,1,2,3 python3 train.py --config config/config_pretrain.json --dataset_config dataset_configs/FVI_all_masks.json
CUDA_VISIBLE_DEVICES=0,1,2,3 python3 train.py --config config/config_finetune.json --dataset_config dataset_configs/FVI_all_masks.json

Testing

Change the train.py in training scripts to test.py, and add -p /pth/to/ckpt to the end.

DAVIS TSAM object removal:

CUDA_VISIBLE_DEVICES=0 python3 test.py --config config/config_finetune_davis.json --dataset_config dataset_configs/DAVIS_removal.json -p /pth/to/ckpt

Citation

@inproceedings{zou2020progressive,
  title={Progressive Temporal Feature Alignment Network for Video Inpainting},
  author={Xueyan Zou and Linjie Yang and Ding Liu and Yong Jae Lee},
  booktitle={CVPR},
  year={2021}
}

Acknowledgement

Part of the code is borrow from https://github.com/amjltc295/Free-Form-Video-Inpainting and https://github.com/researchmm/STTN. Thanks for their great works!