SlimSAM
SlimSAM copied to clipboard
How to use SlimSAM to segment everything and provide bbox?
Hi, I would like to know if there a demo about how to use SlimSAM to segment everything, also, provide the bounding boxes in coco format. Thanks.
Hi @dnnxl , You can find more details at https://github.com/czg1225/SlimSAM/blob/master/inference.py.
Thanks! I just checked, what is the right key for the local running weights? I'm trying to use it with SAM.
import torch
from segment_anything import sam_model_registry
model_type = 'vit_p77'
checkpoint = '/content/SlimSAM-77-uniform.pth'
SlimSAM_model = sam_model_registry[model_type](checkpoint=checkpoint).to("cuda")
SlimSAM_model.to("cuda")
SlimSAM_model.eval()
Do I need to use vit_h
, vit_l
, vit_b
instead? Or is there any additional step to add it in the model sam_model_registry?
KeyError Traceback (most recent call last)
[<ipython-input-6-ecf96a6bd20b>](https://localhost:8080/#) in <cell line: 6>()
4 model_type = 'vit_p77'
5 checkpoint = '/content/SlimSAM-77-uniform.pth'
----> 6 SlimSAM_model = sam_model_registry[model_type](checkpoint=checkpoint).to("cuda")
7 SlimSAM_model.to("cuda")
8 SlimSAM_model.eval()
KeyError: 'vit_p77'
Hi @dnnxl , Please refer to 70th line of https://github.com/czg1225/SlimSAM/blob/master/segment_anything/build_sam.py. We have add the key value 'vit_p50' and 'vit_p77' for our local pruned models.
sam_model_registry = {
"default": build_sam_vit_h,
"vit_h": build_sam_vit_h,
"vit_l": build_sam_vit_l,
"vit_b": build_sam_vit_b,
"vit_p50": build_sam_vit_p50,
"vit_p77": build_sam_vit_p77,
}