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:
@@ -0,0 +1,82 @@
|
||||
//! TRACES: S15 | FR-MRG-8
|
||||
//! Load an ONNX file through the application's own runtime and run it once.
|
||||
//!
|
||||
//! ```sh
|
||||
//! cargo run -p dr-segment --example onnx_probe --release -- model.onnx [1x1x768x1024]
|
||||
//! ```
|
||||
//!
|
||||
//! The F6 check, as a tool. tract's operator coverage is the thing that can
|
||||
//! sink a model choice — `segmentation.md` records a dynamic-shape export it
|
||||
//! could not parse at all — and the only way to know is to load the file
|
||||
//! under the backend the app ships and see. This does that for any model,
|
||||
//! before any Rust is written against its outputs: it prints the declared
|
||||
//! inputs and outputs, runs zeros through at the given shape, and times it.
|
||||
//!
|
||||
//! Written for S15.2 (XFeat), kept because the next model will need it too.
|
||||
|
||||
use std::time::Instant;
|
||||
|
||||
fn main() {
|
||||
let mut args = std::env::args().skip(1);
|
||||
let Some(path) = args.next() else {
|
||||
eprintln!("usage: onnx_probe <model.onnx> [NxCxHxW]");
|
||||
std::process::exit(2);
|
||||
};
|
||||
let shape: Vec<usize> = args
|
||||
.next()
|
||||
.map(|s| s.split('x').map(|d| d.parse().expect("dimension")).collect())
|
||||
.unwrap_or_else(|| vec![1, 1, 768, 1024]);
|
||||
|
||||
let bytes = std::fs::read(&path).expect("read model");
|
||||
println!("{path}: {} bytes", bytes.len());
|
||||
|
||||
let _ = ort::set_api(ort_tract::api());
|
||||
let t = Instant::now();
|
||||
let mut session = match ort::session::Session::builder().and_then(|mut b| b.commit_from_memory(&bytes)) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
println!("FAIL load: {e}");
|
||||
std::process::exit(1);
|
||||
}
|
||||
};
|
||||
println!("ok loaded in {:?}", t.elapsed());
|
||||
for i in session.inputs().iter() {
|
||||
println!(" input {} {:?}", i.name(), i.dtype());
|
||||
}
|
||||
for o in session.outputs().iter() {
|
||||
println!(" output {} {:?}", o.name(), o.dtype());
|
||||
}
|
||||
|
||||
let n: usize = shape.iter().product();
|
||||
|
||||
// Twice: the first run pays for tract's optimisation and plan, the second
|
||||
// is the number that matters. The tensor is built per run rather than
|
||||
// cloned — `Tensor::clone` under the tract backend panics.
|
||||
for pass in 1..=2 {
|
||||
let input =
|
||||
ndarray::Array::from_shape_vec(ndarray::IxDyn(&shape), vec![0.0f32; n]).expect("shape");
|
||||
let tensor = ort::value::Tensor::from_array(input).expect("tensor");
|
||||
let t = Instant::now();
|
||||
let outputs = match session.run(ort::inputs![tensor]) {
|
||||
Ok(o) => o,
|
||||
Err(e) => {
|
||||
println!("FAIL run: {e}");
|
||||
std::process::exit(1);
|
||||
}
|
||||
};
|
||||
println!("ok run {pass} in {:?}", t.elapsed());
|
||||
if pass == 2 {
|
||||
for i in 0..outputs.len() {
|
||||
match outputs[i].try_extract_tensor::<f32>() {
|
||||
Ok((shape, data)) => {
|
||||
let (lo, hi) = data
|
||||
.iter()
|
||||
.fold((f32::MAX, f32::MIN), |(lo, hi), &v| (lo.min(v), hi.max(v)));
|
||||
println!(" output {i}: shape {shape:?}, range {lo:.4}..{hi:.4}");
|
||||
}
|
||||
Err(e) => println!(" output {i}: not f32 ({e})"),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -139,6 +139,27 @@ first (D13's lesson, S15.2).
|
||||
| SuperPoint, SuperGlue, R2D2, SiLK, MASt3R | non-commercial | — | Out on licence |
|
||||
| LightGlue | Apache-2.0 | Transformer over a variable keypoint count | Not until mutual-nearest-neighbour matching fails on a real set |
|
||||
|
||||
**S15.2, 2026-09-19: XFeat loads under tract.** `tools/export-xfeat.sh`
|
||||
exports the network alone at 768×1024 — thirteen operator types, all
|
||||
standard: `Conv`, `InstanceNormalization`, `AveragePool`, `Resize`, `Slice`,
|
||||
`Transpose`, `Reshape`, `Concat`, `Add`, `Relu`, `Sigmoid`, `ReduceMean`,
|
||||
`Unsqueeze` — and
|
||||
[`examples/onnx_probe.rs`](../core/dr-segment/examples/onnx_probe.rs) loads
|
||||
the 2.8 MB file through the app's own `ort`-over-tract backend with nothing
|
||||
unsupported, in 28 ms, and runs it in **~300 ms on the reference desktop's
|
||||
CPU**. The weights ship as `models/keypoints/xfeat-1024.onnx`, recorded in
|
||||
`models/LICENCE.md`. Still to do: the tablet figure (S15.4), and a
|
||||
keypoint-level comparison against the PyTorch reference once the Rust decoder
|
||||
exists — the probe proves the graph runs, not that the numbers match.
|
||||
|
||||
The outputs are three maps at 1/8 resolution, 96×128 for the export size:
|
||||
64-channel descriptors, 65-channel keypoint logits (each 8×8 cell's position
|
||||
plus "none"), and a reliability heatmap. The Rust decoder is: softmax over the
|
||||
65, pixel-shuffle the first 64 to full resolution, 5×5 non-maximum
|
||||
suppression, top-k by reliability, bilinear sampling of the descriptor at
|
||||
each keypoint, L2 normalise. That is `detectAndCompute` in the reference,
|
||||
minus the network.
|
||||
|
||||
Without weights: AKAZE (BSD, `akaze` from rust-cv), which is adequate on
|
||||
well-textured overlaps and worse on sky, repeated structure and exposure
|
||||
drift — which is where a learned detector earns its place.
|
||||
|
||||
@@ -9,18 +9,18 @@ Denominators are parsed from [`requirements.md`](requirements.md) at run time, n
|
||||
|
||||
| Metric | Value |
|
||||
|---|---|
|
||||
| Source files scanned | 355 |
|
||||
| TRACES tags found | 1503 |
|
||||
| Source files scanned | 356 |
|
||||
| TRACES tags found | 1504 |
|
||||
| Requirements defined | 184 |
|
||||
| Requirements deferred (post-v1) | 24 |
|
||||
| Requirements covered | 143 |
|
||||
| **Coverage** | **77.7%** (143/184) |
|
||||
| Requirements covered | 144 |
|
||||
| **Coverage** | **78.3%** (144/184) |
|
||||
|
||||
### By type
|
||||
|
||||
| Type | Covered | Defined |
|
||||
|---|---|---|
|
||||
| FR | 105 | 127 |
|
||||
| FR | 106 | 127 |
|
||||
| NFR | 34 | 50 |
|
||||
| R | 4 | 7 |
|
||||
|
||||
@@ -99,6 +99,7 @@ _None._
|
||||
| FR-EXP-8 | [`core/dr-decode/src/lib.rs:330`](../core/dr-decode/src/lib.rs#L330), [`core/dr-decode/src/lib.rs:354`](../core/dr-decode/src/lib.rs#L354), [`core/dr-decode/src/lib.rs:368`](../core/dr-decode/src/lib.rs#L368), [`core/dr-decode/src/lib.rs:71`](../core/dr-decode/src/lib.rs#L71), [`core/dr-decode/src/lib.rs:79`](../core/dr-decode/src/lib.rs#L79), [`core/dr-decode/src/lib.rs:82`](../core/dr-decode/src/lib.rs#L82), [`core/dr-decode/src/locate.rs:1164`](../core/dr-decode/src/locate.rs#L1164), [`core/dr-decode/src/locate.rs:1223`](../core/dr-decode/src/locate.rs#L1223), [`core/dr-decode/src/locate.rs:316`](../core/dr-decode/src/locate.rs#L316), [`core/dr-decode/src/locate.rs:487`](../core/dr-decode/src/locate.rs#L487), [`core/dr-decode/src/locate.rs:571`](../core/dr-decode/src/locate.rs#L571), [`core/dr-decode/src/locate.rs:584`](../core/dr-decode/src/locate.rs#L584), [`core/dr-decode/src/locate.rs:667`](../core/dr-decode/src/locate.rs#L667), [`core/dr-export/examples/export.rs:99`](../core/dr-export/examples/export.rs#L99), [`core/dr-export/src/encode.rs:117`](../core/dr-export/src/encode.rs#L117), [`core/dr-export/src/encode.rs:161`](../core/dr-export/src/encode.rs#L161), [`core/dr-export/src/encode.rs:1`](../core/dr-export/src/encode.rs#L1), [`core/dr-export/src/encode.rs:206`](../core/dr-export/src/encode.rs#L206), [`core/dr-export/src/encode.rs:235`](../core/dr-export/src/encode.rs#L235), [`core/dr-export/src/encode.rs:311`](../core/dr-export/src/encode.rs#L311), [`core/dr-export/src/encode.rs:325`](../core/dr-export/src/encode.rs#L325), [`core/dr-export/src/encode.rs:408`](../core/dr-export/src/encode.rs#L408), [`core/dr-export/src/encode.rs:456`](../core/dr-export/src/encode.rs#L456), [`core/dr-export/src/encode.rs:70`](../core/dr-export/src/encode.rs#L70), [`core/dr-export/src/encode.rs:795`](../core/dr-export/src/encode.rs#L795), [`core/dr-export/src/encode.rs:809`](../core/dr-export/src/encode.rs#L809), [`core/dr-export/src/encode.rs:850`](../core/dr-export/src/encode.rs#L850), [`core/dr-export/src/encode.rs:898`](../core/dr-export/src/encode.rs#L898), [`core/dr-export/src/exif.rs:1`](../core/dr-export/src/exif.rs#L1), [`core/dr-export/src/lib.rs:136`](../core/dr-export/src/lib.rs#L136), [`core/dr-export/src/metadata.rs:1`](../core/dr-export/src/metadata.rs#L1), [`core/dr-export/src/metadata.rs:41`](../core/dr-export/src/metadata.rs#L41), [`core/dr-export/src/metadata.rs:74`](../core/dr-export/src/metadata.rs#L74), [`core/dr-types/src/lib.rs:655`](../core/dr-types/src/lib.rs#L655), [`core/dr-types/src/settings.rs:647`](../core/dr-types/src/settings.rs#L647), [`ui/dr-ui/src/develop.rs:4539`](../ui/dr-ui/src/develop.rs#L4539), [`ui/dr-ui/src/develop.rs:4651`](../ui/dr-ui/src/develop.rs#L4651), [`ui/dr-ui/src/develop.rs:729`](../ui/dr-ui/src/develop.rs#L729), [`ui/dr-ui/src/export.rs:467`](../ui/dr-ui/src/export.rs#L467), [`ui/dr-ui/src/export.rs:634`](../ui/dr-ui/src/export.rs#L634), [`ui/dr-ui/src/export.rs:680`](../ui/dr-ui/src/export.rs#L680), [`ui/dr-ui/src/export.rs:702`](../ui/dr-ui/src/export.rs#L702), [`ui/dr-ui/src/export.rs:826`](../ui/dr-ui/src/export.rs#L826), [`ui/dr-ui/src/export.rs:844`](../ui/dr-ui/src/export.rs#L844), [`ui/dr-ui/src/lib.rs:262`](../ui/dr-ui/src/lib.rs#L262), [`ui/dr-ui/src/lib.rs:285`](../ui/dr-ui/src/lib.rs#L285), [`ui/dr-ui/src/lib.rs:633`](../ui/dr-ui/src/lib.rs#L633), [`ui/dr-ui/src/settings_ui.rs:1`](../ui/dr-ui/src/settings_ui.rs#L1) |
|
||||
| FR-EXP-9 | [`core/dr-decode/src/lib.rs:510`](../core/dr-decode/src/lib.rs#L510), [`core/dr-export/src/lib.rs:128`](../core/dr-export/src/lib.rs#L128), [`core/dr-export/src/lib.rs:1`](../core/dr-export/src/lib.rs#L1), [`core/dr-gpu/src/adjust.rs:1106`](../core/dr-gpu/src/adjust.rs#L1106), [`ui/dr-ui/examples/face_native.rs:1`](../ui/dr-ui/examples/face_native.rs#L1), [`ui/dr-ui/src/develop.rs:4672`](../ui/dr-ui/src/develop.rs#L4672), [`ui/dr-ui/src/develop.rs:4710`](../ui/dr-ui/src/develop.rs#L4710), [`ui/dr-ui/src/develop.rs:6878`](../ui/dr-ui/src/develop.rs#L6878), [`ui/dr-ui/src/lib.rs:601`](../ui/dr-ui/src/lib.rs#L601), [`ui/dr-ui/src/library.rs:4052`](../ui/dr-ui/src/library.rs#L4052), [`ui/dr-ui/src/library.rs:4522`](../ui/dr-ui/src/library.rs#L4522), [`ui/dr-ui/src/library.rs:4558`](../ui/dr-ui/src/library.rs#L4558), [`ui/dr-ui/tests/export_ignores_the_viewport.rs:1`](../ui/dr-ui/tests/export_ignores_the_viewport.rs#L1) |
|
||||
| FR-MRG-3 | [`core/dr-decode/examples/linear_dng.rs:1`](../core/dr-decode/examples/linear_dng.rs#L1) |
|
||||
| FR-MRG-8 | [`core/dr-segment/examples/onnx_probe.rs:1`](../core/dr-segment/examples/onnx_probe.rs#L1) |
|
||||
| FR-NC-1 | [`core/dr-sync-nextcloud/src/auth.rs:132`](../core/dr-sync-nextcloud/src/auth.rs#L132), [`core/dr-sync-nextcloud/src/auth.rs:44`](../core/dr-sync-nextcloud/src/auth.rs#L44), [`core/dr-sync-nextcloud/src/provider.rs:1`](../core/dr-sync-nextcloud/src/provider.rs#L1), [`core/dr-sync/src/account.rs:361`](../core/dr-sync/src/account.rs#L361), [`ui/dr-ui/src/launch.rs:277`](../ui/dr-ui/src/launch.rs#L277), [`ui/dr-ui/src/launch.rs:61`](../ui/dr-ui/src/launch.rs#L61), [`ui/dr-ui/src/launch_ui.rs:417`](../ui/dr-ui/src/launch_ui.rs#L417) |
|
||||
| FR-NC-10 | [`core/dr-sync/src/account.rs:226`](../core/dr-sync/src/account.rs#L226), [`ui/dr-ui/src/export.rs:1`](../ui/dr-ui/src/export.rs#L1), [`ui/dr-ui/src/lib.rs:674`](../ui/dr-ui/src/lib.rs#L674), [`ui/dr-ui/src/library.rs:1170`](../ui/dr-ui/src/library.rs#L1170), [`ui/dr-ui/src/library.rs:2463`](../ui/dr-ui/src/library.rs#L2463), [`ui/dr-ui/src/library.rs:589`](../ui/dr-ui/src/library.rs#L589), [`ui/dr-ui/src/library.rs:916`](../ui/dr-ui/src/library.rs#L916), [`ui/dr-ui/src/library_ui.rs:2062`](../ui/dr-ui/src/library_ui.rs#L2062), [`ui/dr-ui/src/library_ui.rs:4275`](../ui/dr-ui/src/library_ui.rs#L4275), [`ui/dr-ui/src/library_ui.rs:660`](../ui/dr-ui/src/library_ui.rs#L660), [`ui/dr-ui/src/sidecar_cache.rs:1`](../ui/dr-ui/src/sidecar_cache.rs#L1) |
|
||||
| FR-NC-12 | [`core/dr-sync-folder/src/lib.rs:1`](../core/dr-sync-folder/src/lib.rs#L1), [`core/dr-sync-nextcloud/src/lib.rs:1022`](../core/dr-sync-nextcloud/src/lib.rs#L1022), [`core/dr-sync-nextcloud/src/lib.rs:40`](../core/dr-sync-nextcloud/src/lib.rs#L40), [`core/dr-sync-nextcloud/src/provider.rs:1`](../core/dr-sync-nextcloud/src/provider.rs#L1), [`core/dr-sync/src/account.rs:1`](../core/dr-sync/src/account.rs#L1), [`core/dr-sync/src/account.rs:87`](../core/dr-sync/src/account.rs#L87), [`core/dr-sync/src/lib.rs:218`](../core/dr-sync/src/lib.rs#L218), [`core/dr-sync/src/lib.rs:51`](../core/dr-sync/src/lib.rs#L51), [`core/dr-sync/src/provider.rs:106`](../core/dr-sync/src/provider.rs#L106), [`core/dr-sync/src/provider.rs:1`](../core/dr-sync/src/provider.rs#L1), [`core/dr-sync/src/provider.rs:53`](../core/dr-sync/src/provider.rs#L53), [`core/dr-sync/src/reachability.rs:1`](../core/dr-sync/src/reachability.rs#L1), [`ui/dr-ui/src/remote.rs:1`](../ui/dr-ui/src/remote.rs#L1) |
|
||||
@@ -209,7 +210,7 @@ Defined in `requirements.md` and marked `(post-v1)` on the defining line. Not in
|
||||
|
||||
## Not yet tagged
|
||||
|
||||
41 of 184 requirements have no implementation tag. Expected while the codebase is young; each should gain one as it is built.
|
||||
40 of 184 requirements have no implementation tag. Expected while the codebase is young; each should gain one as it is built.
|
||||
|
||||
<details><summary>Show untagged requirements</summary>
|
||||
|
||||
@@ -229,7 +230,6 @@ Defined in `requirements.md` and marked `(post-v1)` on the defining line. Not in
|
||||
- FR-MRG-5
|
||||
- FR-MRG-6
|
||||
- FR-MRG-7
|
||||
- FR-MRG-8
|
||||
- FR-MRG-9
|
||||
- FR-NC-11
|
||||
- FR-PLAT-AND-1
|
||||
|
||||
@@ -11,6 +11,8 @@ same script:
|
||||
|
||||
Both come from `https://huggingface.co/Ultralytics/YOLO26`. The face weights in
|
||||
`face/` are a separate matter with a separate grant — see `face/README.md`.
|
||||
The keypoint weights in `keypoints/` are a third, and the easiest — see the
|
||||
last section.
|
||||
|
||||
## The grant
|
||||
|
||||
@@ -69,3 +71,28 @@ Neither replaces the other. Keeping both is the deliberate choice.
|
||||
The loader treats each vocabulary as model metadata rather than compiled-in
|
||||
knowledge, which is what made adding the second model a file plus a descriptor
|
||||
rather than a code change — as this document predicted it would be.
|
||||
|
||||
## `keypoints/` — XFeat, Apache-2.0
|
||||
|
||||
| File | Source | Trained on | Used by |
|
||||
|---|---|---|---|
|
||||
| `keypoints/xfeat-1024.onnx` | `weights/xfeat.pt` from `https://github.com/verlab/accelerated_features` | MegaDepth + synthetic warps, by the authors | panorama alignment (FR-MRG-8) |
|
||||
|
||||
Exported by `tools/export-xfeat.sh` at a fixed 768×1024 grayscale input.
|
||||
Only the convolutional network is in the file; the keypoint decoding is Rust.
|
||||
|
||||
**The repository and its weights are Apache-2.0**, read on 2026-09-19 from the
|
||||
`LICENSE` at its root, with no separate grant on the checkpoint and no
|
||||
non-commercial clause anywhere in the tree. Apache-2.0 is GPLv3-compatible
|
||||
one way — code and weights under it may be combined into a GPLv3 work — so
|
||||
this is neither the InsightFace situation (D13, a use restriction that binds
|
||||
every user) nor the Ultralytics one (D14, where the combined work becomes
|
||||
AGPL). It is the licence position this document would have wanted for every
|
||||
model in it, and it was chosen over stronger detectors partly for that reason:
|
||||
SuperPoint and SuperGlue are non-commercial, R2D2 and SiLK are CC BY-NC.
|
||||
|
||||
The training data is the authors' concern, not a licence on the weights:
|
||||
XFeat trains on MegaDepth, which is itself a research dataset, but the weights
|
||||
are released under the repository's licence without a data-derived
|
||||
restriction — unlike the gaze models §7 of the requirements declined, where
|
||||
the dataset licence restricts models trained on it by name.
|
||||
|
||||
Binary file not shown.
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