MobileSAM icon indicating copy to clipboard operation
MobileSAM copied to clipboard

add little code for onnx export and quick demo

Open ChaofWang opened this issue 2 years ago • 1 comments

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>

ChaofWang avatar Jun 28 '23 13:06 ChaofWang

Checkout SAMExporter too: https://github.com/vietanhdev/samexporter.

vietanhdev avatar Jun 28 '23 17:06 vietanhdev