jax-image-tools/sam21-tiny-video-onnx

jax-image-tools/sam21-tiny-video-onnx

ONNX export of facebook/sam2.1-hiera-tiny's promptable video / memory pipeline, for in-browser propagation with onnxruntime-web on WebGPU.

The published promptable-SAM exports stop at encoder + decoder, which segments one frame but cannot propagate a prompt across frames — that needs the memory bank. This repo carries the graphs for it, including the object_pointer output that image-only exports drop (their checkpoints discard object_pointer_proj).

Files

Graphs

vision_encoder     pixel_values (1,3,1024,1024) -> FPN feats + position encodings
mask_decoder       feats + point prompts -> pred_masks, high_res_masks, iou_scores,
                                            object_pointer, object_score_logits
memory_attention   current feats + memory bank -> memory-conditioned feats
memory_encoder     feats + predicted mask -> memory tokens + position encoding

The memory bank is not a graph. Upstream keeps it as dicts keyed by integer frame index and every branch is on a Python int or bool, never a tensor value, so it needs no tracing and belongs in JS. constants.json carries the geometry.

memory_attention takes the spatial blocks and the object pointers as separate, dynamic inputs, so it is exact for every bank state propagation produces — no padding on early frames.

Fidelity

Gated end to end against fp32 PyTorch propagate_in_video_iterator, with every neural call routed through these graphs:

Worst per-frame IoU 1.000000 against fp32 PyTorch propagate_in_video_iterator, worst max|Δ| on mask logits 3.3e-02. Measured over an 8-frame clip at two motion rates — fast video-like drift, and slow drift approximating z-stack slice-to-slice change — with the vision encoder, mask decoder, memory attention and memory encoder all routed through these graphs and the bank assembly left to upstream.

For reference, the published square-zero-labs/sam2.1-tiny-video-onnx export reports worst per-frame IoU 0.9967; it traces a fixed 7-block bank and pads early frames by duplication, whereas memory_attention here takes the spatial blocks and the object pointers as separate dynamic inputs and needs no padding.

Provenance

Exported with `browser-onnx-tools` (export/export_sam_video_onnx.py), gated by export/validate_sam_video.py.

License

Apache-2.0, inherited from `facebook/sam2.1-hiera-tiny`.