| --- |
| base_model: ZhengPeng7/BiRefNet |
| base_model_relation: quantized |
| license: mit |
| pipeline_tag: image-segmentation |
| library_name: onnx |
| tags: |
| - onnx |
| - moonbox |
| - birefnet |
| - background-removal |
| - image-segmentation |
| - dichotomous-image-segmentation |
| --- |
| |
| # BiRefNet ONNX |
|
|
| ONNX builds of [BiRefNet](https://huggingface.co/ZhengPeng7/BiRefNet), made by Moonbox for on-device use with ONNX Runtime (CUDA, DirectML, CPU). The |
| model is a single graph: each precision is one folder holding `model.onnx`. |
|
|
| ## Files |
|
|
| The last column is the benchmark: each precision's output against the fp32 graph's on the same eval pictures, and |
| its own score on the eval set (details under *How to run*). |
|
|
| | Precision | File | Size | sha256 | Benchmark | |
| |---|---|---|---|---| |
| | fp16 | [`fp16/model.onnx`](https://huggingface.co/moonbox-dev/BiRefNet-ONNX/blob/main/fp16/model.onnx) | 0.49 GB | `323232ec73a0` | vs fp32: IoU 0.9981, MAE 0.0002; DIS5K MAE 0.0141, S 0.953 | |
| | int8 | [`int8/model.onnx`](https://huggingface.co/moonbox-dev/BiRefNet-ONNX/blob/main/int8/model.onnx) | 0.31 GB | `4512960563f7` | vs fp32: IoU 0.9984, MAE 0.0002; DIS5K MAE 0.0142, S 0.953 | |
| | int4 | [`int4/model.onnx`](https://huggingface.co/moonbox-dev/BiRefNet-ONNX/blob/main/int4/model.onnx) | 0.24 GB | `d0c5185c9de8` | vs fp32: IoU 0.9899, MAE 0.0016; DIS5K MAE 0.0139, S 0.953 | |
|
|
| ## How to run |
|
|
| BiRefNet is a general-purpose dichotomous segmentation model (objects, people, fine structure). Made from the official fp32 ONNX `BiRefNet-general-epoch_244.onnx` from [BiRefNet's GitHub release v1](https://github.com/ZhengPeng7/BiRefNet/releases/tag/v1) (the general-use model); its single output is the matte (sigmoid applied). |
|
|
| **ONNX Runtime** (1024 x 1024 input, RGB in 0-1, ImageNet-normalized): |
|
|
| ```python |
| import numpy as np, onnxruntime as ort |
| from PIL import Image |
| |
| sess = ort.InferenceSession("int8/model.onnx", providers=["CUDAExecutionProvider", "CPUExecutionProvider"]) |
| img = Image.open("photo.jpg").convert("RGB") |
| x = np.asarray(img.resize((1024, 1024), Image.BILINEAR), np.float32) / 255 |
| x = (x - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225] # ImageNet normalization |
| x = x.transpose(2, 0, 1)[None] |
| matte = sess.run(None, {sess.get_inputs()[0].name: x})[0][0, 0] |
| alpha = Image.fromarray((np.clip(matte, 0, 1) * 255).astype(np.uint8)).resize(img.size, Image.BILINEAR) |
| img.putalpha(alpha); img.save("cutout.png") |
| ``` |
|
|
| **Precisions.** fp16: weights and math in fp16, inputs and outputs fp32. int8: weights int8 per output channel behind DequantizeLinear, activations in float (weight-only: no calibration). int4: MatMul weights 4-bit (MatMulNBits, block 32), the remaining convolution weights int8 weight-only. bf16 is not made: ONNX Runtime's CUDA provider has no bf16 convolution. |
|
|
| **Benchmark.** [DIS5K](https://github.com/xuebinqin/DIS) test sets, 250 pictures (DIS-TE1..4 and DIS-VD, 50 each), ONNX Runtime 1.30 CUDA provider on an RTX 3080 Laptop GPU. The table's last column gives the |
| precision's agreement with the fp32 graph (IoU of the masks at 0.5, mean absolute difference) and its own MAE against the |
| ground truth. Scored with PySODMetrics on masks shrunk to at most 2048 px on the long side (DIS5K's largest are about |
| 6000 x 4000); HCE (human correction effort) is given for the reference precision. |
|
|
| | Precision | Size | vs fp32 (mask IoU, MAE) | MAE | S-measure | max F | Time / picture | Peak VRAM | |
| |---|---|---|---|---|---|---|---| |
| | fp32 (reference, not published) | 973 MB | - | 0.0141 | 0.953 | 0.959 | 728 ms | 9.3 GB | |
| | fp16 | 490 MB | IoU 0.9981, MAE 0.0002 | 0.0141 | 0.953 | 0.959 | 583 ms | 7.7 GB | |
| | int8 | 315 MB | IoU 0.9984, MAE 0.0002 | 0.0142 | 0.953 | 0.959 | 729 ms | 8.6 GB | |
| | int4 | 244 MB | IoU 0.9899, MAE 0.0016 | 0.0139 | 0.953 | 0.960 | 731 ms | 8.7 GB | |
|
|
| ## How these were made |
|
|
| From [BiRefNet](https://huggingface.co/ZhengPeng7/BiRefNet) at revision `e2bf8e4460`, starting from an fp32 ONNX graph (where it came from is said above): |
| fp16 with onnxconverter-common (inputs and outputs kept fp32), int8 with ONNX Runtime static QDQ quantization |
| (per-channel int8 weights, uint8 activations, MinMax calibration on pictures outside the eval set), int4 with ONNX |
| Runtime's MatMulNBits (weight-only, block 32) where the graph is mostly matrix multiplications. bf16 is left out for |
| convolutional graphs: ONNX Runtime's CUDA provider has no bf16 convolution. Every graph was benchmarked before |
| upload; the `sha256` column above is what the files must hash to. |
|
|
| ## License |
|
|
| MIT. These files are a derivative of the original model and carry its license: it is copied in this |
| repository, and it governs your use of these files. The original model card follows unchanged. |
|
|
| --- |
|
|
| # Original model card |
|
|
| <h1 align="center">Bilateral Reference for High-Resolution Dichotomous Image Segmentation</h1> |
|
|
| <div align='center'> |
| <a href='https://scholar.google.com/citations?user=TZRzWOsAAAAJ' target='_blank'><strong>Peng Zheng</strong></a><sup> 1,4,5,6</sup>,  |
| <a href='https://scholar.google.com/citations?user=0uPb8MMAAAAJ' target='_blank'><strong>Dehong Gao</strong></a><sup> 2</sup>,  |
| <a href='https://scholar.google.com/citations?user=kakwJ5QAAAAJ' target='_blank'><strong>Deng-Ping Fan</strong></a><sup> 1*</sup>,  |
| <a href='https://scholar.google.com/citations?user=9cMQrVsAAAAJ' target='_blank'><strong>Li Liu</strong></a><sup> 3</sup>,  |
| <a href='https://scholar.google.com/citations?user=qQP6WXIAAAAJ' target='_blank'><strong>Jorma Laaksonen</strong></a><sup> 4</sup>,  |
| <a href='https://scholar.google.com/citations?user=pw_0Z_UAAAAJ' target='_blank'><strong>Wanli Ouyang</strong></a><sup> 5</sup>,  |
| <a href='https://scholar.google.com/citations?user=stFCYOAAAAAJ' target='_blank'><strong>Nicu Sebe</strong></a><sup> 6</sup> |
| </div> |
| |
| <div align='center'> |
| <sup>1 </sup>Nankai University  <sup>2 </sup>Northwestern Polytechnical University  <sup>3 </sup>National University of Defense Technology  <sup>4 </sup>Aalto University  <sup>5 </sup>Shanghai AI Laboratory  <sup>6 </sup>University of Trento  |
| </div> |
| |
| <div align="center" style="display: flex; justify-content: center; flex-wrap: wrap;"> |
| <a href='https://www.sciopen.com/article/pdf/10.26599/AIR.2024.9150038.pdf'><img src='https://img.shields.io/badge/Journal-Paper-red'></a>  |
| <a href='https://arxiv.org/pdf/2401.03407'><img src='https://img.shields.io/badge/arXiv-BiRefNet-red'></a>  |
| <a href='https://drive.google.com/file/d/1aBnJ_R9lbnC2dm8dqD0-pzP2Cu-U1Xpt/view?usp=drive_link'><img src='https://img.shields.io/badge/中文版-BiRefNet-red'></a>  |
| <a href='https://www.birefnet.top'><img src='https://img.shields.io/badge/Page-BiRefNet-red'></a>  |
| <a href='https://drive.google.com/drive/folders/1s2Xe0cjq-2ctnJBR24563yMSCOu4CcxM'><img src='https://img.shields.io/badge/Drive-Stuff-green'></a>  |
| <a href='LICENSE'><img src='https://img.shields.io/badge/License-MIT-yellow'></a>  |
| <a href='https://huggingface.co/spaces/ZhengPeng7/BiRefNet_demo'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20HF%20Spaces-BiRefNet-blue'></a>  |
| <a href='https://huggingface.co/ZhengPeng7/BiRefNet'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20HF%20Models-BiRefNet-blue'></a>  |
| <a href='https://colab.research.google.com/drive/14Dqg7oeBkFEtchaHLNpig2BcdkZEogba?usp=drive_link'><img src='https://img.shields.io/badge/Single_Image_Inference-F9AB00?style=for-the-badge&logo=googlecolab&color=525252'></a>  |
| <a href='https://colab.research.google.com/drive/1MaEiBfJ4xIaZZn0DqKrhydHB8X97hNXl#scrollTo=DJ4meUYjia6S'><img src='https://img.shields.io/badge/Inference_&_Evaluation-F9AB00?style=for-the-badge&logo=googlecolab&color=525252'></a>  |
| </div> |
|
|
|
|
| | *DIS-Sample_1* | *DIS-Sample_2* | |
| | :------------------------------: | :-------------------------------: | |
| | <img src="https://drive.google.com/thumbnail?id=1ItXaA26iYnE8XQ_GgNLy71MOWePoS2-g&sz=w400" /> | <img src="https://drive.google.com/thumbnail?id=1Z-esCujQF_uEa_YJjkibc3NUrW4aR_d4&sz=w400" /> | |
|
|
| This repo is the official implementation of "[**Bilateral Reference for High-Resolution Dichotomous Image Segmentation**](https://arxiv.org/pdf/2401.03407.pdf)" (___CAAI AIR 2024___). |
|
|
| Visit our GitHub repo: [https://github.com/ZhengPeng7/BiRefNet](https://github.com/ZhengPeng7/BiRefNet) for more details -- **codes**, **docs**, and **model zoo**! |
|
|
| ## How to use |
|
|
| ### 0. Install Packages: |
| ``` |
| pip install -qr https://raw.githubusercontent.com/ZhengPeng7/BiRefNet/main/requirements.txt |
| ``` |
|
|
| ### 1. Load BiRefNet: |
|
|
| #### Use codes + weights from HuggingFace |
| > Only use the weights on HuggingFace -- Pro: No need to download BiRefNet codes manually; Con: Codes on HuggingFace might not be latest version (I'll try to keep them always latest). |
|
|
| ```python |
| # Load BiRefNet with weights |
| from transformers import AutoModelForImageSegmentation |
| birefnet = AutoModelForImageSegmentation.from_pretrained('ZhengPeng7/BiRefNet', trust_remote_code=True) |
| ``` |
|
|
| #### Use codes from GitHub + weights from HuggingFace |
| > Only use the weights on HuggingFace -- Pro: codes are always latest; Con: Need to clone the BiRefNet repo from my GitHub. |
|
|
| ```shell |
| # Download codes |
| git clone https://github.com/ZhengPeng7/BiRefNet.git |
| cd BiRefNet |
| ``` |
|
|
| ```python |
| # Use codes locally |
| from models.birefnet import BiRefNet |
| |
| # Load weights from Hugging Face Models |
| birefnet = BiRefNet.from_pretrained('ZhengPeng7/BiRefNet') |
| ``` |
|
|
| #### Use codes from GitHub + weights from local space |
| > Only use the weights and codes both locally. |
|
|
| ```python |
| # Use codes and weights locally |
| import torch |
| from utils import check_state_dict |
| |
| birefnet = BiRefNet(bb_pretrained=False) |
| state_dict = torch.load(PATH_TO_WEIGHT, map_location='cpu') |
| state_dict = check_state_dict(state_dict) |
| birefnet.load_state_dict(state_dict) |
| ``` |
|
|
| #### Use the loaded BiRefNet for inference |
| ```python |
| # Imports |
| from PIL import Image |
| import matplotlib.pyplot as plt |
| import torch |
| from torchvision import transforms |
| from models.birefnet import BiRefNet |
| |
| birefnet = ... # -- BiRefNet should be loaded with codes above, either way. |
| torch.set_float32_matmul_precision(['high', 'highest'][0]) |
| birefnet.to('cuda') |
| birefnet.eval() |
| birefnet.half() |
| |
| def extract_object(birefnet, imagepath): |
| # Data settings |
| image_size = (1024, 1024) |
| transform_image = transforms.Compose([ |
| transforms.Resize(image_size), |
| transforms.ToTensor(), |
| transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) |
| ]) |
| |
| image = Image.open(imagepath) |
| input_images = transform_image(image).unsqueeze(0).to('cuda').half() |
| |
| # Prediction |
| with torch.no_grad(): |
| preds = birefnet(input_images)[-1].sigmoid().cpu() |
| pred = preds[0].squeeze() |
| pred_pil = transforms.ToPILImage()(pred) |
| mask = pred_pil.resize(image.size) |
| image.putalpha(mask) |
| return image, mask |
| |
| # Visualization |
| plt.axis("off") |
| plt.imshow(extract_object(birefnet, imagepath='PATH-TO-YOUR_IMAGE.jpg')[0]) |
| plt.show() |
| |
| ``` |
|
|
| ### 2. Use inference endpoint locally: |
| > You may need to click the *deploy* and set up the endpoint by yourself, which would make some costs. |
| ``` |
| import requests |
| import base64 |
| from io import BytesIO |
| from PIL import Image |
| |
| |
| YOUR_HF_TOKEN = 'xxx' |
| API_URL = "xxx" |
| headers = { |
| "Authorization": "Bearer {}".format(YOUR_HF_TOKEN) |
| } |
| |
| def base64_to_bytes(base64_string): |
| # Remove the data URI prefix if present |
| if "data:image" in base64_string: |
| base64_string = base64_string.split(",")[1] |
| |
| # Decode the Base64 string into bytes |
| image_bytes = base64.b64decode(base64_string) |
| return image_bytes |
| |
| def bytes_to_base64(image_bytes): |
| # Create a BytesIO object to handle the image data |
| image_stream = BytesIO(image_bytes) |
| |
| # Open the image using Pillow (PIL) |
| image = Image.open(image_stream) |
| return image |
| |
| def query(payload): |
| response = requests.post(API_URL, headers=headers, json=payload) |
| return response.json() |
| |
| output = query({ |
| "inputs": "https://hips.hearstapps.com/hmg-prod/images/gettyimages-1229892983-square.jpg", |
| "parameters": {} |
| }) |
| |
| output_image = bytes_to_base64(base64_to_bytes(output)) |
| output_image |
| ``` |
|
|
|
|
| > 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 "[<ins>Bilateral Reference for High-Resolution Dichotomous Image Segmentation</ins>](https://arxiv.org/pdf/2401.03407)" (_CAAI AIR 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 **Image Inference** on Colab: [](https://colab.research.google.com/drive/14Dqg7oeBkFEtchaHLNpig2BcdkZEogba?usp=drive_link) |
| + **Online Inference with GUI on Hugging Face** with adjustable resolutions: [](https://huggingface.co/spaces/ZhengPeng7/BiRefNet_demo) |
| + **Inference and evaluation** of your given weights: [](https://colab.research.google.com/drive/1MaEiBfJ4xIaZZn0DqKrhydHB8X97hNXl#scrollTo=DJ4meUYjia6S) |
| <img src="https://drive.google.com/thumbnail?id=12XmDhKtO1o2fEvBu4OE4ULVB2BK0ecWi&sz=w1080" /> |
|
|
| ## Acknowledgement: |
|
|
| + Many thanks to @Freepik for their generous support on GPU resources for training higher resolution BiRefNet models and more of my explorations. |
| + Many thanks to @fal for their generous support on GPU resources for training better general 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={CAAI Artificial Intelligence Research}, |
| volume = {3}, |
| pages = {9150038}, |
| year={2024} |
| } |
| ``` |
|
|
