czg1225 / SlimSAM

SlimSAM: 0.1% Data Makes Segment Anything Slim
Apache License 2.0
248 stars 14 forks source link

How to use SlimSAM to segment everything and provide bbox? #19

Open dnnxl opened 1 month ago

dnnxl commented 1 month ago

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.

czg1225 commented 1 month ago

Hi @dnnxl , You can find more details at https://github.com/czg1225/SlimSAM/blob/master/inference.py.

dnnxl commented 1 month ago

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'
czg1225 commented 1 month ago

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,
}