DOCSEARCHbyAlgolia
Submit Docs

SAMExporter

repository·main·Indexed 19 days ago

https://github.com/vietanhdev/samexporter

A tool to export Segment Anything models (SAM, SAM2, SAM2.1, SAM3, and MobileSAM) to ONNX format for dependency-free deployment. It provides utilities to export image encoders, mask decoders, and language encoders, with support for dynamic quantization, GELU approximation, and embedded image pre-processing. Includes an inference module to run segmentation using exported ONNX models with point, rectangle, or text prompts.

Tokens
10.9K
Snippets
29
Records
38
Agent score
65%

What's inside samexporter

  1. Run Inference with SAMExporter

    main

    Use samexporter.inference to run segmentation on an image using exported ONNX models.

    For SAM / SAM2 / SAM2.1

    Pass the encoder and decoder paths. For SAM2 variants, you must also pass --sam_variant sam2.

    python -m samexporter.inference \
    --encoder_model output_models/sam2_hiera_tiny.encoder.onnx \
    --decoder_model output_models/sam2_hiera_tiny.decoder.onnx \
    --image images/truck.jpg \
    --prompt images/truck_prompt.json \
    --sam_variant sam2 \
    --output output_images/sam2_truck.png \
    --show

    For SAM3 (Text-driven)

    SAM3 requires the --sam_variant sam3 flag and the --language_encoder_model path. You must also provide a --text_prompt.

    python -m samexporter.inference \
    --sam_variant sam3 \
    --encoder_model output_models/sam3/sam3_image_encoder.onnx \
    --decoder_model output_models/sam3/sam3_decoder.onnx \
    --language_encoder_model output_models/sam3/sam3_language_encoder.onnx \
    --image images/truck.jpg \
    --prompt images/truck_sam3.json \
    --text_prompt "truck" \
    --output output_images/truck_sam3.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/truck.jpg --prompt images/truck_prompt.json --output output_images/truck.png --show

  2. Install SAMExporter

    main

    SAMExporter requires Python 3.11+. You can install it via pip. Note that for Windows users, if you want to use the onnxsim model simplifier during export, you should install the [export] extra or enable Windows Long Path support.

    Standard installation:

    pip install torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cpu
pip install samexporter

    Windows installation with export support:

    pip install "samexporter[export]"

    To install from source:

    pip install torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cpu
git clone --recurse-submodules https://github.com/vietanhdev/samexporter
cd samexporter
pip install -e .
    pip install torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cpu
pip install samexporter

  3. Convert SAM3 to ONNX

    main

    SAM3 supports text-driven segmentation and exports into three separate ONNX models: an image encoder, a language (text) encoder, and a decoder.

    Option 1: Use Pre-exported Models

    You can use models from vietanhdev/segment-anything-3-onnx-models on HuggingFace, which are downloaded automatically.

    Option 2: Export from PyTorch

    To export manually, you must have the SAM3 source code available as a submodule and the osam package installed.

    # Clone SAM3 source
git submodule update --init sam3

# Install dependencies
pip install osam

# Export
python -m samexporter.export_sam3 \
    --output_dir output_models/sam3 \
    --opset 18
    python -m samexporter.export_sam3 --output_dir output_models/sam3 --opset 18

  4. Convert SAM2 and SAM2.1 to ONNX

    main

    SAM2 and SAM2.1 models are exported using a single command that produces both the encoder and decoder ONNX files.

    1. Prerequisites

    • Download SAM2 checkpoints (e.g., sam2_hiera_tiny.pt, sam2.1_hiera_tiny.pt) into original_models/.
    • Install the SAM2 PyTorch package:
      pip install git+https://github.com/facebookresearch/segment-anything-2.git

    2. Export

    Use samexporter.export_sam2 specifying the --model_type (e.g., sam2_hiera_tiny or sam2.1_hiera_tiny).

    # Example for SAM2 Tiny
