Blame for README.md

---
base_model: ZhengPeng7/BiRefNet
base_model_relation: quantized
license: mit
pipeline_tag: image-segmentation
library_name: mlx
tags:
- mlx
- moonbox
- birefnet
- background-removal
- dichotomous-image-segmentation
---

# BiRefNet MLX

MLX builds of [BiRefNet](https://huggingface.co/ZhengPeng7/BiRefNet), made by Moonbox for Apple Silicon Macs. Each precision is a folder; keep a folder's
files together. Only the parts listed below are converted: the rest of the pipeline comes from the original
repository unless the section below says otherwise.

## Files

The last column is how far each build is from the original PyTorch model on the same inputs (see the section below
for what was measured).

| Precision | File | Size | sha256 | Checked against the original |
|---|---|---|---|---|
| bf16 | [`bf16/config.json`](https://huggingface.co/moonbox-ai/BiRefNet-MLX/blob/main/bf16/config.json) | 0.1 kB | `c13e80f54aa0` | MAE <= 0.0002, mask IoU >= 0.9978 |
| bf16 | [`bf16/mlx_model.py`](https://huggingface.co/moonbox-ai/BiRefNet-MLX/blob/main/bf16/mlx_model.py) | 20.1 kB | `bd3cb5681703` | MAE <= 0.0002, mask IoU >= 0.9978 |
| bf16 | [`bf16/model.safetensors`](https://huggingface.co/moonbox-ai/BiRefNet-MLX/blob/main/bf16/model.safetensors) | 440.4 MB | `f8732f5d4348` | MAE <= 0.0002, mask IoU >= 0.9978 |
| fp16 | [`fp16/config.json`](https://huggingface.co/moonbox-ai/BiRefNet-MLX/blob/main/fp16/config.json) | 0.1 kB | `56e8081f5c47` | MAE <= 0.00003, mask IoU >= 0.9996 |
| fp16 | [`fp16/mlx_model.py`](https://huggingface.co/moonbox-ai/BiRefNet-MLX/blob/main/fp16/mlx_model.py) | 20.1 kB | `bd3cb5681703` | MAE <= 0.00003, mask IoU >= 0.9996 |
| fp16 | [`fp16/model.safetensors`](https://huggingface.co/moonbox-ai/BiRefNet-MLX/blob/main/fp16/model.safetensors) | 440.4 MB | `b3365c802c6d` | MAE <= 0.00003, mask IoU >= 0.9996 |
| 8bit | [`8bit/config.json`](https://huggingface.co/moonbox-ai/BiRefNet-MLX/blob/main/8bit/config.json) | 0.2 kB | `2e5c6e85149b` | MAE <= 0.0003, mask IoU >= 0.9971 |
| 8bit | [`8bit/mlx_model.py`](https://huggingface.co/moonbox-ai/BiRefNet-MLX/blob/main/8bit/mlx_model.py) | 20.1 kB | `bd3cb5681703` | MAE <= 0.0003, mask IoU >= 0.9971 |
| 8bit | [`8bit/model.safetensors`](https://huggingface.co/moonbox-ai/BiRefNet-MLX/blob/main/8bit/model.safetensors) | 257.9 MB | `86c4b7f5817a` | MAE <= 0.0003, mask IoU >= 0.9971 |
| 6bit | [`6bit/config.json`](https://huggingface.co/moonbox-ai/BiRefNet-MLX/blob/main/6bit/config.json) | 0.2 kB | `3ab1929c1dab` | MAE <= 0.0009, mask IoU >= 0.9913 |
| 6bit | [`6bit/mlx_model.py`](https://huggingface.co/moonbox-ai/BiRefNet-MLX/blob/main/6bit/mlx_model.py) | 20.1 kB | `bd3cb5681703` | MAE <= 0.0009, mask IoU >= 0.9913 |
| 6bit | [`6bit/model.safetensors`](https://huggingface.co/moonbox-ai/BiRefNet-MLX/blob/main/6bit/model.safetensors) | 209.3 MB | `f930d6704d5b` | MAE <= 0.0009, mask IoU >= 0.9913 |
| 4bit | [`4bit/config.json`](https://huggingface.co/moonbox-ai/BiRefNet-MLX/blob/main/4bit/config.json) | 0.2 kB | `da0823e07a57` | below threshold (experimental): MAE <= 0.0022, mask IoU >= 0.9772 |
| 4bit | [`4bit/mlx_model.py`](https://huggingface.co/moonbox-ai/BiRefNet-MLX/blob/main/4bit/mlx_model.py) | 20.1 kB | `bd3cb5681703` | below threshold (experimental): MAE <= 0.0022, mask IoU >= 0.9772 |
| 4bit | [`4bit/model.safetensors`](https://huggingface.co/moonbox-ai/BiRefNet-MLX/blob/main/4bit/model.safetensors) | 160.6 MB | `2f9835d6d342` | below threshold (experimental): MAE <= 0.0022, mask IoU >= 0.9772 |

## How to run

BiRefNet (the general model, Swin-L backbone) ported to MLX. Each folder holds `model.safetensors`, `config.json`
and `mlx_model.py`, the port (only `mlx`, `numpy` and `pillow` needed):

```python
import numpy as np, mlx_model
from PIL import Image
model = mlx_model.load("fp16")
rgb = np.array(Image.open("in.jpg").convert("RGB"))
alpha = mlx_model.matte(model, rgb)                   # HxW uint8; the input is resized to 1024x1024 as in handler.py
Image.fromarray(np.dstack([rgb, alpha])).save("cutout.png")
```

The port runs the released configuration (Swin-L at full and half resolution, ASPP-deformable decoder) with the
BatchNorms folded into the convolutions and torchvision's `deform_conv2d` written out in MLX. The quantized folders
quantize the Swin linears (group size 64); convolutions and norms stay bf16.

| Precision | Size | MAE vs PyTorch fp32 (worst image) | Mask IoU (worst) | Status |
|---|---|---|---|---|
| `fp16/` | 440 MB | 0.00003 | 0.9996 | passes |
| `bf16/` | 440 MB | 0.0002 | 0.9978 | passes |
| `8bit/` | 258 MB | 0.0003 | 0.9971 | passes |
| `6bit/` | 209 MB | 0.0009 | 0.9913 | passes |
| `4bit/` | 161 MB | 0.0022 | 0.9772 | **below threshold (IoU >= 0.98), experimental** |

About 1.2 s per image (1024x1024) on an Apple M5.

## How these were made

From [BiRefNet](https://huggingface.co/ZhengPeng7/BiRefNet) at revision `e2bf8e4460`: the original weights loaded in MLX and stored as bf16 or fp16, or
quantized with MLX's affine quantization (`8bit`, `6bit`, `4bit`, group size 64; layers MLX can't quantize, such as
convolutions and norms, stay bf16). Every build was checked against the original 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>,&thinsp;
    <a href='https://scholar.google.com/citations?user=0uPb8MMAAAAJ' target='_blank'><strong>Dehong Gao</strong></a><sup> 2</sup>,&thinsp;
    <a href='https://scholar.google.com/citations?user=kakwJ5QAAAAJ' target='_blank'><strong>Deng-Ping Fan</strong></a><sup> 1*</sup>,&thinsp;
    <a href='https://scholar.google.com/citations?user=9cMQrVsAAAAJ' target='_blank'><strong>Li Liu</strong></a><sup> 3</sup>,&thinsp;
    <a href='https://scholar.google.com/citations?user=qQP6WXIAAAAJ' target='_blank'><strong>Jorma Laaksonen</strong></a><sup> 4</sup>,&thinsp;
    <a href='https://scholar.google.com/citations?user=pw_0Z_UAAAAJ' target='_blank'><strong>Wanli Ouyang</strong></a><sup> 5</sup>,&thinsp;
    <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&ensp;  <sup>2 </sup>Northwestern Polytechnical University&ensp;  <sup>3 </sup>National University of Defense Technology&ensp; <sup>4 </sup>Aalto University&ensp;  <sup>5 </sup>Shanghai AI Laboratory&ensp;  <sup>6 </sup>University of Trento&ensp; 
</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>&ensp; 
  <a href='https://arxiv.org/pdf/2401.03407'><img src='https://img.shields.io/badge/arXiv-BiRefNet-red'></a>&ensp; 
  <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>&ensp; 
  <a href='https://www.birefnet.top'><img src='https://img.shields.io/badge/Page-BiRefNet-red'></a>&ensp; 
  <a href='https://drive.google.com/drive/folders/1s2Xe0cjq-2ctnJBR24563yMSCOu4CcxM'><img src='https://img.shields.io/badge/Drive-Stuff-green'></a>&ensp; 
  <a href='LICENSE'><img src='https://img.shields.io/badge/License-MIT-yellow'></a>&ensp; 
  <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>&ensp; 
  <a href='https://huggingface.co/ZhengPeng7/BiRefNet'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20HF%20Models-BiRefNet-blue'></a>&ensp; 
  <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>&ensp; 
  <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>&ensp; 
</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: [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/14Dqg7oeBkFEtchaHLNpig2BcdkZEogba?usp=drive_link)
+ **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)  
+ **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)
<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}
}
```