--- 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 ```shell # Download Codes git clone https://github.com/ZhengPeng7/BiRefNet.git cd BiRefNet ``` ```python # 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](https://arxiv.org/pdf/2401.03407)" (_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: [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/14Dqg7oeBkFEtchaHLNpig2BcdkZEogba?usp=drive_link) + **Inference and evaluation** of your given weights: [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1MaEiBfJ4xIaZZn0DqKrhydHB8X97hNXl#scrollTo=DJ4meUYjia6S) + **Online Inference with GUI on Hugging Face** with adjustable resolutions: [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/ZhengPeng7/BiRefNet_demo) ## 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} } ```