Instructions to use Pablopigue/sam2-lite with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- sam2
How to use Pablopigue/sam2-lite with sam2:
# Use SAM2 with images import torch from sam2.sam2_image_predictor import SAM2ImagePredictor predictor = SAM2ImagePredictor.from_pretrained("Pablopigue/sam2-lite") with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16): predictor.set_image(<your_image>) masks, _, _ = predictor.predict(<input_prompts>)# Use SAM2 with videos import torch from sam2.sam2_video_predictor import SAM2VideoPredictor predictor = SAM2VideoPredictor.from_pretrained("Pablopigue/sam2-lite") with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16): state = predictor.init_state(<your_video>) # add new prompts and instantly get the output on the same frame frame_idx, object_ids, masks = predictor.add_new_points(state, <your_prompts>) # propagate the prompts to get masklets throughout the video for frame_idx, object_ids, masks in predictor.propagate_in_video(state): ... - Notebooks
- Google Colab
- Kaggle
File size: 6,511 Bytes
3d20082 f3387fa 3d20082 aeaa0d3 3d20082 2a25d2c aeaa0d3 3d20082 aeaa0d3 3d20082 | 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 | ---
license: cc-by-nc-4.0
library_name: pytorch
pipeline_tag: mask-generation
tags:
- video-object-segmentation
- sam2
- knowledge-distillation
- onnx
---
# sam2-lite
**SAM 2.1 video object segmentation 2.13Γ faster on CPU**: SAM 2.1 Hiera-tiny with its image encoder distilled into a MobileNetV4 backbone and its memory attention distilled to attend to 3 memory frames instead of 7. Click an object in the first frame and it is tracked through the video.
Code, training and evaluation: https://github.com/pablopigue/sam2-lite
## Results
DAVIS 2017 val (30 videos, semi-supervised, J&F), per-frame latency of the full video pipeline (median, fp32, 13th Gen Intel(R) Core(TM) i7-13620H).
| | SAM 2.1 Hiera-T (teacher) | **sam2-lite** @1024 | **sam2-lite-mobile** @576 |
|---|---|---|---|
| J&F | 89.1 | 84.4 | 77.1 |
| CPU 6 threads (ms/frame) | 1956 | 920 | β |
| CPU 2 threads (ms/frame) | 4949 | 2270 | 398 |
| Model size, fp32 (MB) | 149 | 74 | 74 |
sam2-lite-mobile is the same weights at a 576Γ576 input; it was measured in a different session from the other latency numbers (~10 % run-to-run variation). 2 CPU threads = a modest 2-core CPU.
### ONNX Runtime encoder
The image encoder is also provided as ONNX (one file per resolution). PyTorch β ONNX Runtime, all measured in one later session (hence e.g. 892 instead of 920 ms for the same PyTorch model: ~5-10 % varies between sessions, so only same-session numbers are compared):
| | Encoder (ms) | Full pipeline (ms) |
|---|---|---|
| 1024, CPU 6 threads | 224 β 142 (1.58Γ) | 892 β 830 (1.07Γ) |
| 1024, CPU 2 threads | 434 β 295 (1.47Γ) | 2236 β 2075 (1.08Γ) |
| 576, CPU 6 threads | 44 β 40 (1.11Γ) | 159 β 151 (1.06Γ) |
| 576, CPU 2 threads | 96 β 75 (1.28Γ) | 392 β 373 (1.05Γ) |
The encoder is faster but the pipeline gains little: the memory attention, which still runs in PyTorch, dominates each frame. Quality is unchanged (5 val videos at 576: J&F 82.35 with PyTorch, 82.35 with ONNX Runtime).
## Files
- `model.safetensors` + `config.yaml`: the whole tracker (student encoder, fine-tuned memory attention, SAM 2.1-tiny memory encoder and mask decoder) and what is needed to rebuild it.
- `student_encoder_r1024.onnx`, `student_encoder_r576.onnx`: the image encoder for ONNX Runtime (CPU).
- `LICENSE`, `LICENSE-APACHE-2.0`, `NOTICE`: see [License](#license).
## Usage
```bash
git clone https://github.com/pablopigue/sam2-lite && cd sam2-lite && uv sync
```
```python
from huggingface_hub import snapshot_download
from sam2lite.export.bundle import load_bundle
path = snapshot_download("Pablopigue/sam2-lite")
predictor = load_bundle(path) # sam2-lite @1024
mobile = load_bundle(
path, image_size=576, onnx=f"{path}/student_encoder_r576.onnx"
) # sam2-lite-mobile, ONNX encoder
# Then use it like SAM 2's video predictor: init_state, add_new_points_or_box, propagate_in_video.
```
## Training
- **Image encoder:** MobileNetV4-Conv-Medium (timm, ImageNet-1k pretrained) + SAM 2's FPN neck, trained to reproduce the teacher encoder's multi-scale features (MSE per FPN level, FitNets-style), on ~3.8k unlabeled frames from 54 DAVIS 2017 train videos. The rest of SAM 2.1 is reused unchanged.
- **Memory attention:** a copy of SAM 2.1's, fine-tuned to reproduce the original's output (7 memory frames) while attending to 3 (stride 4), on DAVIS 2017 train clips.
- DAVIS 2017 val was used only for the final evaluation.
## Limitations
- Accuracy drops by 4.7 J&F points at 1024; a few small or close-up objects explain most of the drop.
- sam2-lite-mobile misses its own quality target (J&F β₯ 80) while meeting its latency target (β€ 600 ms with 2 threads).
- Distilled on a small dataset (DAVIS 2017 train only); the experiments point at the data as the current bottleneck.
- Latency was measured on one laptop CPU; numbers on other hardware will differ.
## License
sam2-lite is an independent project; it is not affiliated with, endorsed by or sponsored by Meta. "SAM 2" refers to the original model by Meta FAIR, on which this work is based.
Released under [CC BY-NC 4.0](https://creativecommons.org/licenses/by-nc/4.0/) for **non-commercial research use only** (`LICENSE`); `NOTICE` details every part:
- It contains SAM 2.1 weights by Meta Platforms, Inc., licensed under the Apache License 2.0 (`LICENSE-APACHE-2.0`). **Modified:** the memory attention (fine-tuned to attend to 3 memory frames instead of 7) and the memory temporal encodings (3 of the 7 kept). The other SAM 2.1 weights (memory encoder, prompt encoder, mask decoder, object pointers) are unchanged and remain available under the Apache License 2.0.
- The image encoder starts from the timm weights [`mobilenetv4_conv_medium.e500_r224_in1k`](https://huggingface.co/timm/mobilenetv4_conv_medium.e500_r224_in1k) (Apache 2.0), **modified** by distillation. They were pretrained on ImageNet-1k, whose terms allow only non-commercial research and educational use.
- It was distilled on frames of [DAVIS 2017](https://davischallenge.org), licensed under CC BY-NC 4.0; about half of its sequences come from third-party sources (mostly YouTube) with their own terms. No DAVIS videos, frames or annotations are included.
## Citations
```bibtex
@article{ravi2024sam2,
title = {SAM 2: Segment Anything in Images and Videos},
author = {Ravi, Nikhila and others},
journal = {arXiv preprint arXiv:2408.00714},
year = {2024}
}
@article{pont2017davis,
title = {The 2017 DAVIS Challenge on Video Object Segmentation},
author = {Pont-Tuset, Jordi and Perazzi, Federico and Caelles, Sergi and Arbel{\'a}ez, Pablo and Sorkine-Hornung, Alexander and Van Gool, Luc},
journal = {arXiv preprint arXiv:1704.00675},
year = {2017}
}
@inproceedings{perazzi2016davis,
title = {A Benchmark Dataset and Evaluation Methodology for Video Object Segmentation},
author = {Perazzi, Federico and Pont-Tuset, Jordi and McWilliams, Brian and Van Gool, Luc and Gross, Markus and Sorkine-Hornung, Alexander},
booktitle = {IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
year = {2016}
}
@article{qin2024mobilenetv4,
title = {MobileNetV4: Universal Models for the Mobile Ecosystem},
author = {Qin, Danfeng and others},
journal = {arXiv preprint arXiv:2404.10518},
year = {2024}
}
```
MLflow runs behind every number: `b44efdd1`, `62e59f5e`, `ca473c99`, `a0b4ebcc`, `a384f235`, `b205e7be`, `af0d9753`, `7baca321`, `79a6f48c`, `005f6d98`, `7be93e7d`, `4e29d891`
|