edgetam-video-onnx

jax-image-tools/edgetam-video-onnx

ONNX export of facebook/EdgeTAM'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 4.1e-04, at both motion rates tested — 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.

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/EdgeTAM.