dtourolle
54290b9540
Twelve real frames from the fixture set now align in 4.5 s — 4.4 s of matching, 118 ms of bundle adjustment — where the first run took 51 s and left the first two frames out. The matcher computes each pair's similarity matrix once, across the cores, with a dot product written to vectorise; both nearest-neighbour directions read it. The frames that failed were portrait: fitted into the landscape input they used 512 of 1024 px, and their thin overlap did not survive at half resolution. The same weights are now exported at 768×1024 as well and the detector picks the shape by aspect. The example aligns from embedded previews and draws the set on a cylinder; on the fixture the sweep is 152° at a fitted 47.9 mm against the EXIF's 50, RMS 1.5 px, and the overlaps show no ghosting.
84 lines
3.6 KiB
Bash
Executable File
84 lines
3.6 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 # both shapes -> models/keypoints/xfeat-{1024,768}.onnx
|
||
#
|
||
# Two files from one set of weights: 1024 wide × 768 tall for landscape
|
||
# frames and 768 × 1024 for portrait, chosen by the frame's aspect at run
|
||
# time. One landscape file would fit a portrait frame at half the width and
|
||
# waste half the detector on padding, which is what the 6D fixture did.
|
||
#
|
||
# 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)"
|
||
|
||
OUT="${REPO}/models/keypoints"
|
||
|
||
# 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 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
|
||
|
||
mkdir -p "${OUT}"
|
||
for shape in "768 1024" "1024 768"; do
|
||
set -- $shape
|
||
H="$1"; W="$2"
|
||
NAME="xfeat-${W}"
|
||
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
|
||
cp "${WORK}/${NAME}.onnx" "${OUT}/${NAME}.onnx"
|
||
echo "==> ${OUT}/${NAME}.onnx"
|
||
done
|