moonbox-dev/BiRefNet-ONNX

BiRefNet ONNX

ONNX builds of 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	0.49 GB	323232ec73a0	vs fp32: IoU 0.9981, MAE 0.0002; DIS5K MAE 0.0141, S 0.953
int8	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	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 (the general-use model); its single output is the matte (sigmoid applied).

ONNX Runtime (1024 x 1024 input, RGB in 0-1, ImageNet-normalized):

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 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 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

Bilateral Reference for High-Resolution Dichotomous Image Segmentation

Peng Zheng ^1,4,5,6,  Dehong Gao ²,  Deng-Ping Fan ^1*,  Li Liu ³,  Jorma Laaksonen ⁴,  Wanli Ouyang ⁵,  Nicu Sebe ⁶

¹ Nankai University  ² Northwestern Polytechnical University  ³ National University of Defense Technology  ⁴ Aalto University  ⁵ Shanghai AI Laboratory  ⁶ University of Trento 

                   

DIS-Sample_1	DIS-Sample_2

This repo is the official implementation of "Bilateral Reference for High-Resolution Dichotomous Image Segmentation" (CAAI AIR 2024).

Visit our GitHub repo: 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).

# 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.

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

# 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.

# 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

# 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 "Bilateral Reference for High-Resolution Dichotomous Image Segmentation" (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:
• Online Inference with GUI on Hugging Face with adjustable resolutions:
• Inference and evaluation of your given weights:

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}
}