File size: 5,891 Bytes
b2bb56a c9ccb85 b2bb56a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | ---
base_model: schirrmacher/ormbg
base_model_relation: quantized
license: apache-2.0
pipeline_tag: image-segmentation
library_name: mlx
tags:
- mlx
- moonbox
- background-removal
- isnet
- image-matting
---
# ORMBG MLX
MLX builds of [ORMBG](https://huggingface.co/schirrmacher/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`](https://huggingface.co/moonbox-ai/ormbg-MLX/blob/main/bf16/config.json) | 0.1 kB | `09fea8fff5a6` | alpha MAE <= 0.0026, mask IoU >= 0.987 |
| bf16 | [`bf16/mlx_model.py`](https://huggingface.co/moonbox-ai/ormbg-MLX/blob/main/bf16/mlx_model.py) | 8.2 kB | `8cf5c652dcbb` | alpha MAE <= 0.0026, mask IoU >= 0.987 |
| bf16 | [`bf16/model.safetensors`](https://huggingface.co/moonbox-ai/ormbg-MLX/blob/main/bf16/model.safetensors) | 88.0 MB | `39025b410188` | alpha MAE <= 0.0026, mask IoU >= 0.987 |
| fp16 | [`fp16/config.json`](https://huggingface.co/moonbox-ai/ormbg-MLX/blob/main/fp16/config.json) | 0.1 kB | `484d65d02eff` | alpha MAE <= 0.0004, mask IoU >= 0.998 |
| fp16 | [`fp16/mlx_model.py`](https://huggingface.co/moonbox-ai/ormbg-MLX/blob/main/fp16/mlx_model.py) | 8.2 kB | `8cf5c652dcbb` | alpha MAE <= 0.0004, mask IoU >= 0.998 |
| fp16 | [`fp16/model.safetensors`](https://huggingface.co/moonbox-ai/ormbg-MLX/blob/main/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):
```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, 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](https://huggingface.co/schirrmacher/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 <<<](https://huggingface.co/spaces/schirrmacher/ormbg)
Join our [Research Discord Group](https://discord.gg/YYZ3D66t)!
This model is a **fully open-source background remover** optimized for images with humans. It is based on [Highly Accurate Dichotomous Image Segmentation research](https://github.com/xuebinqin/DIS).
## Inference
```
python ormbg/inference.py
```
## Training
Install dependencies:
```
conda env create -f environment.yaml
conda activate ormbg
```
Replace dummy dataset with [training dataset](https://huggingface.co/datasets/schirrmacher/humans).
```
python3 ormbg/train_model.py
```
# Research
I started training the model with synthetic images of the [Human Segmentation Dataset](https://huggingface.co/datasets/schirrmacher/humans) crafted with [LayerDiffuse](https://github.com/layerdiffusion/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)](https://arxiv.org/pdf/2112.09290).
### 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](https://github.com/DLR-RM/BlenderProc)
- Removed training data created with [LayerDiffuse](https://github.com/layerdiffusion/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](https://drive.google.com/drive/folders/1Yy9clZ58xCiai1zYESQkEKZCkslSC8eg)
### 05/07/2024
- Added [P3M-10K](https://paperswithcode.com/dataset/p3m-10k) dataset for training and validation
- Added [AIM-500](https://paperswithcode.com/dataset/aim-500) dataset for training and validation
- Added [PPM-100](https://github.com/ZHKKKe/PPM) dataset for training and validation
- Applied [Grid Dropout](https://albumentations.ai/docs/api_reference/augmentations/dropout/grid_dropout/) to make the model smarter
|