MobileSAM
MobileSAM copied to clipboard
add little code for onnx export and quick demo
Hello, add these lines of code in build_sam.py to get the same code behavior as sam like that:
def build_sam_mobile(checkpoint=None):
from mobile_encoder.setup_mobile_sam import setup_model
mobile_sam = setup_model()
if checkpoint is not None:
mobile_sam.load_state_dict(
torch.load(checkpoint), strict=True)
return mobile_sam
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,
'mobile': build_sam_mobile
}
Then you can export onnx model and quick demo like sam:
python scripts/amg.py --checkpoint <path/to/checkpoint> --model-type mobile --input <image_or_folder> --output <path/to/output>
python scripts/export_onnx_model.py --checkpoint <path/to/checkpoint> --model-type mobile --output <path/to/output>
Checkout SAMExporter too: https://github.com/vietanhdev/samexporter.