moonbox-ai/ormbg-MLX

ORMBG MLX

MLX builds of ORMBG, 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	0.1 kB	09fea8fff5a6	alpha MAE <= 0.0026, mask IoU >= 0.987
bf16	bf16/mlx_model.py	8.2 kB	8cf5c652dcbb	alpha MAE <= 0.0026, mask IoU >= 0.987
bf16	bf16/model.safetensors	88.0 MB	39025b410188	alpha MAE <= 0.0026, mask IoU >= 0.987
fp16	fp16/config.json	0.1 kB	484d65d02eff	alpha MAE <= 0.0004, mask IoU >= 0.998
fp16	fp16/mlx_model.py	8.2 kB	8cf5c652dcbb	alpha MAE <= 0.0004, mask IoU >= 0.998
fp16	fp16/model.safetensors	88.0 MB	a7f5a13a20a9	alpha MAE <= 0.0004, mask IoU >= 0.998

How to run

ORMBG ported to MLX. Each folder holds model.safetensors, config.json and mlx_model.py, the port (only mlx and numpy needed):

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, as ormbg/inference.py (input resized to 1024x1024)
Image.fromarray(np.dstack([rgb, alpha])).save("cutout.png")

The BatchNorm layers are folded into the convolutions. There are no 8bit / 6bit / 4bit builds: the network is all convolutions, which MLX doesn't quantize.

Precision	Size	Alpha MAE vs PyTorch fp32 (worst image)	Mask IoU (worst)	Time per image on an M5
fp16/	88 MB	0.0004	0.998	0.17 s
bf16/	88 MB	0.0026	0.987	0.18 s

How these were made

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

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