Hugging Face's logo

moonbox-dev
/
MobileSAM-MLX

Mask Generation
MLX
Safetensors
moonbox
mobilesam
segment-anything
sam
image-segmentation
Model card Files
xet
 Community
MobileSAM-MLX / README.md
jennyxue's picture
jennyxue
Model card and the original license
34d7499 verified
preview code
|
Raw History Blame Contribute Delete
17 kB
---
license: apache-2.0
pipeline_tag: mask-generation
library_name: mlx
tags:
- mlx
- moonbox
- mobilesam
- segment-anything
- sam
- image-segmentation
---
# MobileSAM v2 MLX
MLX builds of [MobileSAM v2](https://github.com/ChaoningZhang/MobileSAM), 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/MobileSAM-MLX/blob/main/bf16/config.json) | 0.3 kB | `18c4a51430fe` | mask IoU 0.996 mean / 0.972 worst |
| bf16 | [`bf16/mlx_model.py`](https://huggingface.co/moonbox-ai/MobileSAM-MLX/blob/main/bf16/mlx_model.py) | 20.7 kB | `57f29a58bfa3` | mask IoU 0.996 mean / 0.972 worst |
| bf16 | [`bf16/model.safetensors`](https://huggingface.co/moonbox-ai/MobileSAM-MLX/blob/main/bf16/model.safetensors) | 19.6 MB | `948290b7a2e3` | mask IoU 0.996 mean / 0.972 worst |
| fp16 | [`fp16/config.json`](https://huggingface.co/moonbox-ai/MobileSAM-MLX/blob/main/fp16/config.json) | 0.3 kB | `d6876abfe8f9` | mask IoU 0.999 mean / 0.988 worst |
| fp16 | [`fp16/mlx_model.py`](https://huggingface.co/moonbox-ai/MobileSAM-MLX/blob/main/fp16/mlx_model.py) | 20.7 kB | `57f29a58bfa3` | mask IoU 0.999 mean / 0.988 worst |
| fp16 | [`fp16/model.safetensors`](https://huggingface.co/moonbox-ai/MobileSAM-MLX/blob/main/fp16/model.safetensors) | 19.6 MB | `e8a52d0c013f` | mask IoU 0.999 mean / 0.988 worst |
| 8bit | [`8bit/config.json`](https://huggingface.co/moonbox-ai/MobileSAM-MLX/blob/main/8bit/config.json) | 0.4 kB | `a5c89d1aea72` | below threshold (experimental): mask IoU 0.989 mean / 0.899 worst |
| 8bit | [`8bit/mlx_model.py`](https://huggingface.co/moonbox-ai/MobileSAM-MLX/blob/main/8bit/mlx_model.py) | 20.7 kB | `57f29a58bfa3` | below threshold (experimental): mask IoU 0.989 mean / 0.899 worst |
| 8bit | [`8bit/model.safetensors`](https://huggingface.co/moonbox-ai/MobileSAM-MLX/blob/main/8bit/model.safetensors) | 12.7 MB | `d5daba02bdec` | below threshold (experimental): mask IoU 0.989 mean / 0.899 worst |
| 6bit | [`6bit/config.json`](https://huggingface.co/moonbox-ai/MobileSAM-MLX/blob/main/6bit/config.json) | 0.4 kB | `6cd61ce1c00d` | below threshold (experimental): mask IoU 0.989 mean / 0.914 worst |
| 6bit | [`6bit/mlx_model.py`](https://huggingface.co/moonbox-ai/MobileSAM-MLX/blob/main/6bit/mlx_model.py) | 20.7 kB | `57f29a58bfa3` | below threshold (experimental): mask IoU 0.989 mean / 0.914 worst |
| 6bit | [`6bit/model.safetensors`](https://huggingface.co/moonbox-ai/MobileSAM-MLX/blob/main/6bit/model.safetensors) | 10.8 MB | `6699f046a09d` | below threshold (experimental): mask IoU 0.989 mean / 0.914 worst |
| 4bit | [`4bit/config.json`](https://huggingface.co/moonbox-ai/MobileSAM-MLX/blob/main/4bit/config.json) | 0.4 kB | `4f556f197ff7` | below threshold (experimental): mask IoU 0.922 mean / 0.012 worst |
| 4bit | [`4bit/mlx_model.py`](https://huggingface.co/moonbox-ai/MobileSAM-MLX/blob/main/4bit/mlx_model.py) | 20.7 kB | `57f29a58bfa3` | below threshold (experimental): mask IoU 0.922 mean / 0.012 worst |
| 4bit | [`4bit/model.safetensors`](https://huggingface.co/moonbox-ai/MobileSAM-MLX/blob/main/4bit/model.safetensors) | 9.0 MB | `9beb710a859d` | below threshold (experimental): mask IoU 0.922 mean / 0.012 worst |
## How to run
MobileSAM (TinyViT-5M image encoder + SAM's prompt encoder and mask decoder, `weights/mobile_sam.pt`) ported to
MLX: the point / box segmentation half of MobileSAM v2. Each folder holds `model.safetensors`, `config.json` and
`mlx_model.py`, the port (only `mlx`, `numpy` and `pillow` needed), with SamPredictor's interface:
```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"))
emb = mlx_model.embed(model, rgb) # once per image
masks, scores, low_res, logits = mlx_model.predict(model, emb, points=[[320, 240]], labels=[1], multimask=True)
masks, scores, _, _ = mlx_model.predict(model, emb, box=[100, 80, 400, 380], multimask=False)
```
**Not included:** MobileSAM v2's object-aware prompt sampler (a YOLOv8 from ultralytics, AGPL-3.0, weights on Google
Drive only) and its prompt-guided decoder (shipped as a pickled PyTorch module). For "segment everything", feed
boxes from any detector to `predict`.
**How they were checked.** Against the original in PyTorch fp32 on 8 photos, one point prompt (3 masks) and one box
prompt each: IoU of each of the 32 masks with PyTorch's. Pass: IoU >= 0.95 on every mask.
| Precision | Size | Mask IoU mean / worst | Status |
|---|---|---|---|
| `fp16/` | 19.7 MB | 0.999 / 0.988 | passes |
| `bf16/` | 19.7 MB | 0.996 / 0.972 | passes |
| `8bit/` | 12.7 MB | 0.989 / 0.899 | **below threshold, experimental** |
| `6bit/` | 10.9 MB | 0.989 / 0.914 | **below threshold, experimental** |
| `4bit/` | 9.0 MB | 0.922 / 0.012 | **below threshold, experimental**: some prompts give a wrong mask |
Use `fp16/`: the model is 10 M parameters, so quantizing saves 7-11 MB and costs accuracy. About 65 ms per image
(encoder + two prompts) on an Apple M5.
## How these were made
From [MobileSAM v2](https://github.com/ChaoningZhang/MobileSAM) at revision `f706ad9c4e`: 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
<p float="center">
<img src="assets/logo2.png?raw=true" width="99.1%" />
</p>
# Faster Segment Anything (MobileSAM) and Everything (MobileSAMv2)
:pushpin: MobileSAMv2, available at [ResearchGate](https://www.researchgate.net/publication/376579294_MobileSAMv2_Faster_Segment_Anything_to_Everything) and [arXiv](https://arxiv.org/abs/2312.09579.pdf), replaces the grid-search prompt sampling in SAM with object-aware prompt sampling for faster **segment everything(SegEvery)**.
:pushpin: MobileSAM, available at [ResearchGate](https://www.researchgate.net/publication/371851844_Faster_Segment_Anything_Towards_Lightweight_SAM_for_Mobile_Applications) and [arXiv](https://arxiv.org/pdf/2306.14289.pdf), replaces the heavyweight image encoder in SAM with a lightweight image encoder for faster **segment anything(SegAny)**.
**Support for ONNX model export**. Feel free to test it on your devices and share your results with us.
**A demo of MobileSAM** running on **CPU** is open at [hugging face demo](https://huggingface.co/spaces/dhkim2810/MobileSAM). On our own Mac i5 CPU, it takes around 3s. On the hugging face demo, the interface and inferior CPUs make it slower but still works fine. Stayed tuned for a new version with more features! You can also run a demo of MobileSAM on [your local PC](https://github.com/ChaoningZhang/MobileSAM/tree/master/app).
:grapes: Media coverage and Projects that adapt from SAM to MobileSAM (Thank you all!)
* **2023/07/03**: [joliGEN](https://github.com/jolibrain/joliGEN) supports MobileSAM for faster and lightweight mask refinement for image inpainting with Diffusion and GAN.
* **2023/07/03**: [MobileSAM-in-the-Browser](https://github.com/akbartus/MobileSAM-in-the-Browser) shows a demo of running MobileSAM on the browser of your local PC or Mobile phone.
* **2023/07/02**: [Inpaint-Anything](https://github.com/qiaoyu1002/Inpaint-Anything) supports MobileSAM for faster and lightweight Inpaint Anything
* **2023/07/02**: [Personalize-SAM](https://github.com/qiaoyu1002/Personalize-SAM) supports MobileSAM for faster and lightweight Personalize Segment Anything with 1 Shot
* **2023/07/01**: [MobileSAM-in-the-Browser](https://github.com/akbartus/MobileSAM-in-the-Browser) makes an example implementation of MobileSAM in the browser.
* **2023/06/30**: [SegmentAnythingin3D](https://github.com/Jumpat/SegmentAnythingin3D) supports MobileSAM to segment anything in 3D efficiently.
* **2023/06/30**: MobileSAM has been featured by [AK](https://twitter.com/_akhaliq?lang=en) for the second time, see the link [AK's MobileSAM tweet](https://twitter.com/_akhaliq/status/1674410573075718145). Welcome to retweet.
* **2023/06/29**: [AnyLabeling](https://github.com/vietanhdev/anylabeling) supports MobileSAM for auto-labeling.
* **2023/06/29**: [SonarSAM](https://github.com/wangsssky/SonarSAM) supports MobileSAM for Image encoder full-finetuing.
* **2023/06/29**: [Stable Diffusion WebUIv](https://github.com/continue-revolution/sd-webui-segment-anything) supports MobileSAM.
* **2023/06/28**: [Grounding-SAM](https://github.com/IDEA-Research/Grounded-Segment-Anything) supports MobileSAM with [Grounded-MobileSAM](https://github.com/IDEA-Research/Grounded-Segment-Anything/tree/main/EfficientSAM).
* **2023/06/27**: MobileSAM has been featured by [AK](https://twitter.com/_akhaliq?lang=en), see the link [AK's MobileSAM tweet](https://twitter.com/_akhaliq/status/1673585099097636864). Welcome to retweet.
![MobileSAM](assets/model_diagram.jpg?raw=true)
:star: **How is MobileSAM trained?** MobileSAM is trained on a single GPU with 100k datasets (1% of the original images) for less than a day. The training code will be available soon.
:star: **How to Adapt from SAM to MobileSAM?** Since MobileSAM keeps exactly the same pipeline as the original SAM, we inherit pre-processing, post-processing, and all other interfaces from the original SAM. Therefore, by assuming everything is exactly the same except for a smaller image encoder, those who use the original SAM for their projects can **adapt to MobileSAM with almost zero effort**.
:star: **MobileSAM performs on par with the original SAM (at least visually)** and keeps exactly the same pipeline as the original SAM except for a change on the image encoder. Specifically, we replace the original heavyweight ViT-H encoder (632M) with a much smaller Tiny-ViT (5M). On a single GPU, MobileSAM runs around 12ms per image: 8ms on the image encoder and 4ms on the mask decoder.
* The comparison of ViT-based image encoder is summarzed as follows:
Image Encoder | Original SAM | MobileSAM
:-----------------------------------------:|:---------|:-----:
Parameters | 611M | 5M
Speed | 452ms | 8ms
* Original SAM and MobileSAM have exactly the same prompt-guided mask decoder:
Mask Decoder | Original SAM | MobileSAM
:-----------------------------------------:|:---------|:-----:
Parameters | 3.876M | 3.876M
Speed | 4ms | 4ms
* The comparison of the whole pipeline is summarized as follows:
Whole Pipeline (Enc+Dec) | Original SAM | MobileSAM
:-----------------------------------------:|:---------|:-----:
Parameters | 615M | 9.66M
Speed | 456ms | 12ms
:star: **Original SAM and MobileSAM with a point as the prompt.**
<p float="left">
<img src="assets/mask_point.jpg?raw=true" width="99.1%" />
</p>
:star: **Original SAM and MobileSAM with a box as the prompt.**
<p float="left">
<img src="assets/mask_box.jpg?raw=true" width="99.1%" />
</p>
:muscle: **Is MobileSAM faster and smaller than FastSAM? Yes!**
MobileSAM is around 7 times smaller and around 5 times faster than the concurrent FastSAM.
The comparison of the whole pipeline is summarzed as follows:
Whole Pipeline (Enc+Dec) | FastSAM | MobileSAM
:-----------------------------------------:|:---------|:-----:
Parameters | 68M | 9.66M
Speed | 64ms |12ms
:muscle: **Does MobileSAM aign better with the original SAM than FastSAM? Yes!**
FastSAM is suggested to work with multiple points, thus we compare the mIoU with two prompt points (with different pixel distances) and show the resutls as follows. Higher mIoU indicates higher alignment.
mIoU | FastSAM | MobileSAM
:-----------------------------------------:|:---------|:-----:
100 | 0.27 | 0.73
200 | 0.33 |0.71
300 | 0.37 |0.74
400 | 0.41 |0.73
500 | 0.41 |0.73
## Installation
The code requires `python>=3.8`, as well as `pytorch>=1.7` and `torchvision>=0.8`. Please follow the instructions [here](https://pytorch.org/get-started/locally/) to install both PyTorch and TorchVision dependencies. Installing both PyTorch and TorchVision with CUDA support is strongly recommended.
Install Mobile Segment Anything:
```
pip install git+https://github.com/ChaoningZhang/MobileSAM.git
```
or clone the repository locally and install with
```
git clone git@github.com:ChaoningZhang/MobileSAM.git
cd MobileSAM; pip install -e .
```
## Demo
Once installed MobileSAM, you can run demo on your local PC or check out our [HuggingFace Demo](https://huggingface.co/spaces/dhkim2810/MobileSAM).
It requires latest version of [gradio](https://gradio.app).
```
cd app
python app.py
```
## <a name="GettingStarted"></a>Getting Started
The MobileSAM can be loaded in the following ways:
```
from mobile_sam import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor
model_type = "vit_t"
sam_checkpoint = "./weights/mobile_sam.pt"
device = "cuda" if torch.cuda.is_available() else "cpu"
mobile_sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
mobile_sam.to(device=device)
mobile_sam.eval()
predictor = SamPredictor(mobile_sam)
predictor.set_image(<your_image>)
masks, _, _ = predictor.predict(<input_prompts>)
```
or generate masks for an entire image:
```
from mobile_sam import SamAutomaticMaskGenerator
mask_generator = SamAutomaticMaskGenerator(mobile_sam)
masks = mask_generator.generate(<your_image>)
```
## <a name="GettingStarted"></a>Getting Started (MobileSAMv2)
Download the model weights from the [checkpoints](https://drive.google.com/file/d/1dE-YAG-1mFCBmao2rHDp0n-PP4eH7SjE/view?usp=sharing).
After downloading the model weights, faster SegEvery with MobileSAMv2 can be simply used as follows:
```
cd MobileSAMv2
bash ./experiments/mobilesamv2.sh
```
## ONNX Export
**MobileSAM** now supports ONNX export. Export the model with
```
python scripts/export_onnx_model.py --checkpoint ./weights/mobile_sam.pt --model-type vit_t --output ./mobile_sam.onnx
```
Also check the [example notebook](https://github.com/ChaoningZhang/MobileSAM/blob/master/notebooks/onnx_model_example.ipynb) to follow detailed steps.
We recommend to use `onnx==1.12.0` and `onnxruntime==1.13.1` which is tested.
## BibTex of our MobileSAM
If you use MobileSAM in your research, please use the following BibTeX entry. :mega: Thank you!
```bibtex
@article{mobile_sam,
title={Faster Segment Anything: Towards Lightweight SAM for Mobile Applications},
author={Zhang, Chaoning and Han, Dongshen and Qiao, Yu and Kim, Jung Uk and Bae, Sung-Ho and Lee, Seungkyu and Hong, Choong Seon},
journal={arXiv preprint arXiv:2306.14289},
year={2023}
}
```
## Acknowledgement
This work was supported by Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government(MSIT) (No.RS-2022-00155911, Artificial Intelligence Convergence Innovation Human Resources Development (Kyung Hee University)),(No.2019-0-01287, Evolvable Deep Learning Model Generation Platform for Edge Computing)
<details>
<summary>
<a href="https://github.com/facebookresearch/segment-anything">SAM</a> (Segment Anything) [<b>bib</b>]
</summary>
```bibtex
@article{kirillov2023segany,
title={Segment Anything},
author={Kirillov, Alexander and Mintun, Eric and Ravi, Nikhila and Mao, Hanzi and Rolland, Chloe and Gustafson, Laura and Xiao, Tete and Whitehead, Spencer and Berg, Alexander C. and Lo, Wan-Yen and Doll{\'a}r, Piotr and Girshick, Ross},
journal={arXiv:2304.02643},
year={2023}
}
```
</details>
<details>
<summary>
<a href="https://github.com/microsoft/Cream/tree/main/TinyViT">TinyViT</a> (TinyViT: Fast Pretraining Distillation for Small Vision Transformers) [<b>bib</b>]
</summary>
```bibtex
@InProceedings{tiny_vit,
title={TinyViT: Fast Pretraining Distillation for Small Vision Transformers},
author={Wu, Kan and Zhang, Jinnian and Peng, Houwen and Liu, Mengchen and Xiao, Bin and Fu, Jianlong and Yuan, Lu},
booktitle={European conference on computer vision (ECCV)},
year={2022}
```
</details>