DOCSEARCHbyAlgolia
Submit Docs

BiRefNet: High-Resolution Dichotomous Image Segmentation

repository·main·Indexed 23 days ago

https://github.com/zhengpeng7/birefnet

BiRefNet is a high-resolution dichotomous image segmentation model designed for robust performance in tasks such as matting and general-purpose segmentation. It supports integration via HuggingFace Transformers, ONNX conversion, and GGUF for lightweight C++ inference. The project provides a variety of model variants for benchmarks (DIS, COD, HRSOD) and practical applications (General Use, Portrait Matting), along with tools for training, fine-tuning, and video inference.

Tokens
5.9K
Snippets
11
Records
24
Agent score
86%

What's inside BiRefNet

  1. Explore Third-Party BiRefNet Integrations

    main

    BiRefNet is integrated into various third-party applications and frameworks:

    • ComfyUI: Official integration by Comfy-Org, and specialized nodes by lldacing/ComfyUI_BiRefNet_ll, MoonHugo/ComfyUI-BiRefNet-Hugo, ZHO-ZHO-ZHO/ComfyUI-BiRefNet-ZHO, and viperyl/ComfyUI-BiRefNet.
    • Stable Diffusion WebUI: Added via dimitribarbot/sd-webui-birefnet in the Extras tab.
    • Invoke AI: Integrated as a node via veeliks/invoke_birefnet.
    • Blender: Used as an add-on for 2D asset generation via tin2tin/2D_Asset_Generator.
    • Online Services: Available via CNAPS.AI (workflow-based), unscreen.pro (video background removal), and fal.ai (API and UI).
    • Other: Lightweight C++ inference via GGUF, Rust implementation via Burn, and TensorRT deployment for high-speed inference.

  4. Prepare Datasets for BiRefNet

    main

    Datasets for training or testing can be downloaded from the following sources:

    • Organized Combined Sets (DIS, COD, HRSOD): Available via Google Drive or BaiduDisk.
    • Single Official Datasets: Available in the single_ones folder or via their official pages.

    Google Drive Links:

    BaiduDisk Links:

  5. Fine-tune BiRefNet on custom data

    main

    To fine-tune the model on your own dataset, follow these steps:

    1. Dataset Structure: Organize your data in ${data_root_dir}/TASK_NAME/DATASET_NAME. Both im (images) and gt (ground truth) folders must exist within each dataset folder.
    2. Task Configuration: Replace all instances of 'General' in the project with your custom task name.
    3. Adapt Settings in config.py:
      • sys_home_dir: The root folder containing codes, datasets, and weights.
      • testsets: Define your validation set.
      • training_set: Define your training set.
      • lambdas_pix_last: Adjust loss weights if switching between segmentation (classification) and matting (regression).
    4. Resume from existing weights: Use the resume argument in train.py.
      • Note on Epochs: Training continues from the epoch indicated in the filename. For example, if using BiRefNet-general-epoch_244.pth and you want to train for 50 more epochs, set your target epoch to 294.

  7. Download BiRefNet models from the Model Zoo

    main

    BiRefNet provides several model variants depending on your use case.

    Benchmark Models

    Use these models for academic comparison on standard benchmarks:

    • DIS: Trained on DIS5K-TR (Backbone: swin_v1_large)
    • COD: Trained on COD10K-TR, CAMO-TR (Backbone: swin_v1_large)
    • HRSOD: Various versions trained on DUTS-TR, HRSOD-TR, and UHRSD-TR (Backbone: swin_v1_large)

    Practical Application Models

    Use these models for real-world applications (general use or matting):

    • General Use (2048x2048): High-resolution model trained on a large mix of datasets (AIM-500, DIS, HIM2K, etc.) using swin_v1_large.
    • General Use: Standard resolution models available with swin_v1_large or swin_v1_tiny backbones.
    • General Matting: Optimized for matting tasks using swin_v1_large.
    • Portrait Matting: Specifically optimized for portrait segmentation using swin_v1_large.

    All models can be downloaded via the provided Google Drive links in the repository's Model Zoo section.

  9. Load BiRefNet using HuggingFace Transformers

    main

    You can load the BiRefNet model in a single line of code using the HuggingFace transformers library. This is the recommended way to quickly integrate the model into your projects. Ensure you set trust_remote_code=True to allow the loading of the model's custom architecture.

    from transformers import AutoModelForImageSegmentation
birefnet = AutoModelForImageSegmentation.from_pretrained('zhengpeng7/BiRefNet', trust_remote_code=True)

  10. Install BiRefNet via Conda and Pip

    main

    To set up the environment for BiRefNet, create a new Conda environment with Python 3.11 and install the required dependencies using requirements.txt. Note that PyTorch >= 2.5.0 is recommended for faster training performance via compilation.

    conda create -n birefnet python=3.11 -y && conda activate birefnet
pip install -r requirements.txt
    # PyTorch>=2.5.0 (I try to make everything as latest as possible) is used for faster training (~40%) with compilation.
conda create -n birefnet python=3.11 -y && conda activate birefnet
pip install -r requirements.txt

  11. Try BiRefNet Online Demos

    main

    Several interactive environments are available for testing BiRefNet without local setup:

    • Inference and Evaluation: Use Google Colab to run inference and evaluation with specific weights. Open in Colab.
    • GUI Inference: Use the Hugging Face Space for online inference with a graphical user interface and adjustable resolutions. Hugging Face Spaces.
    • Batch Inference: Use Google Colab for online multiple image inference. Open in Colab.

  12. Convert BiRefNet weights to ONNX

    main

    You can convert BiRefNet .pth weights to .onnx format for deployment.

    Performance Notes:

    • SwinL (default): ONNX inference takes ~165ms on an A100 GPU (at 1024x1024 resolution).
    • SwinT (lightweight): ONNX inference takes ~93.8ms on an A100 GPU.
    • Accuracy: Results from ONNX files may differ slightly from original .pth files, which is considered acceptable.

    Compatibility Warning: Ensure compatibility between onnxruntime-gpu, CUDA, and CUDNN. The tested environment uses torch==2.0.1 and cuda=11.8.