python -m samexporter.export_sam2 \
    --checkpoint original_models/sam2_hiera_tiny.pt \
    --output_encoder output_models/sam2_hiera_tiny.encoder.onnx \
    --output_decoder output_models/sam2_hiera_tiny.decoder.onnx \
    --model_type sam2_hiera_tiny
    python -m samexporter.export_sam2 --checkpoint original_models/sam2_hiera_tiny.pt --output_encoder output_models/sam2_hiera_tiny.encoder.onnx --output_decoder output_models/sam2_hiera_tiny.decoder.onnx --model_type sam2_hiera_tiny

  5. Convert SAM and MobileSAM to ONNX

    main

    To convert original SAM or MobileSAM models to ONNX, follow these steps:

    1. Prepare Checkpoints

    Place your .pth or .pt checkpoints in an original_models/ directory. Supported files include sam_vit_b_01ec64.pth, sam_vit_l_0b3195.pth, sam_vit_h_4b8939.pth, and mobile_sam.pt.

    2. Export Encoder

    Use samexporter.export_encoder to export the image encoder. Use --model-type to specify the architecture (e.g., vit_h or vit_b).

    # Example for SAM ViT-H
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

    3. Export Decoder

    Use samexporter.export_decoder to export the mask decoder. Use --return-single-mask to return only one mask proposal, or omit it to return multiple.

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

  6. Format prompts for SAM inference

    main

    The get_input_points method (internal helper used by run_decoder) converts high-level prompt dictionaries into the numpy arrays required by the ONNX model. It supports two types of prompts:

    1. Points: Uses the data field for [x, y] coordinates and the label field for the point type.
    2. Rectangles: Converts the bounding box [x1, y1, x2, y2] into two distinct points (top-left and bottom-right) with specific labels (2 and 3) used by the SAM decoder.
    prompt = [
    {"type": "point", "data": [100, 150], "label": 1},
    {"type": "rectangle", "data": [50, 50, 150, 150]}
]

  7. How SAM3 language encoding works

    main

    SAM3 supports language-conditioned segmentation if a language_encoder_path is provided to the SegmentAnything3ONNX constructor.

    When a language encoder is present, the encode() method uses the SAM3LanguageEncoder to process the text_prompt. This produces three specific tensors used by the decoder:

    1. language_mask: A boolean attention mask [1, seq_len].
    2. language_features: Float features [seq_len, 1, 256].
    3. language_embeds: Float embeddings [seq_len, 1, 1024].

    Note on Tokenization: SAM3LanguageEncoder attempts to use osam._models.yoloworld.clip.tokenize for correct tokenization. If osam is not installed, it falls back to a minimal zero-filled tokenizer, which will result in near-random language features. For correct behavior, ensure osam is installed.

  8. Define prompts for SAM2 prediction

    main

    When using SegmentAnything2ONNX.predict_masks, prompts must be provided as a list of dictionaries. The supported types are:

    • point: Requires data as [x, y] and a label (typically 1 for foreground or 0 for background).
    • rectangle: Requires data as [x1, y1, x2, y2] (top-left and bottom-right coordinates). The class internally converts these into two points with labels 2 and 3.
    prompts = [
    {"type": "point", "data": [100, 150], "label": 1},
    {"type": "rectangle", "data": [50, 50, 200, 200]}
]

  9. Prompt JSON Format

    main

    Prompts are defined in JSON files as a list of mark objects. Supported types include point, rectangle, and text (for SAM3).

    [
  {"type": "point",     "data": [x, y],           "label": 1},
  {"type": "rectangle", "data": [x1, y1, x2, y2]},
  {"type": "text",      "data": "object description"}
]
    • label: 1: Foreground point.
    • label: 0: Background point.
    • type: "text": Specific to SAM3, though it is often more convenient to use the --text_prompt CLI flag during inference.
    [
  {"type": "point",     "data": [x, y],           "label": 1},
  {"type": "rectangle", "data": [x1, y1, x2, y2]},
  {"type": "text",      "data": "object description"}
]

  10. Prepare SAM3 for ONNX by replacing complex RoPE buffers

    main

    ONNX does not support complex-valued tensors. To export SAM3 models that use Rotary Positional Embeddings (RoPE), you must replace the complex freqs_cis buffers with separate real (cosine) and imaginary (sine) float buffers.

    Use the get_replace_freqs_cis function to recursively traverse the model and perform this replacement.

    def get_replace_freqs_cis(module: torch.nn.Module) -> None:
    """Replace complex freqs_cis buffers with separate real/imag float buffers.

ONNX does not support complex-valued tensors, so the complex RoPE
    (rotary positional embedding) buffer must be split into its real
    (cosine) and imaginary (sine) components before export.
    """
    if hasattr(module, "freqs_cis"):
        freqs_cos = module.freqs_cis.real.float()
        freqs_sin = module.freqs_cis.imag.float()
        module.register_buffer("freqs_cos", freqs_cos)
        module.register_buffer("freqs_sin", freqs_sin)
        del module.freqs_cis
    for child in module.children():
        get_replace_freqs_cis(child)

  11. Transform masks back to original size

    main

    The transform_masks method in SegmentAnything2ONNX uses affine transformations to map masks back to the original image dimensions, which is useful if the image was previously transformed (e.g., via rotation or scaling) before being passed to the encoder.

    • Parameters:
      • masks: The predicted masks array.
      • original_size: A (height, width) tuple of the target size.
      • transform_matrix: A 2x3 affine transformation matrix.
    • Returns: A np.ndarray of transformed masks.
    # Example usage within SegmentAnything2ONNX context
transformed_masks = model.transform_masks(
    masks=masks, 
    original_size=(720, 1280), 
    transform_matrix=my_affine_matrix
)

  12. Use run_export() to programmatically export SAM encoders

    main

    The run_export function allows for programmatic conversion of SAM encoders to ONNX. It handles model loading (via sam_model_registry or setup_model for mobile), wraps the model in ImageEncoderOnnxModel to handle pre-processing, and manages the ONNX export process.

    Note on ViT-H: For vit_h models, the function automatically handles large weight files by exporting to a temporary directory and then using convert_model_to_external_data to split the weights into a separate .bin file alongside the .onnx file.

    Parameters:

    • model_type (str): One of ['default', 'vit_h', 'vit_l', 'vit_b', 'mobile'].
    • checkpoint (str): Path to the model weights.
    • output (str): Target path for the .onnx file.
    • use_preprocess (bool): Whether to include normalization in the graph.
    • opset (int): ONNX opset version.
    • gelu_approximate (bool): Whether to use tanh for GELU operations (default: False).
    from samexporter.export_encoder import run_export

run_export(
    model_type="vit_h",
    checkpoint="sam_vit_h.pth",
    output="sam_vit_h.onnx",
    use_preprocess=True,
    opset=18,
    gelu_approximate=True
)