Export Segment Anything Models to ONNX
MIT License
Exporting Segment Anything, MobileSAM, and Segment Anything 2 into ONNX format for easy deployment.
Supported models:
Requirements:
From PyPi:
pip install torch==2.4.0 torchvision --index-url https://download.pytorch.org/whl/cpu
pip install samexporter
From source:
pip install torch==2.4.0 torchvision --index-url https://download.pytorch.org/whl/cpu
git clone https://github.com/vietanhdev/samexporter
cd samexporter
pip install -e .
original_models
+ sam_vit_b_01ec64.pth
+ sam_vit_h_4b8939.pth
+ sam_vit_l_0b3195.pth
+ mobile_sam.pt
...
python -m samexporter.export_encoder --checkpoint original_models/sam_vit_h_4b8939.pth \
--output output_models/sam_vit_h_4b8939.encoder.onnx \
--model-type vit_h \
--quantize-out output_models/sam_vit_h_4b8939.encoder.quant.onnx \
--use-preprocess
python -m samexporter.export_decoder --checkpoint original_models/sam_vit_h_4b8939.pth \
--output output_models/sam_vit_h_4b8939.decoder.onnx \
--model-type vit_h \
--quantize-out output_models/sam_vit_h_4b8939.decoder.quant.onnx \
--return-single-mask
Remove --return-single-mask
if you want to return multiple masks.
python -m samexporter.inference \
--encoder_model output_models/sam_vit_h_4b8939.encoder.onnx \
--decoder_model output_models/sam_vit_h_4b8939.decoder.onnx \
--image images/truck.jpg \
--prompt images/truck_prompt.json \
--output output_images/truck.png \
--show
python -m samexporter.inference \
--encoder_model output_models/sam_vit_h_4b8939.encoder.onnx \
--decoder_model output_models/sam_vit_h_4b8939.decoder.onnx \
--image images/plants.png \
--prompt images/plants_prompt1.json \
--output output_images/plants_01.png \
--show
python -m samexporter.inference \
--encoder_model output_models/sam_vit_h_4b8939.encoder.onnx \
--decoder_model output_models/sam_vit_h_4b8939.decoder.onnx \
--image images/plants.png \
--prompt images/plants_prompt2.json \
--output output_images/plants_02.png \
--show
Short options:
bash convert_all_meta_sam.sh
bash convert_mobile_sam.sh
cd original_models
bash download_sam2.sh
The models will be downloaded to the original_models
folder:
original_models
+ sam2_hiera_tiny.pt
+ sam2_hiera_small.pt
+ sam2_hiera_base_plus.pt
+ sam2_hiera_large.pt
...
pip install git+https://github.com/facebookresearch/segment-anything-2.git
bash convert_all_meta_sam2.sh
python -m samexporter.inference \
--encoder_model output_models/sam2_hiera_tiny.encoder.onnx \
--decoder_model output_models/sam2_hiera_tiny.decoder.onnx \
--image images/plants.png \
--prompt images/truck_prompt_2.json \
--output output_images/plants_prompt_2_sam2.png \
--sam_variant sam2 \
--show
This package was originally developed for auto labeling feature in AnyLabeling project. However, you can use it for other purposes.
This project is licensed under the MIT License - see the LICENSE file for details.