|
--- |
|
license: apache-2.0 |
|
pipeline_tag: mask-generation |
|
library_name: sam2 |
|
--- |
|
|
|
Repository for SAM 2: Segment Anything in Images and Videos, a foundation model towards solving promptable visual segmentation in images and videos from FAIR. See the [SAM 2 paper](https://arxiv.org/abs/2408.00714) for more information. |
|
|
|
The official code is publicly release in this [repo](https://github.com/facebookresearch/segment-anything-2/). |
|
|
|
## Usage |
|
|
|
For image prediction: |
|
|
|
```python |
|
import torch |
|
from sam2.sam2_image_predictor import SAM2ImagePredictor |
|
|
|
predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-small") |
|
|
|
with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16): |
|
predictor.set_image(<your_image>) |
|
masks, _, _ = predictor.predict(<input_prompts>) |
|
``` |
|
|
|
For video prediction: |
|
|
|
```python |
|
import torch |
|
from sam2.sam2_video_predictor import SAM2VideoPredictor |
|
|
|
predictor = SAM2VideoPredictor.from_pretrained("facebook/sam2-hiera-small") |
|
|
|
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_or_box(state, <your_prompts>): |
|
|
|
# propagate the prompts to get masklets throughout the video |
|
for frame_idx, object_ids, masks in predictor.propagate_in_video(state): |
|
... |
|
``` |
|
|
|
Refer to the [demo notebooks](https://github.com/facebookresearch/segment-anything-2/tree/main/notebooks) for details. |
|
|
|
### Citation |
|
|
|
To cite the paper, model, or software, please use the below: |
|
``` |
|
@article{ravi2024sam2, |
|
title={SAM 2: Segment Anything in Images and Videos}, |
|
author={Ravi, Nikhila and Gabeur, Valentin and Hu, Yuan-Ting and Hu, Ronghang and Ryali, Chaitanya and Ma, Tengyu and Khedr, Haitham and R{\"a}dle, Roman and Rolland, Chloe and Gustafson, Laura and Mintun, Eric and Pan, Junting and Alwala, Kalyan Vasudev and Carion, Nicolas and Wu, Chao-Yuan and Girshick, Ross and Doll{\'a}r, Piotr and Feichtenhofer, Christoph}, |
|
journal={arXiv preprint arXiv:2408.00714}, |
|
url={https://arxiv.org/abs/2408.00714}, |
|
year={2024} |
|
} |
|
``` |