dtourolle
e4b6b6c935
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.
79 lines
3.3 KiB
Bash
Executable File
79 lines
3.3 KiB
Bash
Executable File
#!/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"
|