--- 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)! ![](examples/image/image01_no_background.png) 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