jennyxue commited on
Commit
fa9b642
·
verified ·
1 Parent(s): 3fe428e

Model card and the original license

Browse files
Files changed (1) hide show
  1. README.md +292 -0
README.md ADDED
@@ -0,0 +1,292 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ base_model: ZhengPeng7/BiRefNet
3
+ base_model_relation: quantized
4
+ license: mit
5
+ pipeline_tag: image-segmentation
6
+ library_name: onnx
7
+ tags:
8
+ - onnx
9
+ - moonbox
10
+ - birefnet
11
+ - background-removal
12
+ - image-segmentation
13
+ - dichotomous-image-segmentation
14
+ ---
15
+
16
+ # BiRefNet ONNX
17
+
18
+ ONNX builds of [BiRefNet](https://huggingface.co/ZhengPeng7/BiRefNet), made by Moonbox for on-device use with ONNX Runtime (CUDA, DirectML, CPU). The
19
+ model is a single graph: each precision is one folder holding `model.onnx`.
20
+
21
+ ## Files
22
+
23
+ The last column is the benchmark: each precision's output against the fp32 graph's on the same eval pictures, and
24
+ its own score on the eval set (details under *How to run*).
25
+
26
+ | Precision | File | Size | sha256 | Benchmark |
27
+ |---|---|---|---|---|
28
+ | 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 |
29
+
30
+ ## How to run
31
+
32
+ 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).
33
+
34
+ **ONNX Runtime** (1024 x 1024 input, RGB in 0-1, ImageNet-normalized):
35
+
36
+ ```python
37
+ import numpy as np, onnxruntime as ort
38
+ from PIL import Image
39
+
40
+ sess = ort.InferenceSession("int8/model.onnx", providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
41
+ img = Image.open("photo.jpg").convert("RGB")
42
+ x = np.asarray(img.resize((1024, 1024), Image.BILINEAR), np.float32) / 255
43
+ x = (x - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225] # ImageNet normalization
44
+ x = x.transpose(2, 0, 1)[None]
45
+ matte = sess.run(None, {sess.get_inputs()[0].name: x})[0][0, 0]
46
+ alpha = Image.fromarray((np.clip(matte, 0, 1) * 255).astype(np.uint8)).resize(img.size, Image.BILINEAR)
47
+ img.putalpha(alpha); img.save("cutout.png")
48
+ ```
49
+
50
+ **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.
51
+
52
+ **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
53
+ precision's agreement with the fp32 graph (IoU of the masks at 0.5, mean absolute difference) and its own MAE against the
54
+ ground truth. Scored with PySODMetrics on masks shrunk to at most 2048 px on the long side (DIS5K's largest are about
55
+ 6000 x 4000); HCE (human correction effort) is given for the reference precision.
56
+
57
+ | Precision | Size | vs fp32 (mask IoU, MAE) | MAE | S-measure | max F | Time / picture | Peak VRAM |
58
+ |---|---|---|---|---|---|---|---|
59
+ | fp32 (reference, not published) | 973 MB | - | 0.0141 | 0.953 | 0.959 | 728 ms | 9.3 GB |
60
+ | fp16 | 490 MB | IoU 0.9981, MAE 0.0002 | 0.0141 | 0.953 | 0.959 | 583 ms | 7.7 GB |
61
+ | int8 | 315 MB | IoU 0.9984, MAE 0.0002 | 0.0142 | 0.953 | 0.959 | 729 ms | 8.6 GB |
62
+ | int4 | 244 MB | IoU 0.9899, MAE 0.0016 | 0.0139 | 0.953 | 0.960 | 731 ms | 8.7 GB |
63
+
64
+ ## How these were made
65
+
66
+ From [BiRefNet](https://huggingface.co/ZhengPeng7/BiRefNet) at revision `e2bf8e4460`, starting from an fp32 ONNX graph (where it came from is said above):
67
+ fp16 with onnxconverter-common (inputs and outputs kept fp32), int8 with ONNX Runtime static QDQ quantization
68
+ (per-channel int8 weights, uint8 activations, MinMax calibration on pictures outside the eval set), int4 with ONNX
69
+ Runtime's MatMulNBits (weight-only, block 32) where the graph is mostly matrix multiplications. bf16 is left out for
70
+ convolutional graphs: ONNX Runtime's CUDA provider has no bf16 convolution. Every graph was benchmarked before
71
+ upload; the `sha256` column above is what the files must hash to.
72
+
73
+ ## License
74
+
75
+ MIT. These files are a derivative of the original model and carry its license: it is copied in this
76
+ repository, and it governs your use of these files. The original model card follows unchanged.
77
+
78
+ ---
79
+
80
+ # Original model card
81
+
82
+ <h1 align="center">Bilateral Reference for High-Resolution Dichotomous Image Segmentation</h1>
83
+
84
+ <div align='center'>
85
+ <a href='https://scholar.google.com/citations?user=TZRzWOsAAAAJ' target='_blank'><strong>Peng Zheng</strong></a><sup> 1,4,5,6</sup>,&thinsp;
86
+ <a href='https://scholar.google.com/citations?user=0uPb8MMAAAAJ' target='_blank'><strong>Dehong Gao</strong></a><sup> 2</sup>,&thinsp;
87
+ <a href='https://scholar.google.com/citations?user=kakwJ5QAAAAJ' target='_blank'><strong>Deng-Ping Fan</strong></a><sup> 1*</sup>,&thinsp;
88
+ <a href='https://scholar.google.com/citations?user=9cMQrVsAAAAJ' target='_blank'><strong>Li Liu</strong></a><sup> 3</sup>,&thinsp;
89
+ <a href='https://scholar.google.com/citations?user=qQP6WXIAAAAJ' target='_blank'><strong>Jorma Laaksonen</strong></a><sup> 4</sup>,&thinsp;
90
+ <a href='https://scholar.google.com/citations?user=pw_0Z_UAAAAJ' target='_blank'><strong>Wanli Ouyang</strong></a><sup> 5</sup>,&thinsp;
91
+ <a href='https://scholar.google.com/citations?user=stFCYOAAAAAJ' target='_blank'><strong>Nicu Sebe</strong></a><sup> 6</sup>
92
+ </div>
93
+
94
+ <div align='center'>
95
+ <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;
96
+ </div>
97
+
98
+ <div align="center" style="display: flex; justify-content: center; flex-wrap: wrap;">
99
+ <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;
100
+ <a href='https://arxiv.org/pdf/2401.03407'><img src='https://img.shields.io/badge/arXiv-BiRefNet-red'></a>&ensp;
101
+ <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;
102
+ <a href='https://www.birefnet.top'><img src='https://img.shields.io/badge/Page-BiRefNet-red'></a>&ensp;
103
+ <a href='https://drive.google.com/drive/folders/1s2Xe0cjq-2ctnJBR24563yMSCOu4CcxM'><img src='https://img.shields.io/badge/Drive-Stuff-green'></a>&ensp;
104
+ <a href='LICENSE'><img src='https://img.shields.io/badge/License-MIT-yellow'></a>&ensp;
105
+ <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;
106
+ <a href='https://huggingface.co/ZhengPeng7/BiRefNet'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20HF%20Models-BiRefNet-blue'></a>&ensp;
107
+ <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;
108
+ <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;
109
+ </div>
110
+
111
+
112
+ | *DIS-Sample_1* | *DIS-Sample_2* |
113
+ | :------------------------------: | :-------------------------------: |
114
+ | <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" /> |
115
+
116
+ 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___).
117
+
118
+ Visit our GitHub repo: [https://github.com/ZhengPeng7/BiRefNet](https://github.com/ZhengPeng7/BiRefNet) for more details -- **codes**, **docs**, and **model zoo**!
119
+
120
+ ## How to use
121
+
122
+ ### 0. Install Packages:
123
+ ```
124
+ pip install -qr https://raw.githubusercontent.com/ZhengPeng7/BiRefNet/main/requirements.txt
125
+ ```
126
+
127
+ ### 1. Load BiRefNet:
128
+
129
+ #### Use codes + weights from HuggingFace
130
+ > 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).
131
+
132
+ ```python
133
+ # Load BiRefNet with weights
134
+ from transformers import AutoModelForImageSegmentation
135
+ birefnet = AutoModelForImageSegmentation.from_pretrained('ZhengPeng7/BiRefNet', trust_remote_code=True)
136
+ ```
137
+
138
+ #### Use codes from GitHub + weights from HuggingFace
139
+ > Only use the weights on HuggingFace -- Pro: codes are always latest; Con: Need to clone the BiRefNet repo from my GitHub.
140
+
141
+ ```shell
142
+ # Download codes
143
+ git clone https://github.com/ZhengPeng7/BiRefNet.git
144
+ cd BiRefNet
145
+ ```
146
+
147
+ ```python
148
+ # Use codes locally
149
+ from models.birefnet import BiRefNet
150
+
151
+ # Load weights from Hugging Face Models
152
+ birefnet = BiRefNet.from_pretrained('ZhengPeng7/BiRefNet')
153
+ ```
154
+
155
+ #### Use codes from GitHub + weights from local space
156
+ > Only use the weights and codes both locally.
157
+
158
+ ```python
159
+ # Use codes and weights locally
160
+ import torch
161
+ from utils import check_state_dict
162
+
163
+ birefnet = BiRefNet(bb_pretrained=False)
164
+ state_dict = torch.load(PATH_TO_WEIGHT, map_location='cpu')
165
+ state_dict = check_state_dict(state_dict)
166
+ birefnet.load_state_dict(state_dict)
167
+ ```
168
+
169
+ #### Use the loaded BiRefNet for inference
170
+ ```python
171
+ # Imports
172
+ from PIL import Image
173
+ import matplotlib.pyplot as plt
174
+ import torch
175
+ from torchvision import transforms
176
+ from models.birefnet import BiRefNet
177
+
178
+ birefnet = ... # -- BiRefNet should be loaded with codes above, either way.
179
+ torch.set_float32_matmul_precision(['high', 'highest'][0])
180
+ birefnet.to('cuda')
181
+ birefnet.eval()
182
+ birefnet.half()
183
+
184
+ def extract_object(birefnet, imagepath):
185
+ # Data settings
186
+ image_size = (1024, 1024)
187
+ transform_image = transforms.Compose([
188
+ transforms.Resize(image_size),
189
+ transforms.ToTensor(),
190
+ transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
191
+ ])
192
+
193
+ image = Image.open(imagepath)
194
+ input_images = transform_image(image).unsqueeze(0).to('cuda').half()
195
+
196
+ # Prediction
197
+ with torch.no_grad():
198
+ preds = birefnet(input_images)[-1].sigmoid().cpu()
199
+ pred = preds[0].squeeze()
200
+ pred_pil = transforms.ToPILImage()(pred)
201
+ mask = pred_pil.resize(image.size)
202
+ image.putalpha(mask)
203
+ return image, mask
204
+
205
+ # Visualization
206
+ plt.axis("off")
207
+ plt.imshow(extract_object(birefnet, imagepath='PATH-TO-YOUR_IMAGE.jpg')[0])
208
+ plt.show()
209
+
210
+ ```
211
+
212
+ ### 2. Use inference endpoint locally:
213
+ > You may need to click the *deploy* and set up the endpoint by yourself, which would make some costs.
214
+ ```
215
+ import requests
216
+ import base64
217
+ from io import BytesIO
218
+ from PIL import Image
219
+
220
+
221
+ YOUR_HF_TOKEN = 'xxx'
222
+ API_URL = "xxx"
223
+ headers = {
224
+ "Authorization": "Bearer {}".format(YOUR_HF_TOKEN)
225
+ }
226
+
227
+ def base64_to_bytes(base64_string):
228
+ # Remove the data URI prefix if present
229
+ if "data:image" in base64_string:
230
+ base64_string = base64_string.split(",")[1]
231
+
232
+ # Decode the Base64 string into bytes
233
+ image_bytes = base64.b64decode(base64_string)
234
+ return image_bytes
235
+
236
+ def bytes_to_base64(image_bytes):
237
+ # Create a BytesIO object to handle the image data
238
+ image_stream = BytesIO(image_bytes)
239
+
240
+ # Open the image using Pillow (PIL)
241
+ image = Image.open(image_stream)
242
+ return image
243
+
244
+ def query(payload):
245
+ response = requests.post(API_URL, headers=headers, json=payload)
246
+ return response.json()
247
+
248
+ output = query({
249
+ "inputs": "https://hips.hearstapps.com/hmg-prod/images/gettyimages-1229892983-square.jpg",
250
+ "parameters": {}
251
+ })
252
+
253
+ output_image = bytes_to_base64(base64_to_bytes(output))
254
+ output_image
255
+ ```
256
+
257
+
258
+ > This BiRefNet for standard dichotomous image segmentation (DIS) is trained on **DIS-TR** and validated on **DIS-TEs and DIS-VD**.
259
+
260
+ ## 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_).
261
+
262
+ This repo contains the weights of BiRefNet proposed in our paper, which has achieved the SOTA performance on three tasks (DIS, HRSOD, and COD).
263
+
264
+ Go to my GitHub page for BiRefNet codes and the latest updates: https://github.com/ZhengPeng7/BiRefNet :)
265
+
266
+
267
+ #### Try our online demos for inference:
268
+
269
+ + 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)
270
+ + **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)
271
+ + **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)
272
+ <img src="https://drive.google.com/thumbnail?id=12XmDhKtO1o2fEvBu4OE4ULVB2BK0ecWi&sz=w1080" />
273
+
274
+ ## Acknowledgement:
275
+
276
+ + Many thanks to @Freepik for their generous support on GPU resources for training higher resolution BiRefNet models and more of my explorations.
277
+ + Many thanks to @fal for their generous support on GPU resources for training better general BiRefNet models.
278
+ + Many thanks to @not-lain for his help on the better deployment of our BiRefNet model on HuggingFace.
279
+
280
+
281
+ ## Citation
282
+
283
+ ```
284
+ @article{zheng2024birefnet,
285
+ title={Bilateral Reference for High-Resolution Dichotomous Image Segmentation},
286
+ author={Zheng, Peng and Gao, Dehong and Fan, Deng-Ping and Liu, Li and Laaksonen, Jorma and Ouyang, Wanli and Sebe, Nicu},
287
+ journal={CAAI Artificial Intelligence Research},
288
+ volume = {3},
289
+ pages = {9150038},
290
+ year={2024}
291
+ }
292
+ ```