README.md

metadata

library_name: birefnet
tags:
  - background-removal
  - mask-generation
  - Dichotomous Image Segmentation
  - Camouflaged Object Detection
  - Salient Object Detection
  - pytorch_model_hub_mixin
  - model_hub_mixin
repo_url: https://github.com/ZhengPeng7/BiRefNet
pipeline_tag: image-segmentation

This model has been pushed to the Hub using birefnet:

• Repo: https://github.com/ZhengPeng7/BiRefNet
• Docs: https://www.birefnet.top

How to use

# Download Codes
git clone https://github.com/ZhengPeng7/BiRefNet.git
cd BiRefNet

# Imports
from PIL import Image
import matplotlib.pyplot as plt
import torch
from torchvision import transforms

from models.birefnet import BiRefNet


# Load Model
device = 'cuda'
torch.set_float32_matmul_precision(['high', 'highest'][0])
model = BiRefNet.from_pretrained('zhengpeng7/birefnet')
model.to(device)
model.eval()
print('BiRefNet is ready to use.')

# Input Data
transform_image = transforms.Compose([
    transforms.Resize((1024, 1024)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

imagepath = 'PATH-TO-YOUR_IMAGE.jpg'
image = Image.open(imagepath)
input_images = transform_image(image).unsqueeze(0).to('cuda')

# Prediction
with torch.no_grad():
    preds = model(input_images)[-1].sigmoid().cpu()
pred = preds[0].squeeze()

# Show Results
plt.imshow(transforms.ToPILImage()(pred).resize(image.size), cmap='gray')
plt.show()

This BiRefNet for standard dichotomous image segmentation (DIS) is trained on DIS-TR and validated on DIS-TEs and DIS-VD.

This repo holds the official model weights of "Bilateral Reference for High-Resolution Dichotomous Image Segmentation" (arXiv 2024).

This repo contains the weights of BiRefNet proposed in our paper, which has achieved the SOTA performance on three tasks (DIS, HRSOD, and COD).

Go to my GitHub page for BiRefNet codes and the latest updates: https://github.com/ZhengPeng7/BiRefNet :)

Try our online demos for inference:

• Online Single Image Inference on Colab:
• Inference and evaluation of your given weights:
• Online Inference with GUI on Hugging Face with adjustable resolutions:

Acknowledgement:

• Many thanks to @fal for their generous support on GPU resources for training better BiRefNet models.
• Many thanks to @not-lain for his help on the better deployment of our BiRefNet model on Huggingface.

Citation

@article{zheng2024birefnet,
  title={Bilateral Reference for High-Resolution Dichotomous Image Segmentation},
  author={Zheng, Peng and Gao, Dehong and Fan, Deng-Ping and Liu, Li and Laaksonen, Jorma and Ouyang, Wanli and Sebe, Nicu},
  journal={arXiv},
  year={2024}
}