Add ONNX export and an ONNX Runtime example - #5

Stacked on #11 (refactor for exporters). This PR adds the ONNX export of FlexGlue that produces the models of the v0.1 release used by COLMAP, an example of inference with ONNX Runtime, and a CI workflow that checks the export against PyTorch.

Contents

• onnx/export.py exports two models:
  • flexglue_descriptor.onnx: image (1, 3, 512, 512) → dense features (1, 256, 64, 64). It runs once per image; the caller squashes the image to 512×512 with a bilinear resize without antialiasing.
  • flexglue_matcher.onnx / flexglue_matcher_fp16.onnx: the sparse mode (match(..., sparse=True)), with float32 or float16 transformer layers.
    • Inputs: features0/1, kpts0/1 (pixels, top-left-corner origin, dynamic count ≥ 1), image_size0/1, and the runtime thresholds filter_threshold, max_distance_px and min_covisibility.
    • Unlike in PyTorch, max_distance_px = 0 rejects all matches; inf disables the check. This is documented in the README, the export contract and match_pair.py.
    • Outputs: matches0, mscores0 and covisibility_score. All matches are rejected in the graph if covisibility_score < min_covisibility.
• No duplicated model code. The graphs are traced from the PyTorch implementation itself, using the hooks added in Refactor the model code for exporters (no behavior change) #11:
  • DINOv3Descriptor.encode, FlexGlue.encode_pair, the transformer layers, pair_context, keypoint_correspondences and sparse_matches.
  • The exporter only configures the exported copy of the matcher: tensor thresholds, exact_nn, unchunked heads, use_blocked_qkv(), and modules.attention replaced by ONNX Runtime's fused com.microsoft::MultiHeadAttention (--no_fused_attention exports standard operators).
  • A thin wrapper module adds the graph inputs.
• onnx/check.py checks exported models against a PyTorch checkpoint: --weights (default: the release checkpoint), --descriptor and/or --matcher paths to ONNX models. It can check any model, e.g. those of the release, not only a fresh export.
  • The options of the matcher (float16 layers, fused operators) and the input resolution are read from the models.
  • It compares ONNX Runtime with PyTorch on the two asset pairs, for 512 and 2048 SIFT keypoints and min_covisibility 0 and 0.8:
    • the dense features;
    • the matches (flipped matches are counted) and the match and covisibility scores;
    • the predicted points of all keypoints that are confident enough to be matched, and the confidences of all keypoints. These are not outputs of the matcher: they come from a debug graph that check.py exports from the checkpoint with the same options as the matcher.
  • A few points can jump between two near-equal peaks of their correspondence scores, so the check counts large deviations instead of using a maximum. Tolerances are per device and precision (see TOLERANCES).
• CI (.github/workflows/onnx_export.yml):
  • Runs on changes to flexglue/, onnx/ or requirements.txt.
  • On a CPU runner with pinned torch / onnx / onnxscript / onnxruntime versions, it exports the descriptor and the float32 matcher (export.py), then checks them (check.py), in separate steps.
  • The float16 matcher needs a CUDA GPU and is checked manually.
• onnx/match_pair.py: a standalone example that matches two images with the released models using only numpy, OpenCV (SIFT keypoints) and ONNX Runtime, without PyTorch.
  • It downloads the models from the release and verifies their sha256.
  • It uses the fp16 matcher on CUDA and the fp32 matcher on CPU.
  • It can save a visualization of the matches.
• README: new "ONNX export" section.

Not exported: early stopping (#7), refinement, the dense and sampled modes, and CUDA graphs.

Validation

• CI run on this PR: passed in 6.5 min. Installation takes 0.5 min; the export and check take 5.5 min, including the first download of the weights.

• check.py results (ONNX Runtime 1.23, A100 / CPU):

  	Features	Flipped matches	Points, median error	Points off by >1 px	Confidence, p99 error
  CPU fp32	≤ 3e-5	0	≤ 3e-5 px	0%	≤ 2e-6
  CUDA fp32 (ONNX Runtime uses TF32)	≤ 3.5e-3	≤ 1	≤ 0.013 px	≤ 0.4%	≤ 1.4e-3
  CUDA fp16	≤ 3.5e-3	≤ 2	≤ 0.016 px	≤ 0.5%	≤ 2.9e-3

• Same results as the released models: models exported with this code give the same matches and covisibility scores as the v0.1 release models on the asset pairs (2k and 4k keypoints, fp32 and fp16), so the release models are unchanged. check.py also passes on the release models directly (descriptor, fp32 and fp16 matchers on CUDA).

• Same speed: fp16 on A100, 32.7 vs. 32.6 ms per pair at 2k and 46.5 vs. 46.1 ms at 8k (device I/O).

• onnx/match_pair.py gives the same matches and covisibility score as FlexGlue.match(..., sparse=True) on the same SIFT keypoints.

• Example output on the two asset pairs:

  $ python onnx/match_pair.py assets/berlin1.jpg assets/berlin2.jpg --output matches.png
  flexglue_matcher_fp16.onnx: 139 matches, covisibility score 0.993
  $ python onnx/match_pair.py assets/lamar_cab_1446.jpg assets/lamar_cab_1450.jpg --min_covisibility 0.8
  flexglue_matcher_fp16.onnx: 0 matches, covisibility score 0.064
  The pair is rejected: covisibility score < 0.8.
  

COLMAP integration: colmap/colmap#4908.