#!/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"