feat(preprocess): delegate model branch to ptychoml + geometry/DC knobs

AGENTS.md

@@ -850,16 +850,40 @@ from `PtychoViTInferenceOp`, the chunking loop is misbehaving — check that
   either redundant or wrong. Worth confirming with the beamline team and
   potentially reducing pipeline complexity.
 
-* **`auto_center_dp` (config field, default `true`) — one-shot diffraction
+* **`auto_center_dp` (config field, default `false`) — one-shot diffraction
   centering via scipy segmentation.** `ImagePreprocessorOp` averages the
   first batch (typically 64 frames), masks hot pixels at detector
   saturation, thresholds at 5% of peak, runs `scipy.ndimage.label` to find
   connected components, takes the centroid of the largest one, and shifts
   every subsequent batch (and the intensity tap) so that centroid lands
-  at the canvas centre. Averaging over the first batch protects against
-  the odd empty/saturated first frame. If no object passes the threshold
-  (truly blank first batch), no shift is applied. Set to `false` if the
-  operator has already centered manually via `batch_x0`/`batch_y0`.
+  at the canvas centre. **Default off:** this shift is not part of
+  ptychoml's preprocessing pipeline and moves the beam away from the
+  position the model was trained on. Enable only when the detector ROI is
+  too far off-centre to fix with `batch_x0`/`batch_y0`.
+
+* **Diffraction geometry + normalization knobs (config fields).** The model
+  input branch of `ImagePreprocessorOp` delegates to
+  `ptychoml.preprocess_diffraction` (normalize → hot-pixel mask → sqrt →
+  D4 → fftshift); the intensity tap is oriented with `ptychoml.apply_d4`.
+  All are settable from the scan JSON:
+  * `tap_orient` (default `antitranspose`) — D4 element applied to the saved
+    intensity tap (`/dp`). The default reproduces the historical HXN
+    anti-diagonal flip; use `identity` to save raw detector frames.
+  * `dp_orient` (default `rot90_cw`) — D4 element on the model-input branch.
+    The default reproduces the prior hardcoded chain (verified equivalent);
+    orientation auto-detect will set this automatically once wired up.
+  * `fftshift_dp` (default unset → `None`) — DC convention for the model
+    input. `None` lets ptychoml auto-detect via `detect_dc_at_corner` and
+    shift only when the central beam sits at the corners. Override with
+    `true`/`false` only when a specific dataset misbehaves.
+  * `vit_normalization` (default `1e5`) — per-scan max intensity (hot pixels
+    excluded). In live mode the full DP stack isn't available, so this must
+    come from the scan JSON (precomputed offline, or via
+    `ptychoml.compute_intensity_normalization` on a representative subset).
+    Without the right value the amplitude scale drifts from training.
+  * `vit_scale` (default `1e4`) — model-input amplitude scale factor.
+  * `hot_pixel_count_threshold` (default unset → disabled) — photon-count
+    threshold for hot-pixel zeroing; matches `hxn_to_vit.py` when enabled.
 
 * **`mosaic_overshoot_factor` (config field, default 1.2) — canvas safety
   margin for the ViT mosaic.** Sized as `max(observed_range, commanded_range

holoptycho/preprocess.py

@@ -6,16 +6,36 @@
 import cupy as cp
 
 from ptychoml.preprocess import (
+    apply_d4,
     apply_intensity_floor,
     crop_to_roi,
     inpaint_bad_pixels,
+    preprocess_diffraction,
 )
 
 from holoscan.core import Operator, OperatorSpec, ConditionType, IOSpec
 from holoscan.schedulers import GreedyScheduler, MultiThreadScheduler, EventBasedScheduler
 from holoscan.logger import LogLevel, set_log_level
 from holoscan.decorator import create_op, Input
 
+
+def crop_flipped_roi(image, roi):
+    """Crop a horizontally-flipped Eiger2 frame, mirroring the column ROI.
+
+    ``roi[1]`` stores the column window in raw-frame (pre-flip) coordinates,
+    measured from the left. ``np.flip(image, 1)`` reverses the columns, so
+    the raw col start/stop must be mirrored to positions measured from the
+    right of the flipped frame, using the actual post-flip frame width
+    (only known at compute time, not when the ROI is set at compose time).
+    Rows (``roi[0]``) are unaffected by the horizontal flip.
+    """
+    image = np.flip(image, 1)
+    W = image.shape[1]
+    r0, r1 = int(roi[0, 0]), int(roi[0, 1])
+    c0_raw, c1_raw = int(roi[1, 0]), int(roi[1, 1])
+    return image[r0:r1, W - c1_raw : W - c0_raw]
+
+
 class ImageBatchOp(Operator):
     def __init__(self, *args, **kwargs):
         super().__init__(*args,**kwargs)
@@ -74,12 +94,11 @@ def compute(self, op_input, op_output, context):
         if self.roi is None:
             return
 
-        # For Eiger2 detector
+        # For Eiger2 detector: horizontal flip mirrors the column ROI.
         if self.flip_image:
-            image = np.flip(image,1)
-
-
-        image = crop_to_roi(image, self.roi)
+            image = crop_flipped_roi(image, self.roi)
+        else:
+            image = crop_to_roi(image, self.roi)
 
         # Remove Bad pixels (-1 to unsigned int)
         image[image==np.iinfo(image.dtype).max] = 0
@@ -109,15 +128,44 @@ def __init__(self, *args, **kwargs):
         # batch via scipy connected-component segmentation; same shift then
         # applied to every subsequent batch. ``None`` = not yet computed,
         # ``False`` = disabled by config. See ``_compute_centering_shift``.
-        self.auto_center = True
+        # Disabled by default: this shift is NOT part of ptychoml's preprocessing
+        # pipeline and moves the diffraction beam away from the position the model
+        # was trained on. Enable only if the detector ROI is uncalibrated and the
+        # beam is significantly off-centre.
+        self.auto_center = False
         self._center_shift: tuple[int, int] | None = None
-        # Extra transpose on the model-input branch, applied AFTER rot90 +
-        # fftshift. Historical (was commented out); re-enabled as a knob
-        # because some training runs expected the transposed orientation
-        # and feeding the wrong one yields garbage predictions. Affects
-        # only the model input — the intensity tap (saved dp) stays in
-        # detector orientation.
-        self.dp_transpose = True
+        # Geometry knobs for the two output branches.
+        #   tap_orient: D4 element applied to the intensity tap (saved /dp).
+        #               Default 'antitranspose' matches the historical HXN
+        #               anti-diagonal flip applied to bring data into the
+        #               beamline operator's view; the saved DP looks how the
+        #               operator expects.
+        #   dp_orient:  D4 element on the model-input branch. Default
+        #               'rot90_cw' reproduces the old chain
+        #               (antidiag flip ∘ rot90 ∘ transpose) for backwards
+        #               compatibility; orientation auto-detect will overwrite
+        #               this once it's wired up.
+        #   fftshift_dp: model-input DC-convention control. ``None``
+        #               (default) lets ptychoml auto-detect via
+        #               ``detect_dc_at_corner`` and shift only when the
+        #               central beam is at the corners. Override with
+        #               ``True``/``False`` from the scan config if a
+        #               specific dataset misbehaves; otherwise leave alone.
+        # All three are settable from the scan config; see ptycho_holo.py.
+        self.tap_orient = 'antitranspose'
+        self.dp_orient = 'rot90_cw'
+        self.fftshift_dp: bool | None = None
+        # Intensity normalization passed straight through to
+        # ptychoml.preprocess_diffraction so each DP gets scaled by the
+        # same constant the offline pipeline used. ``normalization`` is the
+        # per-scan max intensity (hot pixels excluded) — see
+        # ptychoml.compute_intensity_normalization. The default of 1e5 is a
+        # placeholder; in production the scan JSON overrides it per-scan.
+        self.normalization = 1.0e5
+        self.scale = 1.0e4
+        # Photon-count threshold for hot-pixel zeroing (None = disabled).
+        # 50000 matches hxn_to_vit.py's default.
+        self.hot_pixel_count_threshold = None
         super().__init__(*args, **kwargs)
 
         # Per-second compute() throughput counters. See note in EigerZmqRxOp.
@@ -240,32 +288,33 @@ def compute(self, op_input, op_output, context):
             dy, dx = self._center_shift
             processed_images = self._apply_shift(processed_images, dy, dx)
 
-        # Anti-diagonal flip per frame: HXN eiger data lands in the pipeline
-        # reflected across the anti-diagonal relative to the beamline
-        # operator's view. Applying the flip here means both the saved dp
-        # tap (next line) and the model-input branch downstream operate on
-        # data in the operator orientation, which is what the ViT model was
-        # trained on. Force a contiguous copy — the chain afterwards
-        # (rot90/fftshift/transpose/sqrt) is much slower on a strided view,
-        # which we saw back up ImagePreprocessorOp and starve point_proc.
-        processed_images = np.ascontiguousarray(
-            processed_images.transpose(0, 2, 1)[:, ::-1, ::-1]
-        )
-
-        # Tap detector-frame intensity before rot90/fftshift — ptycho-vit's
-        # training loader expects intensity in detector orientation. Contiguous
-        # copy so downstream rot90 (returns a view) doesn't alias the emitted
-        # buffer.
-        op_output.emit(np.ascontiguousarray(processed_images), "intensity")
-
-        # processed_images = processed_images[:, self.roi[0,0]:self.roi[0,1], self.roi[1,0]:self.roi[1,1]]
-        processed_images = np.rot90(processed_images, axes=(2,1))
-        processed_images = np.fft.fftshift(processed_images, axes=(1,2))
-        if self.dp_transpose:
-            processed_images = np.transpose(processed_images, [0, 2, 1])
+        # Tap branch: apply the configured D4 to put the saved intensity
+        # into whatever orientation the operator wants to see on the
+        # dashboard. Default 'antitranspose' reproduces the historical HXN
+        # anti-diagonal flip (transpose ∘ flip-both-axes); set
+        # tap_orient='identity' to save raw detector frames. Contiguous
+        # copy so the emitted buffer doesn't alias processed_images, which
+        # is mutated in place below.
+        tap = np.ascontiguousarray(apply_d4(processed_images, self.tap_orient))
+        op_output.emit(tap, "intensity")
+
+        # Model branch: delegate the entire normalize → mask → sqrt → D4 →
+        # fftshift sequence to ptychoml.preprocess_diffraction. Bad pixels
+        # are already inpainted above; the intensity floor (low-threshold)
+        # stays a holoptycho-side knob applied before the call because
+        # preprocess_diffraction doesn't expose it. fftshift=None (the
+        # default for this op) lets ptychoml auto-detect the central beam
+        # position and shift only when needed.
         if self.detmap_threshold > 0:
             apply_intensity_floor(processed_images, self.detmap_threshold)
-        diff_amp = np.sqrt(processed_images, dtype = np.float32 ,order='C')
+        diff_amp = preprocess_diffraction(
+            processed_images,
+            normalization=self.normalization,
+            scale=self.scale,
+            hot_pixel_count_threshold=self.hot_pixel_count_threshold,
+            dp_orient=self.dp_orient,
+            fftshift=self.fftshift_dp,
+        )
 
         op_output.emit(diff_amp, "diff_amp")
         op_output.emit(indices, "image_indices")

holoptycho/ptycho_holo.py

@@ -650,16 +650,41 @@ def compose(self):
         self.image_batch = ImageBatchOp(self, name="image_batch")
         self.image_proc = ImagePreprocessorOp(self, name="image_proc")
         # Auto-center the diffraction pattern via scipy segmentation on the
-        # average of the first batch (default on). Set `auto_center_dp=false`
-        # in the config to disable (e.g. if the operator has already set
-        # batch_x0/batch_y0 manually and doesn't want extra refinement).
-        self.image_proc.auto_center = bool(getattr(self.param, "auto_center_dp", True))
-        # Extra `np.transpose([0, 2, 1])` on the model-input branch after
-        # rot90 + fftshift (default on). Some model training runs expected
-        # the transposed orientation; flipping this knob is the fastest way
-        # to test "is the model getting garbage because the orientation is
-        # wrong?". Affects only the model input, not the saved dp.
-        self.image_proc.dp_transpose = bool(getattr(self.param, "dp_transpose", True))
+        # average of the first batch. Default off: this shift is not part of
+        # ptychoml's preprocessing pipeline and displaces the beam from the
+        # model's expected position. Enable via auto_center_dp=true in the
+        # scan JSON only when the detector ROI is too far off-centre to use
+        # batch_x0/batch_y0 correction instead.
+        self.image_proc.auto_center = bool(getattr(self.param, "auto_center_dp", False))
+        # Geometry + normalization for the two output branches. See
+        # ImagePreprocessorOp docstrings for what each does. Defaults reproduce
+        # the prior hardcoded HXN chain for the D4 transforms (antidiag tap +
+        # rot90_cw model branch); orientation auto-detect (when wired up) will
+        # set ``dp_orient`` automatically. ``fftshift_dp`` defaults to None so
+        # ptychoml's auto-detector picks the right DC convention per batch;
+        # override via the scan JSON only when a specific dataset misbehaves.
+        self.image_proc.tap_orient = str(getattr(self.param, "tap_orient", "antitranspose"))
+        self.image_proc.dp_orient = str(getattr(self.param, "dp_orient", "rot90_cw"))
+        _fftshift_dp = getattr(self.param, "fftshift_dp", None)
+        self.image_proc.fftshift_dp = (
+            bool(_fftshift_dp) if _fftshift_dp is not None else None
+        )
+        # Per-scan ViT normalization: the max intensity across all DPs in
+        # this scan, with hot pixels excluded. In live mode we don't have
+        # the full DP stack to compute it from, so it must come from the
+        # scan JSON (operator pre-computes it offline from a prior scan, or
+        # ptychoml.compute_intensity_normalization can produce it from a
+        # representative subset before the live run starts). Without it the
+        # amplitude scale drifts from what the model was trained against and
+        # predictions are systematically off-scale.
+        self.image_proc.normalization = float(getattr(self.param, "vit_normalization", 1.0e5))
+        self.image_proc.scale = float(getattr(self.param, "vit_scale", 1.0e4))
+        # Photon-count threshold for hot-pixel zeroing (None disables). Matches
+        # hxn_to_vit.py's default when enabled.
+        _hot_pix = getattr(self.param, "hot_pixel_count_threshold", None)
+        self.image_proc.hot_pixel_count_threshold = (
+            float(_hot_pix) if _hot_pix is not None else None
+        )
         self.image_send = ImageSendOp(self, name="image_send")
         self.point_proc = PointProcessorOp(
             self,
@@ -803,12 +828,15 @@ def compose(self):
         # Prefer a second GPU for PyCUDA/TRT when available, but fall back to
         # the recon GPU on single-GPU nodes instead of hard-failing.
         vit_gpu = self.param.gpus[1] if len(self.param.gpus) > 1 else self.param.gpus[0]
-        # Live mode: ImagePreprocessorOp applies fftshift — undo it for model
+        # DC-convention handling is auto-detected end-to-end by ptychoml
+        # (ImagePreprocessorOp's preprocess_diffraction centers the beam, and
+        # the session's own detect_dc_at_corner check verifies it), so no
+        # manual fftshift flag is needed here — leave fftshift at its default
+        # None.
         self.vit = PtychoViTInferenceOp(
             self,
             engine_path=self.engine_path,
             gpu=vit_gpu,
-            data_is_shifted=True,
             name="vit_inference",
         )
         # SaveViTResult publishes positions_um alongside each batch and

holoptycho/vit_inference.py

@@ -68,8 +68,13 @@ class PtychoViTInferenceOp(Operator):
         engine_path:       Path to .engine file (must match batch size B)
         gpu:               CUDA device ordinal (default 1; leave 0 for PtychoRecon)
         output_save_dir:   Directory for saving predictions (default /data/users/Holoscan)
-        data_is_shifted:   If True, input diff_amp has been fftshift'd and
-                           should be undone before inference.
+        fftshift:          DC-convention override for the session. Default
+                           ``None`` lets ptychoml auto-detect the central beam
+                           location per batch (via
+                           ``ptychoml.detect_dc_at_corner``) and shift only
+                           when needed — robust against an upstream op changing
+                           its own fftshift policy. Pass ``True``/``False`` to
+                           force a fixed convention (rarely needed).
     """
 
     def __init__(
@@ -79,15 +84,15 @@ def __init__(
         engine_path: str,
         gpu: int = 1,
         output_save_dir: str = "/data/users/Holoscan",
-        data_is_shifted: bool = False,
+        fftshift: bool | None = None,
         **kwargs,
     ):
         super().__init__(fragment, *args, **kwargs)
         self._logger = logging.getLogger("PtychoViTInferenceOp")
         self.engine_path = engine_path
         self.gpu = gpu
         self.output_save_dir = output_save_dir
-        self._data_is_shifted = data_is_shifted
+        self._fftshift = fftshift
 
         # Lazy-initialized on first compute()
         self._session = None
@@ -115,7 +120,7 @@ def _init_session(self):
         self._session = PtychoViTInference(
             engine_path=self.engine_path,
             gpu=self.gpu,
-            data_is_shifted=self._data_is_shifted,
+            fftshift=self._fftshift,
         )
         self._session._init_engine()
         self.engine_batch_size = int(self._session.expected_input_shape[0])

tests/test_preprocess_geometry.py

@@ -0,0 +1,48 @@
+"""Tests for holoptycho-specific preprocessing geometry.
+
+Covers the Eiger2 horizontal-flip ROI mirroring in ImageBatchOp — detector
+handling that lives in holoptycho (the generic crop is ptychoml.crop_to_roi).
+"""
+import numpy as np
+import pytest
+
+# preprocess.py imports holoscan + cupy at module level.
+pytest.importorskip("holoscan")
+pytest.importorskip("cupy")
+
+from ptychoml.preprocess import crop_to_roi
+from holoptycho.preprocess import crop_flipped_roi
+
+
+def _labelled_frame(h, w):
+    """Frame where pixel (r, c) == r*100 + c, so the crop's provenance is
+    readable from the values."""
+    r = np.arange(h)[:, None] * 100
+    c = np.arange(w)[None, :]
+    return (r + c).astype(np.int32)
+
+
+def test_crop_flipped_roi_selects_mirrored_window():
+    image = _labelled_frame(4, 10)
+    roi = np.array([[1, 3], [2, 5]])  # rows 1:3, raw cols 2:5
+
+    out = crop_flipped_roi(image, roi)
+
+    assert out.shape == (2, 3)
+    # rows 1,2 preserved; raw cols {2,3,4} selected but in reversed order
+    # because the frame was flipped before cropping.
+    assert out[0].tolist() == [104, 103, 102]
+    assert out[1].tolist() == [204, 203, 202]
+
+
+def test_crop_flipped_roi_equals_crop_then_flip():
+    """Flipping the whole frame then taking the mirrored window is the same
+    physical window as cropping first then reversing its columns."""
+    rng = np.random.default_rng(0)
+    image = rng.integers(0, 1000, size=(6, 12)).astype(np.int32)
+    roi = np.array([[2, 5], [3, 9]])
+
+    out = crop_flipped_roi(image, roi)
+    expected = np.flip(crop_to_roi(image, roi), 1)
+
+    np.testing.assert_array_equal(out, expected)