E2FGVI: End-to-End Framework for Flow-Guided Video Inpainting

repository·master·Indexed 22 days ago

https://github.com/mcg-nku/e2fgvi

An end-to-end framework for flow-guided video inpainting designed for high-quality video completion with faster speeds and lower computational costs. The project includes support for a standard model and a high-quality (HQ) model that supports arbitrary resolutions. It provides tools for inference via test.py, evaluation using evaluate.py, and training through train.py with provided JSON configurations.

Tokens
1.4K
Snippets
6
Records
6
Agent score
28%

What's inside E2FGVI

  1. Install E2FGVI

    master

    To install E2FGVI, clone the repository and use Conda to create an environment from the provided environment.yml file.

    Requirements:

    • Python >= 3.7
    • PyTorch >= 1.5
    • CUDA >= 9.2
    • mmcv-full (install following the official mmcv pipeline)

    If environment.yml fails, refer to the project's issue tracker for manual installation steps.

    git clone https://github.com/MCG-NKU/E2FGVI.git
    conda env create -f environment.yml
    conda activate e2fgvi
  2. Train E2FGVI models

    master

    Training is configured via JSON files. Use train.py with the appropriate config file.

    Configurations:

    • configs/train_e2fgvi.json: For the standard E2FGVI model.
    • configs/train_e2fgvi_hq.json: For the E2FGVI-HQ model.

    Commands:

    # Train E2FGVI
    python train.py -c configs/train_e2fgvi.json
    
    # Train E2FGVI-HQ
    python train.py -c configs/train_e2fgvi_hq.json

    To resume training, run the same command. To visualize training loss, use Tensorboard:

    tensorboard --logdir release_model
    python train.py -c configs/train_e2fgvi.json
  3. Prepare pretrained models for E2FGVI

    master

    Download the pretrained weights and place them in a directory named release_model.

    Models available:

    • E2FGVI-CVPR22.pth: Standard model.
    • E2FGVI-HQ-CVPR22.pth: High-quality model that supports arbitrary resolutions (does not resize input).
    • i3d_rgb_imagenet.pt: Required for evaluating the VFID metric.

    Required directory structure:

    release_model
       |- E2FGVI-CVPR22.pth
       |- E2FGVI-HQ-CVPR22.pth
       |- i3d_rgb_imagenet.pt
       |- README.md
    release_model
       |- E2FGVI-CVPR22.pth
       |- E2FGVI-HQ-CVPR22.pth
       |- i3d_rgb_imagenet.pt (for evaluating VFID metric)
       |- README.md
  4. Prepare datasets for training and evaluation

    master

    To use datasets like YouTube-VOS or DAVIS, follow this structure:

    1. Place JPEGImages in datasets/<dataset_name>.
    2. Run sh datasets/zip_dir.sh (after editing paths) to compress each video in JPEGImages into a .zip file.
    3. Unzip masks into the datasets directory.

    Target Directory Structure:

    datasets
       |- davis
          |- JPEGImages
             |- <video_name>.zip
          |- test_masks
             |- <video_name>
                |- 00000.png
          |- train.json
          |- test.json
       |- youtube-vos
          |- ...
    datasets
       |- davis
          |- JPEGImages
             |- <video_name>.zip
          |- test_masks
             |- <video_name>
                |- 00000.png
          |- train.json
          |- test.json
       |- youtube-vos
          |- ...
       |- zip_file.sh
  5. Evaluate trained models

    master

    Use evaluate.py to validate model performance on datasets. Results are saved in results/<model_name>_<dataset_name>.

    Arguments:

    • --model: e2fgvi or e2fgvi_hq.
    • --dataset: Name of the dataset (e.g., davis or youtube-vos).
    • --data_root: Path to the datasets/ directory.
    • --ckpt: Path to the checkpoint.
    • --save_results: (Optional) Add this flag if you need to validate temporal warping error.

    Commands:

    # Evaluate E2FGVI
    python evaluate.py --model e2fgvi --dataset <dataset_name> --data_root datasets/ --ckpt release_model/E2FGVI-CVPR22.pth
    
    # Evaluate E2FGVI-HQ
    python evaluate.py --model e2fgvi_hq --dataset <dataset_name> --data_root datasets/ --ckpt release_model/E2FGVI-HQ-CVPR22.pth
    python evaluate.py --model e2fgvi --dataset <dataset_name> --data_root datasets/ --ckpt release_model/E2FGVI-CVPR22.pth
  6. Run video inpainting inference

    master

    Use test.py to perform video inpainting. You can provide either a sequence of video frames or an .mp4 video file, along with its corresponding masks.

    Key Arguments:

    • --model: Choose e2fgvi or e2fgvi_hq.
    • --video: Path to the video file or directory of frames.
    • --mask: Path to the mask file or directory of masks.
    • --ckpt: Path to the pretrained checkpoint.
    • --set_size: (Optional) Required when using e2fgvi_hq to define custom output resolution.
    • --width / --height: (Optional) Used with --set_size to specify output dimensions (e.g., 1280x720).

    Note on Resolution:

    • e2fgvi scales input to a fixed 432x240 resolution.
    • e2fgvi_hq maintains the input resolution unless --set_size is used.

    Results are saved in the results directory.

    # Example 1: Using video frames
    python test.py --model e2fgvi --video examples/tennis --mask examples/tennis_mask --ckpt release_model/E2FGVI-CVPR22.pth
    
    # Example 2: Using an mp4 video
    python test.py --model e2fgvi_hq --video examples/schoolgirls.mp4 --mask examples/schoolgirls_mask --ckpt release_model/E2FGVI-HQ-CVPR22.pth
    
    # Example 3: Using HQ model with custom 720p resolution
    python test.py --model e2fgvi_hq --video <video_path> --mask <mask_path> --ckpt release_model/E2FGVI-HQ-CVPR22.pth --set_size --width 1280 --height 720