Views
No views yet
1import torch
2from sam2.sam2_image_predictor import SAM2ImagePredictor
3
4predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-small")
5
6with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
7 predictor.set_image(<your_image>)
8 masks, _, _ = predictor.predict(<input_prompts>)1import torch
2from sam2.sam2_video_predictor import SAM2VideoPredictor
3
4predictor = SAM2VideoPredictor.from_pretrained("facebook/sam2-hiera-small")
5
6with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
7 state = predictor.init_state(<your_video>)
8
9 # add new prompts and instantly get the output on the same frame
10 frame_idx, object_ids, masks = predictor.add_new_points_or_box(state, <your_prompts>):
11
12 # propagate the prompts to get masklets throughout the video
13 for frame_idx, object_ids, masks in predictor.propagate_in_video(state):
14 ...@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}
}