sanskar003/SAM3-Split-Encoder-Decoder

SAM3 — Split Encoder / Decoder (ONNX)

Text-prompted, open-vocabulary image segmentation. Give it an image and a text concept (e.g. "person", "red car") and it returns instance masks, boxes and scores for that concept.

This is the two-stage split of SAM3:

File	Role	Runs
sam3_encoder.onnx	Vision backbone: image → 3 feature tensors	once per image
sam3_decoder.onnx	Detection head: features + text → masks/boxes/scores	once per prompt

Because the backbone is separated out, you encode the image once and then decode many prompts cheaply against the cached features — instead of re-running the whole model for every prompt.

Precision: the ONNX graph stores weights in fp32 (with a few int8-quantized tensors) and has fp32 inputs/outputs (feed a normal float32 image, read float32 logits/boxes/masks). It is meant to run at fp16 on GPU — build/enable fp16 in your runtime (e.g. TensorRT). The latency and VRAM figures below are for that fp16 GPU deployment.

---

Inputs & outputs

Encoder — sam3_encoder.onnx

Tensor	Shape	dtype	Notes
image (in)	[1, 3, 924, 924]	float32	RGB, CHW, normalized (x/255 - 0.5) / 0.5
/Add_179_output_0 (out)	[1, 4356, 256]	float32	feature — pass straight to decoder
/fused_229/Conv_output_0 (out)	[1, 256, 264, 264]	float32	feature — pass straight to decoder
/fused_232/Conv_output_0 (out)	[1, 256, 132, 132]	float32	feature — pass straight to decoder

Decoder — sam3_decoder.onnx

Tensor	Shape	dtype	Notes
3 feature tensors (in)	(above)	float32	fed by name from the encoder outputs
tokenized_text (in)	[1, 32]	int64	CLIP-style input_ids, zero-padded to 32
pred_masks (out)	[1, 200, 264, 264]	float32	per-query mask logits
pred_logits (out)	[1, 200, 1]	float32	per-query class logit
pred_boxes (out)	[1, 200, 4]	float32	per-query box, normalized cxcywh
presence_logit_dec (out)	[1, 1]	float32	global concept-presence logit

The encoder's output tensor names are identical to the decoder's feature input names, so features pass through by name — no remapping needed.

---

How to use

Requires: onnxruntime-gpu (or onnxruntime), opencv-python, numpy, transformers. The tokenizer comes from facebook/sam3.

import cv2
import numpy as np
import onnxruntime as ort
from transformers import AutoTokenizer

INPUT_SIZE = 924
TEXT_LEN = 32
SCORE_THRESHOLD = 0.25

# ── load once ────────────────────────────────────────────────────────────────
providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
encoder = ort.InferenceSession("sam3_encoder.onnx", providers=providers)
decoder = ort.InferenceSession("sam3_decoder.onnx", providers=providers)

tokenizer = AutoTokenizer.from_pretrained("facebook/sam3")
tokenizer.pad_token_id = 0  # CLIP-style zero padding

enc_out_names = [o.name for o in encoder.get_outputs()]      # the 3 feature tensors
dec_in_names  = [i.name for i in decoder.get_inputs()]
feature_names = [n for n in dec_in_names if n != "tokenized_text"]


# ── preprocessing ─────────────────────────────────────────────────────────────
def preprocess_image(frame_rgb):
    """HxWx3 uint8 RGB -> [1,3,924,924] float32."""
    resized = cv2.resize(frame_rgb, (INPUT_SIZE, INPUT_SIZE), interpolation=cv2.INTER_LINEAR)
    img = resized.astype(np.float32) / 255.0
    img = (img - 0.5) / 0.5
    img = img.transpose(2, 0, 1)[np.newaxis]
    return np.ascontiguousarray(img)

def tokenize(prompt):
    ids = tokenizer(prompt, padding="max_length", max_length=TEXT_LEN,
                    truncation=True, return_tensors="np")["input_ids"]
    return ids.astype(np.int64)


# ── postprocessing ────────────────────────────────────────────────────────────
def sigmoid(x):
    return 1.0 / (1.0 + np.exp(-np.clip(x, -50, 50)))

def cxcywh_to_xyxy(b):
    cx, cy, w, h = b[..., 0], b[..., 1], b[..., 2], b[..., 3]
    return np.stack([cx - 0.5 * w, cy - 0.5 * h, cx + 0.5 * w, cy + 0.5 * h], axis=-1)

def postprocess(out, orig_hw):
    H, W = orig_hw
    scores = sigmoid(out["pred_logits"][0, :, 0]) * sigmoid(out["presence_logit_dec"][0])
    keep = scores > SCORE_THRESHOLD
    boxes = cxcywh_to_xyxy(out["pred_boxes"][0][keep]) * np.array([W, H, W, H], np.float32)
    dets = []
    for score, box, m in zip(scores[keep], boxes, out["pred_masks"][0][keep]):
        mask = sigmoid(cv2.resize(m, (W, H), interpolation=cv2.INTER_LINEAR)) > 0.5
        dets.append({"score": float(score), "box_xyxy": box, "mask": mask})
    return dets


# ── inference: encode ONCE, decode MANY prompts ───────────────────────────────
frame_bgr = cv2.imread("image.jpg")
frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
orig_hw = frame_rgb.shape[:2]

# Stage 1 — backbone (once)
feats = encoder.run(enc_out_names, {"image": preprocess_image(frame_rgb)})
feat_feed = dict(zip(enc_out_names, feats))

# Stage 2 — decode each prompt against the cached features
for prompt in ["person", "car", "dog", "traffic light"]:
    feed = {n: feat_feed[n] for n in feature_names}
    feed["tokenized_text"] = tokenize(prompt)
    outs = decoder.run(None, feed)
    out = dict(zip([o.name for o in decoder.get_outputs()], outs))
    dets = postprocess(out, orig_hw)
    print(f"{prompt}: {len(dets)} detection(s)")

Scoring recap: score = sigmoid(pred_logits) * sigmoid(presence_logit_dec), keep score > 0.25 (tune to taste). Boxes are cxcywh normalized → convert to xyxy and scale by [W, H, W, H]. Masks are logits at 264×264 → resize to source size, sigmoid(...) > 0.5.

---

Benchmarks

• NVIDIA RTX 4070 Ti — DeepStream (TensorRT, fp16): ~60 ms end-to-end for 1 image + 4 prompts (one encoder pass + four decoder passes).

• VRAM: ~1800 MB with both encoder and decoder sessions loaded.

Because the encoder runs once and decoders are cheap, added prompts cost far less than a full model pass each — the per-prompt cost is dominated by the small decoder, not the backbone.

To reproduce on your own hardware, run the encoder once per image and the decoder per prompt (as in the example above), and time each stage separately (encoder vs. decoder).

---

Notes

• Precision: fp32 weights (a few int8-quantized tensors) with fp32 I/O; intended to run at fp16 on GPU. Enable fp16 in your runtime (e.g. TensorRT) for the benchmarked latency/VRAM. Running the raw fp32 graph (e.g. plain ONNX Runtime) is correct but slower and uses more memory.
• ONNX opset 17 (IR version 8) — use a matching onnxruntime (≥ 1.14).
• Fixed input resolution 924×924 — the position embeddings are baked to this grid; resizing the input is not supported.
• Single execution context per session: don't run one session from multiple threads at once; serialize calls or use one session per thread.

License & attribution

Derived from facebook/sam3. Use of these weights is subject to the SAM3 license — review and accept it before use.