ProPainter: Improving Propagation and Transformer for Video Inpainting

Shangchen Zhou Chongyi Li Kelvin C.K. Chan Chen Change Loy

S-Lab, Nanyang Technological University

ICCV 2023

⭐ If ProPainter is helpful to your projects, please help star this repo. Thanks! 🤗

📖 For more visual results, go checkout our project page

Update

2023.09.07: Our code and model are publicly available. 🐳
2023.09.01: This repo is created.

Results

🏂 Object Removal

🌈 Watermark Removal

🎨 Video Completion

Overview

Dependencies and Installation

Clone Repo

git clone https://github.com/sczhou/ProPainter.git

Create Conda Environment and Install Dependencies
```
conda env create -f environment.yaml
conda activate propainter
```
- Python >= 3.7
- PyTorch >= 1.6.0
- CUDA >= 9.2
- mmcv-full (refer the command table to install v1.4.8)

Get Started

Prepare pretrained models

Download our pretrained models from Releases V0.1.0 to the weights folder. (All pretrained models can also be automatically downloaded during the first inference.)

The directory structure will be arranged as:

weights
   |- ProPainter.pth
   |- recurrent_flow_completion.pth
   |- raft-things.pth
   |- i3d_rgb_imagenet.pt (for evaluating VFID metric)
   |- README.md

Quick test

We provide some examples in the inputs folder. Run the following commands to try it out:

# The first example (object removal)
python inference_propainter.py --video inputs/object_removal/bmx-trees --mask inputs/object_removal/bmx-trees_mask 
# The second example (watermark removal)
python inference_propainter.py --video inputs/watermark_removal/running_car.mp4 --mask inputs/watermark_removal/mask.png

The results will be saved in the results folder. To test your own videos, please prepare the input mp4 video (or split frames) and frame-wise mask(s).

If you want to specify the video resolution for processing or avoid running out of memory, you can use the --set_size flag and set the video size of --width and --height:

# process a 576x320 video
python inference_propainter.py --video inputs/watermark_removal/running_car.mp4 --mask inputs/watermark_removal/mask.png --set_size --height 320 --width 576

Regarding the issue of running out of memory, you can mitigate it by reducing the number of local neighbors through decreasing the --neighbor_length (default 20) or reducing the number of global references by increasing the --ref_stride (default 10).

Dataset preparation for training and evaluation

Dataset	YouTube-VOS	DAVIS
Description	For training (3,471) and evaluation (508)	For evaluation (50 in 90)
Images	[Official Link] (Download train and test all frames)	[Official Link] (2017, 480p, TrainVal)
Masks	[Google Drive] [Baidu Disk] (For reproducing paper results; provided in E²FGVI paper)

The training and test split files are provided in datasets/<dataset_name>. For each dataset, you should place JPEGImages to datasets/<dataset_name>. Resize all video frames to size 432x240 for training. Unzip downloaded mask files to datasets.

The datasets directory structure will be arranged as: (Note: please check it carefully)

datasets
   |- davis
      |- JPEGImages_432_240
         |- <video_name>
            |- 00000.jpg
            |- 00001.jpg
      |- test_masks
         |- <video_name>
            |- 00000.png
            |- 00001.png   
      |- train.json
      |- test.json
   |- youtube-vos
      |- JPEGImages_432_240
         |- <video_name>
            |- 00000.jpg
            |- 00001.jpg
      |- test_masks
         |- <video_name>
            |- 00000.png
            |- 00001.png
      |- train.json
      |- test.json

Evaluation

Run one of the following commands for evaluation:

 # For evaluating flow completion model
 python scripts/evaluate_flow_completion.py --dataset <dataset_name> --video_root <video_root> --mask_root <mask_root> --save_results
 # For evaluating ProPainter model
 python scripts/evaluate_propainter.py --dataset <dataset_name> --video_root <video_root> --mask_root <mask_root> --save_results

The scores and results will also be saved in the results_eval folder. Please --save_results for further evaluating temporal warping error.

Training

Our training configures are provided in train_flowcomp.json (for Recurrent Flow Completion Network) and train_propainter.json (for ProPainter).

Run one of the following commands for training:

 # For training Recurrent Flow Completion Network
 python train.py -c configs/train_flowcomp.json
 # For training ProPainter
 python train.py -c configs/train_propainter.json

You can run the same command to resume your training.

Citation

If you find our repo useful for your research, please consider citing our paper:

@inproceedings{zhou2023propainter,
   title={{ProPainter}: Improving Propagation and Transformer for Video Inpainting},
   author={Zhou, Shangchen and Li, Chongyi and Chan, Kelvin C.K and Loy, Chen Change},
   booktitle={Proceedings of IEEE International Conference on Computer Vision (ICCV)},
   year={2023}
}

License

This project is licensed under NTU S-Lab License 1.0. Redistribution and use should follow this license.

Contact

If you have any questions, please feel free to reach me out at [email protected].

Acknowledgement

This code is based on E²FGVI and STTN. Some code are brought from BasicVSR++. Thanks for their awesome works.

elwingaofork / propainter Goto Github PK

propainter's Introduction