moonbox-dev/ormbg-ONNX

ORMBG ONNX

ONNX builds of ORMBG, 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.09 GB	f44c17d92d57	vs fp32: IoU 0.9986, MAE 0.0003; DIS5K MAE 0.1461, S 0.720
int8	int8/model.onnx	0.04 GB	fe1df769e1fc	vs fp32: IoU 0.9785, MAE 0.0030; DIS5K MAE 0.1467, S 0.719

How to run

ORMBG is trained for pictures of people; it does poorly on general objects (see the DIS5K numbers), so use it for portraits. Made from the official ormbg.onnx in this model's repository (fp32); output 0 is the matte (the others are side outputs and features).

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

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.transpose(2, 0, 1)[None]
matte = sess.run(None, {sess.get_inputs()[0].name: x})[0][0, 0]
matte = (matte - matte.min()) / (matte.max() - matte.min())   # as the original inference does
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; a calibrated static int8 build agreed with fp32 at only IoU 0.82 and was replaced). int4: not made: the graph has no MatMul weights. 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)	176 MB	-	0.1462	0.720	0.673	197 ms	1.7 GB
fp16	88 MB	IoU 0.9986, MAE 0.0003	0.1461	0.720	0.673	176 ms	0.8 GB
int8	44 MB	IoU 0.9785, MAE 0.0030	0.1467	0.719	0.673	194 ms	1.9 GB

How these were made

From ORMBG at revision 6253b31824, 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

Apache-2.0. 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

Open Remove Background Model (ormbg)

>>> DEMO <<<

Join our Research Discord Group!

This model is a fully open-source background remover optimized for images with humans. It is based on Highly Accurate Dichotomous Image Segmentation research.

Inference

python ormbg/inference.py

Training

Install dependencies:

conda env create -f environment.yaml
conda activate ormbg

Replace dummy dataset with training dataset.

python3 ormbg/train_model.py

Research

I started training the model with synthetic images of the Human Segmentation Dataset crafted with LayerDiffuse. However, I noticed that the model struggles to perform well on real images.

Synthetic datasets have limitations for achieving great segmentation results. This is because artificial lighting, occlusion, scale or backgrounds create a gap between synthetic and real images. A "model trained solely on synthetic data generated with naïve domain randomization struggles to generalize on the real domain", see PEOPLESANSPEOPLE: A Synthetic Data Generator for Human-Centric Computer Vision (2022).

Next steps:

• Expand dataset with synthetic and real images
• Research on state of the art loss functions

Latest changes (26/07/2024):

• Created synthetic dataset with 10k images, crafted with BlenderProc
• Removed training data created with LayerDiffuse, since it lacks the accuracy needed
• Improved model performance (after 100k iterations):
  • F1: 0.9888 -> 0.9932
  • MAE: 0.0113 -> 0.008
  • Scores based on this validation dataset

05/07/2024

• Added P3M-10K dataset for training and validation
• Added AIM-500 dataset for training and validation
• Added PPM-100 dataset for training and validation
• Applied Grid Dropout to make the model smarter