diff --git a/AGENTS.md b/AGENTS.md index 479fbd4..af88ad6 100644 --- a/AGENTS.md +++ b/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 diff --git a/holoptycho/preprocess.py b/holoptycho/preprocess.py index 26fcaac..3cf841e 100644 --- a/holoptycho/preprocess.py +++ b/holoptycho/preprocess.py @@ -6,9 +6,11 @@ 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 @@ -16,6 +18,24 @@ 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") diff --git a/holoptycho/ptycho_holo.py b/holoptycho/ptycho_holo.py index 943c1c9..b3983b8 100644 --- a/holoptycho/ptycho_holo.py +++ b/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 diff --git a/holoptycho/vit_inference.py b/holoptycho/vit_inference.py index 76f52f3..06640d2 100644 --- a/holoptycho/vit_inference.py +++ b/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,7 +84,7 @@ 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) @@ -87,7 +92,7 @@ def __init__( 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]) diff --git a/tests/test_preprocess_geometry.py b/tests/test_preprocess_geometry.py new file mode 100644 index 0000000..fcd7837 --- /dev/null +++ b/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) diff --git a/tests/test_vit_session_contract.py b/tests/test_vit_session_contract.py new file mode 100644 index 0000000..f3caa18 --- /dev/null +++ b/tests/test_vit_session_contract.py @@ -0,0 +1,33 @@ +"""Guard the ptychoml.PtychoViTInference call contract used by vit_inference. + +``PtychoViTInferenceOp._init_session`` constructs the session with a +``fftshift`` keyword. ptychoml renamed this parameter (it was +``data_is_shifted``), and bumping the ptychoml pin across that rename without +updating the call site silently broke the ViT op at the first batch — not at +import, so the smoke tests missed it. This pins the contract from holoptycho's +side so a future rename fails loudly here instead of on the beamline. + +Constructs ``PtychoViTInference`` directly (ptychoml only — no holoscan/TILED), +so it runs in the plain CI environment. +""" +import inspect + +from ptychoml import PtychoViTInference + + +def test_session_accepts_fftshift_kwarg(): + params = inspect.signature(PtychoViTInference.__init__).parameters + assert "fftshift" in params, ( + "ptychoml.PtychoViTInference no longer accepts 'fftshift'; " + "PtychoViTInferenceOp._init_session passes it" + ) + + +def test_session_construction_matches_op_call(): + # Mirror PtychoViTInferenceOp._init_session's call exactly. __init__ only + # stores config (the engine loads lazily in _init_engine), so this must + # not raise on the keyword even with a non-existent engine path. + session = PtychoViTInference( + engine_path="/tmp/nonexistent.engine", gpu=0, fftshift=None + ) + assert session is not None