DOCSEARCHbyAlgolia
Submit Docs

ViTMatte Documentation

repository·main·Indexed 19 days ago

https://github.com/hustvl/vitmatte

An efficient and robust image matting system leveraging pretrained Vision Transformers (ViTs) with a hybrid attention mechanism and detail capture module. The project provides two model sizes, ViTMatte-S and ViTMatte-B, and includes tools for training on Adobe and COCO datasets, running inference on Composition-1k, and evaluating results via a Detectron2-based CLI.

Tokens
2K
Snippets
9
Records
12
Agent score
69%

What's inside ViTMatte

  1. Train ViTMatte-S or ViTMatte-B

    main

    ViTMatte is available in two sizes: ViTMatte-S and ViTMatte-B.

    Before training, you can modify the specific configuration files:

    • configs/ViTMatte_S_100ep.py for the Small model.
    • configs/ViTMatte_B_100ep.py for the Base model.

    Training outputs are saved to ViTMatte/output_of_train/.

    # To train ViTMatte-S
python main.py \
    --config-file configs/ViTMatte_S_100ep.py \
    --num-gpus 2

# To train ViTMatte-B
python main.py \
    --config-file configs/ViTMatte_B_100ep.py \
    --num-gpus 2

  2. Install ViTMatte via Conda and Pip

    main

    To set up the ViTMatte environment, create a dedicated Conda virtual environment with Python 3.8.8, activate it, and install the required dependencies from the requirements.txt file located in the repository root.

    conda create -n ViTMatte python==3.8.8
conda activate ViTMatte

cd path/to/ViTMatte
pip install -r requriments.txt

  3. Get started with ViTMatte

    main

    To begin using ViTMatte, refer to the following guides for specific tasks:

    • Installation: Set up the environment and dependencies.
    • Train: Instructions for training the model.
    • Test: Instructions for evaluating the model on datasets.

  4. Prepare the training dataset for ViTMatte

    main

    To train ViTMatte, you need to prepare the Adobe Image Matting Dataset and the COCO dataset.

    1. Adobe Image Matting Dataset: Download the dataset and merge the Adobe-licensed images and Other folders to include all 431 foregrounds and alphas.
    2. COCO Dataset: Download the COCO dataset.
    3. Configure Paths: Update the data paths in the configuration files located at ViTMatte/configs/common/dataloader to point to your local dataset locations.

  6. Prepare pretrained weights for ViTMatte

    main

    ViTMatte requires pretrained weights from DINO (for ViT-S) or MAE (for ViT-B). You can either source these manually or use the provided preprocessing script to download and prepare them automatically.

    To download and preprocess weights, run:

    cd ViTMatte/pretrained
python preprocess.py

  8. Run inference on the Composition-1k dataset

    main

    To perform inference using ViTMatte, use the inference.py script. You must provide paths to your configuration files, model checkpoints, the directory where you want to save results, and the directory containing the dataset.

    Prerequisite: You need the preprocessed Composition-1k Dataset, which can be obtained from the MatteFormer repository.

    python inference.py \
    --config-dir path/to/config \
    --checkpoint-dir path/to/checkpoint \
    --inference-dir path/to/inference \
    --data-dir path/to/dataset

  9. Evaluate inference results with evaluation.py

    main

    You can perform a quick evaluation of your inference results using evaluation.py.

    Important Note on Quantitative Results: The results produced by evaluation.py are for quick evaluation only. For official quantitative comparisons and fair benchmarking, use the official MATLAB code from DIM (Deep-Image-Matting) instead of this script.

    python evaluation.py \
    --pred-dir path/to/inference \
    --label-dir path/to/composition_1k/alpha \
    --trimap-dir path/to/composition_1k/trimap

  10. Configure training via LazyConfig overrides

    main

    The main.py script utilizes LazyConfig.apply_overrides to allow users to modify configuration parameters directly from the command line without editing the Python config files. This is done using the --opts flag followed by key-value pairs.

    Example of overriding the maximum iterations and the output directory:

    python main.py --config-file configs/example.py --opts train.max_iter 5000 train.output_dir ./output_test

  11. Train or evaluate ViTMatte using main.py

    main

    The main.py script serves as the CLI entrypoint for training and evaluating the ViTMatte model using Detectron2's LazyConfig system. It supports both full training runs and evaluation-only modes.

    Usage Modes

    1. Training: By default, the script runs the training loop defined in do_train. It instantiates the model, optimizer, and dataloaders from the provided configuration file and manages checkpoints, learning rate scheduling, and logging.
    2. Evaluation Only: Use the --eval-only flag to skip training and only run inference on the test dataset to evaluate model performance.

    Configuration Requirements

    The script expects a LazyConfig python file that defines the following structure:

    • model: An instantiable module.
    • dataloader.train: An instantiable training dataloader.
    • dataloader.test: An instantiable test dataloader.
    • dataloader.evaluator: An instantiable evaluator for the test set.
    • optimizer: An instantiable optimizer.
    • lr_multiplier: An instantiable fvcore scheduler.
    • train: A dictionary containing:
      • output_dir (str): Directory for logs and checkpoints.
      • init_checkpoint (str): Path to the initial checkpoint.
      • amp.enabled (bool): Whether to use Automatic Mixed Precision.
      • max_iter (int): Total number of training iterations.
      • eval_period (int): Iteration interval for evaluation.
      • log_period (int): Iteration interval for logging.
      • device (str): Target device (e.g., 'cuda').
      • checkpointer (dict): Configuration for the periodic checkpointer.
      • ddp (dict): Configuration for Distributed Data Parallel.

    CLI Arguments

    The script uses default_argument_parser() from Detectron2, which typically includes:

    • --config-file: Path to the Python configuration file.
    • --eval-only: Flag to run evaluation instead of training.
    • --resume: Flag to resume training from the last checkpoint.
    • --num-gpus: Number of GPUs to use.
    • --opts: Arbitrary configuration overrides (e.g., train.max_iter=1000).
    # Example: Train the model
python main.py --config-file configs/your_config.py

# Example: Evaluate a checkpoint
python main.py --config-file configs/your_config.py --eval-only

# Example: Resume training
python main.py --config-file configs/your_config.py --resume

  12. Run a demo with ViTMatte

    main

    You can run a quick demo to mat an image using its corresponding trimap. The script run_one_image.py will process the demo images and save the results in the ./demo directory. You can also use this script with your own image and trimap by providing the correct paths.

    To run the demo, use the following command structure:

    python run_one_image.py \
    --model vitmatte-s \
    --checkpoint-dir path/to/checkpoint