Skip to content

feat(vit): derive ViT inner_crop from the ONNX probe geometry#42

Merged
Garrett Bischof (gwbischof) merged 1 commit into
mainfrom
h4a/onnx-inner-crop
Jun 9, 2026
Merged

feat(vit): derive ViT inner_crop from the ONNX probe geometry#42
Garrett Bischof (gwbischof) merged 1 commit into
mainfrom
h4a/onnx-inner-crop

Conversation

@gwbischof

Copy link
Copy Markdown
Collaborator

H4a of the #37 decomposition — derive the ViT-patch inner_crop from the model's baked probe instead of always auto-deriving it. (H4b, true vit-only mode, follows separately.)

Changes

New holoptycho/engine_probe.py:

  • find_onnx_for_engine — map a compiled {name}_v{ver}.engine to its source onnx/{name}/{ver}/*.onnx (the model_manager cache convention).
  • inner_crop_from_onnx — load the ONNX, extract the complex probe (probe_real/probe_imag initializers, or two 2-D constant graph outputs), and hand it to ptychoml.inner_crop_from_probe (the inscribed-square geometry, added in feat(preprocess): add inner_crop_from_probe — ViT-patch crop from probe geometry ptychoml#12). Falls back to None when onnx is unavailable or the probe can't be extracted.

Wired in PtychoApp.compose(): inner_crop precedence is explicit config > ONNX-probe-derived > None (then SaveViTResult auto-derives min(H,W)//4).

Why a separate module

The helpers are pure (Path/re/onnx/ptychoml — no Holoscan/Tiled), and vit_inference doesn't use them internally (only compose() does). Putting them in engine_probe.py keeps them out of vit_inference's import-time Tiled setup so they're unit-testable.

Tests

tests/test_engine_probe.py (7) fully covers find_onnx_for_engine — the holoptycho-specific filename logic (match, first-sorted, multi-digit version + underscored name, bad name, missing dir, empty dir) — plus the graceful-None fallback. The inscribed-square geometry is tested in ptychoml; the ONNX extraction path runs in real model loads (onnx isn't in the CI test env, so inner_crop_from_onnx returns None there).

pixi run test   # 85 passed; the 2 vit_inference/ptycho_holo smoke failures are
                # pre-existing (module load needs a reachable TILED_BASE_URL).

Depends on

NSLS2/ptychoml#12 (merged); pixi.lock bumped.

Add holoptycho/engine_probe.py:
- find_onnx_for_engine: map a compiled {name}_v{ver}.engine to its source
  onnx/{name}/{ver}/*.onnx (the model_manager cache convention).
- inner_crop_from_onnx: load the ONNX, extract the complex probe
  (probe_real/probe_imag, or 2-D constant graph outputs), and hand the probe
  to ptychoml.inner_crop_from_probe (geometry now lives in ptychoml). Falls
  back to None when onnx is unavailable or the probe can't be extracted.

Wire it in PtychoApp.compose(): inner_crop precedence is explicit config >
ONNX-probe-derived > None (SaveViTResult auto-derives min(H,W)//4).

The helpers live in their own module (not vit_inference) so they're free of
the Holoscan/Tiled import-time setup and can be unit tested. find_onnx_for_engine
is fully covered in tests/test_engine_probe.py (filename convention, sorting,
bad name, missing/empty dir) plus the graceful-None fallback; the inscribed-
square geometry is tested in ptychoml, and the ONNX *extraction* path runs in
real model loads (onnx isn't in the CI test env).

First of the two H4 PRs (true vit-only mode follows).

Co-authored-by: Himanshu Goel <4122621+himanshugoel2797@users.noreply.github.com>
@gwbischof Garrett Bischof (gwbischof) merged commit b542069 into main Jun 9, 2026
5 checks passed
@gwbischof Garrett Bischof (gwbischof) deleted the h4a/onnx-inner-crop branch June 9, 2026 16:22
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant