Background Matting v2

Use these settings to balance performance and quality based on your target resolution:

• HD: backbone_scale=0.25, refine_sample_pixels=80000
• 4K: backbone_scale=0.125, refine_sample_pixels=320000

The repository includes several Python scripts for experimenting with the model:

• inference_images.py: Performs matting on a directory of images.
• inference_video.py: Performs matting on a video file.
• inference_webcam.py: Provides an interactive matting demo using your webcam.
• inference_speed_test.py: Measures the tensor throughput of the model to verify real-time capability.

Note on Performance: The inference_video.py script is not optimized for real-time use because video encoding/decoding lacks hardware acceleration and parallelization. For production real-time performance, you must implement hardware encoding/decoding and parallel frame loading to the GPU.

To train the model, you must configure data_path.pth to point to your dataset.

Following the original paper's methodology:

1. Use train_base.pth to train only the base model until convergence.
2. Use train_refine.pth to train the entire network end-to-end.

The model can be used with onnxruntime. Note that ONNX inference is reported to be slower than PyTorch/TorchScript. Pre-exported models are available for HD and 4K configurations.

Python Example:

import onnxruntime
import numpy as np

sess = onnxruntime.InferenceSession('PATH_TO_MODEL.onnx')

src = np.random.normal(size=(1, 3, 1080, 1920)).astype(np.float32)
bgr = np.random.normal(size=(1, 3, 1080, 1920)).astype(np.float32)

pha, fgr = sess.run(['pha', 'fgr'], {'src': src, 'bgr': bgr})

Compatibility Note: The architecture uses cropping and patch replacement which may cause issues in some backends. The provided ONNX models use roi_align for cropping and scatter_element for replacing patches. Other configurations can be exported using export_onnx.py.

Model Weights

• GitHub Releases
• Google Drive

Datasets

• Official Dataset Download Page

The Background Matting v2 model is compatible with multiple runtimes. You can integrate it into your workflow using:

• PyTorch
• TorchScript
• TensorFlow
• ONNX

For detailed integration instructions, refer to the doc/model_usage.md file.

For research purposes, you can use the architecture defined in the /model directory. You must manually instantiate the MattingRefine class and load the state dictionary from a checkpoint file.

Requirements:

• Access to the model module in the repository.
• A trained checkpoint file (.pth).

import torch
from model import MattingRefine

device = torch.device('cuda')
precision = torch.float32

model = MattingRefine(backbone='mobilenetv2',
                      backbone_scale=0.25,
                      refine_mode='sampling',
                      refine_sample_pixels=80_000)

model.load_state_dict(torch.load('PATH_TO_CHECKPOINT.pth'))
model = model.eval().to(precision).to(device)

src = torch.rand(1, 3, 1080, 1920).to(precision).to(device)
bgr = torch.rand(1, 3, 1080, 1920).to(precision).to(device)

with torch.no_grad():
    pha, fgr = model(src, bgr)[:2]

Model Inputs

• src: (B, 3, H, W) - Source image with RGB channels normalized to 0 ~ 1.
• bgr: (B, 3, H, W) - Background image with RGB channels normalized to 0 ~ 1.

Model Outputs

For regular use cases, only pha and fgr are required. To composite the result, use: com = pha * fgr + (1 - pha) * bgr.

• pha: (B, 1, H, W) - Alpha matte (0 ~ 1).
• fgr: (B, 3, H, W) - Foreground (0 ~ 1).

Intermediate outputs (training/debugging):

• pha_sm: (B, 1, Hc, Wc) - Coarse alpha matte.
• fgr_sm: (B, 3, Hc, Wc) - Coarse foreground.
• err_sm: (B, 1, Hc, Wc) - Coarse error prediction map.
• ref_sm: (B, 1, H/4, W/4) - Refinement regions (1 denotes a refined 4x4 patch).