10sai-studio/bootstapir-onnx

BootsTAPIR, as ONNX

BootsTAPIR by Google DeepMind (the bootstapir_checkpoint_v2 PyTorch checkpoint), exported to ONNX for a window of 48 frames and one query point. 10sai's companion runs it on the CPU with ONNX Runtime to follow a point a person clicked through a clip.

Not ours. The model, its code and its checkpoint are Google DeepMind's, released under the Apache License 2.0; this repository only re-hosts an export of them so an installer can download one pinned file. See LICENSE.

How it was made

tapnet at 730cda1c730877cfedbe01bf87fb1cadb78a565d, torch 2.14.1 (CPU), Python 3.12, with two changes to the PyTorch model that do not change its result, both needed by the exporter:

• torch.randperm(num_queries) in TAPIR.estimate_trajectories replaced by torch.arange(num_queries): the order the queries are processed in is undone right after, and the exporter has no random permutation.
• The converting frame count is not supported check in torch/utils.py convert_grid_coordinates removed: it compares tensors in Python, which the exporter cannot trace.

model = tapir_model.TAPIR(pyramid_level=1)
model.load_state_dict(torch.load("bootstapir_checkpoint_v2.pt", map_location="cpu"))
torch.onnx.export(model, (video, query), "bootstapir-48x1.onnx", dynamo=True, opset_version=20, ...)

File	bootstapir-48x1.onnx, 222 387 291 bytes, weights inside the file
sha256	0ce28fb62d7d1ee29428807be7ed53b0afbc5a0d33286f305e122b6e19db34f4
Opset	20 (GridSample on 5-D input needs it)
Inputs	video: float32 [1, 48, 256, 256, 3], RGB in −1…1 · query: float32 [1, 1, 3] as frame, y, x in pixels of the 256 × 256 frames
Outputs	tracks [1, 1, 48, 2] x, y in pixels · occlusion [1, 1, 48] and expected_dist [1, 1, 48] logits; a point is visible where (1 − σ(occlusion)) · (1 − σ(expected_dist)) > 0.5

The frame count and the number of points are fixed: the model splits frames and queries into chunks in Python, which the exporter turns into constants. A longer clip is followed in overlapping windows.

Checked

Against the PyTorch model on the same inputs: tracks within 0.01 px. In ONNX Runtime 1.30 (Node 24), on a Ryzen 9 9900X with 8 threads: 0.8 s to load, 2.9 s per 48-frame window.