S15.2: XFeat exports at a fixed shape and loads under tract

tools/export-xfeat.sh exports the convolutional network alone at 768×1024
grayscale, on the pattern of export-seg-model.sh: thirteen standard
operator types, no dynamic axes, the keypoint decoding left to Rust.
examples/onnx_probe loads it through the ort-over-tract backend the app
ships with nothing unsupported and runs it in ~300 ms on the desktop CPU.

The weights are Apache-2.0, read from the repository's LICENSE, with no
grant on the checkpoint — recorded in models/LICENCE.md before they land,
as FR-MRG-8 asks. The probe stays: the next model will need the same
check.
This commit is contained in:
2026-09-19 15:24:10 +02:00
parent 1ded5afbaa
commit e4b6b6c935
6 changed files with 218 additions and 7 deletions
+78
View File
@@ -0,0 +1,78 @@
#!/usr/bin/env bash
# Re-export the keypoint model that ships in models/keypoints/.
#
# The .onnx is committed (D14, FR-MRG-8), so this is not part of any build —
# it exists so the committed artefact is reproducible rather than a binary
# someone once produced and nobody can regenerate. Run it when bumping the
# model or changing its input size.
#
# ./tools/export-xfeat.sh # 768×1024 -> models/keypoints/xfeat-1024.onnx
# ./tools/export-xfeat.sh 576 768 # another fixed size
#
# Requires `uv`. Everything else is fetched into a throwaway venv, including
# a CPU-only torch — the export needs no GPU and the CUDA wheels are 2 GB.
#
# ## What is exported, and what is not
#
# Only `XFeatModel.forward`: the convolutions from a normalised grayscale
# image to three dense maps at 1/8 resolution — 64-channel descriptors,
# 65-channel keypoint logits (an 8×8 cell plus "no keypoint"), and a
# 1-channel reliability heatmap. Everything `detectAndCompute` does after that
# — softmax over the 65 logits, pixel-shuffle to full resolution, 5×5 NMS,
# top-k, bilinear sampling of descriptors, L2 normalisation — is decoded in
# Rust, as yolo26's heads are. Those steps are cheap, shape-dependent and
# exactly the kind of graph tract parses badly.
#
# The input is one channel rather than three: `forward` takes the mean over
# channels first, and feeding grayscale makes that a no-op rather than an
# exported ReduceMean over data the app would have had to build.
#
# ## Why a fixed shape
#
# Same reason as export-seg-model.sh: tract cannot parse a dynamic-shape
# graph. 768×1024 is the proxy size the merge aligns at (panorama.md §4); a
# different size is a different file.
set -euo pipefail
HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO="$(cd "${HERE}/.." && pwd)"
H="${1:-768}"
W="${2:-1024}"
OUT="${REPO}/models/keypoints"
NAME="xfeat-${W}"
# Not `mktemp -d` under /tmp: a tmpfs, and torch is a gigabyte.
WORK="$(mktemp -d -p "${TMPDIR:-/var/tmp}")"
trap 'rm -rf "${WORK}"' EXIT
echo "==> exporting XFeat at ${H}×${W} in ${WORK}"
cd "${WORK}"
git clone -q --depth 1 https://github.com/verlab/accelerated_features.git xfeat
uv venv -q --python 3.12 venv
VIRTUAL_ENV="${WORK}/venv" uv pip install -q --index-url https://download.pytorch.org/whl/cpu torch
VIRTUAL_ENV="${WORK}/venv" uv pip install -q onnx onnxslim
VIRTUAL_ENV="${WORK}/venv" "${WORK}/venv/bin/python" - "${WORK}/xfeat" "${WORK}/${NAME}.onnx" "${H}" "${W}" <<'PY'
import sys, torch, onnx, onnxslim
sys.path.insert(0, sys.argv[1])
from modules.model import XFeatModel
repo, out, H, W = sys.argv[1], sys.argv[2], int(sys.argv[3]), int(sys.argv[4])
net = XFeatModel().eval()
net.load_state_dict(torch.load(f"{repo}/weights/xfeat.pt", map_location="cpu"))
torch.onnx.export(net, torch.zeros(1, 1, H, W), out, opset_version=17, dynamo=False,
input_names=["image"], output_names=["feats", "keypoints", "heatmap"],
dynamic_axes=None, do_constant_folding=True)
m = onnxslim.slim(onnx.load(out))
onnx.checker.check_model(m)
onnx.save(m, out)
print("ops:", sorted({n.op_type for n in m.graph.node}))
for o in m.graph.output:
print("out", o.name, [d.dim_value for d in o.type.tensor_type.shape.dim])
PY
mkdir -p "${OUT}"
cp "${WORK}/${NAME}.onnx" "${OUT}/${NAME}.onnx"
echo "==> ${OUT}/${NAME}.onnx"