Segment Anything (SAM)

The SA-1B dataset saves masks per image as a JSON file. When loaded as a dictionary in Python, it follows this structure:

• image: Contains image_id (int), width (int), height (int), and file_name (str).
• annotations: A list of annotation objects, each containing:
  • id: Annotation ID.
  • segmentation: Mask saved in COCO RLE format.
  • bbox: Box around the mask in [x, y, w, h] (XYWH) format.
  • area: Area in pixels.
  • predicted_iou: Model's prediction of mask quality.
  • stability_score: Measure of mask quality.
  • crop_box: The crop used for generation in [x, y, w, h] (XYWH) format.
  • point_coords: Input point coordinates [[x, y]] used to generate the mask.

{
    "image"                 : image_info,
    "annotations"           : [annotation],
}

image_info {
    "image_id"              : int,
    "width"                 : int,
    "height"                : int,
    "file_name"             : str,
}

annotation {
    "id"                    : int,
    "segmentation"          : dict, # COCO RLE format
    "bbox"                  : [x, y, w, h],
    "area"                  : int,
    "predicted_iou"         : float,
    "stability_score"       : float,
    "crop_box"              : [x, y, w, h],
    "point_coords"          : [[x, y]],
}

To use a custom image in the web demo, you must generate a corresponding .npy embedding file using the SAM Python API. This ensures the browser only needs to run the lightweight ONNX model rather than the full image encoder.

1. Initialize the SamPredictor with your chosen checkpoint.
2. Use predictor.set_image(image) to process your image.
3. Extract the embedding using predictor.get_image_embedding() and save it as a .npy file.
4. Place the image and the .npy file in the src/assets/data directory of the demo project.

import cv2
import numpy as np
from segment_anything import sam_model_registry, SamPredictor

# Initialize the predictor
checkpoint = "sam_vit_h_4b8939.pth"
model_type = "vit_h"
sam = sam_model_registry[model_type](checkpoint=checkpoint)
sam.to(device='cuda')
predictor = SamPredictor(sam)

# Set the new image and export the embedding
image = cv2.imread('src/assets/dogs.jpg')
predictor.set_image(image)
image_embedding = predictor.get_image_embedding().cpu().numpy()
np.save("dogs_embedding.npy", image_embedding)

The web demo requires a quantized ONNX model. You can generate this using the ONNX Model Example notebook.

Use quantize_dynamic to create the quantized version. Once generated, download sam_onnx_quantized_example.onnx and place it in the /model/ directory of the demo project.

Note: If you change the ONNX model (e.g., by using a different checkpoint), you must also re-export the image embeddings to match.

from onnxruntime.quantization import quantize_dynamic, QuantType

onnx_model_path = "sam_onnx_example.onnx"
onnx_model_quantized_path = "sam_onnx_quantized_example.onnx"

quantize_dynamic(
    model_input=onnx_model_path,
    model_output=onnx_model_quantized_path,
    optimize_model=True,
    per_channel=False,
    reduce_range=False,
    weight_type=QuantType.QUInt8,
)

Because the ONNX model expects a fixed input structure, you must manually format prompts:

• Points only: Concatenate a padding point with label -1 and coordinates (0.0, 0.0) to the end of your input points.
• Boxes: Encode the box as two points: the top-left corner (label 2) and the bottom-right corner (label 3). If a box is provided, you do not need the padding point.
• Transformation: Always use predictor.transform.apply_coords() to scale coordinates to the model's expected input size.

# Example: Point + Padding
# input_point = [[500, 375]], input_label = [1]
onnx_coord = np.concatenate([input_point, np.array([[0.0, 0.0]])], axis=0)[None, :, :]
onnx_label = np.concatenate([input_label, np.array([-1])], axis=0)[None, :].astype(np.float32)

# Example: Box (corners) + Point
# input_box = [x0, y0, x1, y1], input_point = [[px, py]], input_label = [0]
onnx_box_coords = input_box.reshape(2, 2)
onnx_box_labels = np.array([2, 3]) # 2=top-left, 3=bottom-right

# Combine point and box
onnx_coord = np.concatenate([input_point, onnx_box_coords], axis=0)[None, :, :]
onnx_label = np.concatenate([input_label, onnx_box_labels], axis=0)[None, :].astype(np.float32)

# Apply coordinate transformation
# onnx_coord = predictor.transform.apply_coords(onnx_coord, image.shape[:2]).astype(np.float32)

To use the exported ONNX model, you must first generate image embeddings using the heavy-weight SAM image encoder (typically on a GPU). Once you have the embeddings, you can run the lightweight ONNX model using an onnxruntime.InferenceSession.

Workflow:

1. Load the image and use SamPredictor.set_image(image).
2. Retrieve embeddings via predictor.get_image_embedding().
3. Prepare prompt inputs (point_coords, point_labels, etc.) following the specific ONNX input signature.
4. Run the session with ort_session.run().
5. Threshold the output logits at 0.0 to get a binary mask.

import onnxruntime
import numpy as np

# 1. Setup session
ort_session = onnxruntime.InferenceSession(onnx_model_path)

# 2. Get image embeddings (requires full SAM model)
# predictor = SamPredictor(sam)
# predictor.set_image(image)
# image_embedding = predictor.get_image_embedding().cpu().numpy()

# 3. Prepare inputs (example for points)
# ... (prepare onnx_coord, onnx_label, etc.) ...

# 4. Run inference
# masks, _, low_res_logits = ort_session.run(None, ort_inputs)

# 5. Threshold
# masks = masks > 0.0