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:
Executable
+78
@@ -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"
|
||||
Reference in New Issue
Block a user