| --- |
| 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 |
|
|
