Image Segmentation
MLX
Safetensors
BiRefNet
moonbox
background-removal
dichotomous-image-segmentation
Instructions to use moonbox-ai/BiRefNet-MLX with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- MLX
How to use moonbox-ai/BiRefNet-MLX with MLX:
# Download the model from the Hub pip install huggingface_hub[hf_xet] hf download moonbox-ai/BiRefNet-MLX --local-dir BiRefNet-MLX
- BiRefNet
How to use moonbox-ai/BiRefNet-MLX with BiRefNet:
# Option 1: use with transformers from transformers import AutoModelForImageSegmentation birefnet = AutoModelForImageSegmentation.from_pretrained("moonbox-ai/BiRefNet-MLX", trust_remote_code=True)# Option 2: use with BiRefNet # Install from https://github.com/ZhengPeng7/BiRefNet from models.birefnet import BiRefNet model = BiRefNet.from_pretrained("moonbox-ai/BiRefNet-MLX") - Notebooks
- Google Colab
- Kaggle
- Local Apps Settings
- LM Studio
- Atomic Chat
|
Download README.md from moonbox-ai/BiRefNet-MLX: direct link, hf CLI and curl.
- Browser
- Download file 15.1 kB
-
https://huggingface.co/moonbox-ai/BiRefNet-MLX/resolve/main/README.md
- Command line
-
hf download hf://moonbox-ai/BiRefNet-MLX/README.md
-
curl -L -o README.md https://huggingface.co/moonbox-ai/BiRefNet-MLX/resolve/main/README.md
15.1 kB
| 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>,  | |
| <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} | |
| } | |
| ``` | |