diff --git a/AGENTS.md b/AGENTS.md index 68c75020..90a84961 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -144,6 +144,7 @@ Coverage reports are generated automatically (terminal, HTML in `htmlcov/`, and ### Naming Conventions - Functions/methods: `snake_case` +- Prefer imperative verb phrases for function names (for example, `get_source_dataarray`, `reconstruct_layer_dataarray`, `validate_inputs`), not noun phrases. - Classes: `PascalCase` - Constants: `UPPER_CASE` - Private functions/methods: leading underscore `_function_name` @@ -334,6 +335,17 @@ request, e.g. `([#123](https://github.com/confusius-tools/confusius/pull/123))`. - Keep tests fast by using small array sizes. - Use seeded random number generators for reproducibility. +### Napari Plugin Tests +- Follow napari's plugin testing guidelines: + . +- Main message: prefer **small unit tests** over full GUI/integration tests. Trust napari + to deliver callbacks/events correctly; test our plugin logic and observable widget/viewer + state directly instead of trying to simulate every napari interaction end-to-end. +- Use napari pytest fixtures such as `make_napari_viewer` / `make_napari_viewer_proxy` + rather than building custom viewer setup/teardown by hand. +- For ConfUSIus napari tests, prefer assertions on public/observable behavior + (widget state, layer state, metadata, outputs) over module-private helper return values. + ### Visual Regression Tests - Use `@pytest.mark.mpl_image_compare` for plot output tests. - Run `just generate-baselines` to regenerate baseline images after intentional plot changes. diff --git a/docs/gui/overview.md b/docs/gui/overview.md index 4e85dcac..01a5f93f 100644 --- a/docs/gui/overview.md +++ b/docs/gui/overview.md @@ -46,9 +46,10 @@ There are two ways to start the plugin: uvx -p 3.13 confusius ``` -The widget contains three collapsible panels—[Data I/O](plugin.md#data-io-panel), -[Signals](plugin.md#signals-panel), and [QC](plugin.md#qc-panel)—that can each be -expanded or collapsed independently. If you want a quick walkthrough inside napari +The widget contains five collapsible panels—[Data I/O](plugin.md#data-io-panel), +[Video](plugin.md#video-panel), [Signals](plugin.md#signals-panel), +[QC](plugin.md#qc-panel), and [Registration](plugin.md#registration-panel)—that can +be expanded or collapsed independently. If you want a quick walkthrough inside napari itself, click **Take a Tour** in the upper-right corner of the sidebar header. !!! tip "Running napari programmatically" diff --git a/docs/gui/plugin.md b/docs/gui/plugin.md index 5a87d19a..51b2db1c 100644 --- a/docs/gui/plugin.md +++ b/docs/gui/plugin.md @@ -14,6 +14,7 @@ introduction, click **Take a Tour** in the sidebar header. - [**Events**](#events-panel) — annotate periods of time (BIDS events) and shade them on the signal plot. - [**QC**](#qc-panel) — compute DVARS, carpet, CV, tSNR for a selected layer. +- [**Registration**](#registration-panel) — run between-scan or within-scan registration, inspect progress, and save/apply transforms. ## Data I/O Panel @@ -193,7 +194,8 @@ signals (from the current source mode) and imported signals: ## Events Panel -The Events Panel annotates *periods of time*—not individual frames—following the [BIDS +The Events Panel allows annotating *periods of time*—not individual frames—following the +[BIDS events](https://bids-specification.readthedocs.io/en/stable/modality-agnostic-files/events.html) convention (`onset`, `duration`, and an optional `trial_type`). Annotated events shade the [signal plot](#signals-panel) and are named in the time overlay while they are @@ -285,3 +287,108 @@ Select a layer from the **Layer** dropdown, check the metrics you want, and clic shadow zones behind the skull can appear bright. CV correctly highlights regions with high temporal variability. See the [Quality Control guide](../user-guide/quality-control.md#temporal-snr) for a full explanation. + +## Registration Panel + +The Registration Panel runs the ConfUSIus registration workflows directly from napari. +Use **Between scans** for registering different recordings, or **Within-scan** for +volume-wise motion correction within a single recording. The panel supports modifying +registration parameters, live preview, and saving/loading/applying computed transforms. + +### Between scans + +Use **Between scans** when you want to register one layer onto another, for example two +recordings from different animals or a functional recording onto an angiography. + +1. Select the **Moving layer** and **Fixed layer**. +2. Choose a transform model. +3. Optionally choose a **Scale** for to compress the intensity dynamics. +3. Optionally choose an **Initialization** transform if layers are very misaligned. +4. Click **Run registration**. + +Available transform models are: + +- `translation` for x/y/z-only shifts, +- `rigid` for translations and rotations, +- `affine` for translations, rotations, scaling, and shear, +- `bspline` for non-linear local deformations. + +For `bspline`, a staged workflow usually works best: first run `rigid` or `affine`, then +run `bspline` and select the previous transform in **Initialization**. This lets the +B-spline model refine a good global alignment instead of trying to solve both +large-scale and local deformation at once. + +#### Main parameters + +| Parameter | What it does | When it is useful | +|---|---|---| +| **Transform** | Chooses the motion model being optimized. | Start with `translation` or `rigid` for simple alignment; use `affine` for global scale/shear differences; use `bspline` only after a good global initialization. | +| **Mesh size** | Sets the B-spline control-grid density. | Increase it only when `bspline` needs to capture finer local mismatches; too fine a grid can lead to unrealistic warping. | +| **Metric** | Chooses the similarity criterion (`correlation` or `mattes_mi`). | `correlation` is a good default for power Doppler data; `mattes_mi` is more robust when intensity distributions differ. | +| **Scale** | Applies optional intensity scaling before registration. | Useful for power Doppler data where large vessels are typically overbright compared to finer structures. | +| **Initialization** | Sets the starting transform before optimization. | Use `center_geometry` or `center_moments` for coarse setup; reuse a saved/manual affine transform when you already have a good approximate alignment. | +| **Learning rate** | Sets the optimizer step size. | Lower values are safer but slower; higher values can converge faster but may create instabilities. | +| **Iterations** | Maximum number of optimizer steps. | Increase it when alignment is still improving near the end of a run. | + +#### Advanced parameters + +| Parameter | What it does | When it is useful | +|---|---|---| +| **Histogram bins** | Number of bins used by `mattes_mi`. | Tune only when using mutual information; more bins can capture finer intensity structure but may be noisier. | +| **Convergence minimum value** | Minimum optimizer improvement required to keep iterating. | Lower it when you want stricter convergence. | +| **Convergence window size** | Number of recent iterations used to test convergence. | Increase it to make convergence detection less sensitive to noise. | +| **Multi-resolution** | Runs registration from coarse to fine scales. | Usually helpful for difficult B-spline alignments or large initial offsets. | +| **Shrink factors** | Downsampling factors for each resolution level. | Use larger coarse levels when the initial mismatch is large. | +| **Smoothing sigmas** | Gaussian smoothing at each resolution level. | Helps emphasize global structure before fine alignment. | +| **Resample interp.** | Interpolation used for the registered output and previews. | `linear` is the usual default; `bspline` can give smoother resampled images. | +| **Fill value** | Value written outside the moving field of view after resampling. | Useful for controlling the appearance of padded background. | + +The animation below uses between-session angiography volumes from the same animal across +different days. ConfUSIus keeps the original layers untouched, adds dedicated preview +layers for inspection, and stores the final registered result as a new layer when the +run completes. + +![ConfUSIus Registration panel — rigid between-session angiography run](../images/gui/plugin-registration.gif) + +### Within-scan + +Use **Within-scan** for motion correction inside a single time series. + +1. Switch **Mode** to **Within-scan**. +2. Select the time-series **Moving layer**. +3. Choose the **Reference volume** index used as the registration target. +4. Pick a transform model (`translation`, `rigid`, or `affine`). +5. Click **Run registration**. + +#### Main parameters + +| Parameter | What it does | When it is useful | +|---|---|---| +| **Reference volume** | Chooses the volume index used as the motion-correction target. | Pick a representative, sharp frame with little motion. | +| **Transform** | Chooses the volume-wise motion model. | `rigid` is the safest starting point; `affine` is available when motion is more complex. | +| **Metric** | Chooses the volume-to-reference similarity criterion. | `correlation` is usually a good default for within-recording motion correction. | +| **Scale** | Applies optional preprocessing before registration. | Useful when an intensity transform makes anatomy more stable across time for the optimizer. | +| **Initialization** | Sets the initial volume-wise centering transform. | Most runs can use no initialization. | +| **Learning rate** | Sets the optimizer step size for each frame. | Reduce it if updates look unstable; increase it if frames are already close and convergence is too slow. | +| **Iterations** | Maximum optimizer steps per frame. | Increase it for harder motion or more flexible transforms. | + +#### Advanced parameters + +| Parameter | What it does | When it is useful | +|---|---|---| +| **Histogram bins** | Number of bins used by `mattes_mi`. | Only relevant when using mutual information. | +| **Convergence minimum value** | Minimum optimizer improvement required to keep iterating. | Lower it when you want stricter volume-wise convergence. | +| **Convergence window size** | Number of recent iterations used to test convergence. | Increase it when convergence decisions look too jittery. | +| **Multi-resolution** | Runs each frame registration from coarse to fine scales. | Helpful when motion is large or frames are noisy. | +| **Shrink factors** | Downsampling factors for each resolution level. | Useful for coarse-to-fine motion correction. | +| **Smoothing sigmas** | Gaussian smoothing at each resolution level. | Helps stabilize coarse registration before fine refinement. | +| **Resample interp.** | Interpolation used for the motion-corrected output. | Controls output smoothness. | +| **Fill value** | Value used outside the field of view after resampling. | Mostly useful for controlling output background appearance. | +| **Parallel jobs** | Number of workers used for volume-wise registration. | Increase it to speed up long runs; reduce it if your machine is already busy. `-1` uses all available CPUs | +| **Keep full traces** | Stores full volume-wise optimizer diagnostics. | Enable it only when you want detailed debugging or later inspection. | + +This workflow returns a `Motion corrected` layer and updates the progress bar as frames +finish. The animation below uses a short open-field recording chunk and shows the result +filling in progressively. + +![ConfUSIus Registration panel — within-scan motion correction](../images/gui/plugin-registration-volumewise.gif) diff --git a/docs/images/gui/generate.py b/docs/images/gui/generate.py index c78e3fc1..1efca47e 100644 --- a/docs/images/gui/generate.py +++ b/docs/images/gui/generate.py @@ -26,6 +26,8 @@ - `plugin-events-create.gif` — GIF of creating events with the Start/End workflow. - `plugin-qc.png` — QC panel with DVARS, carpet, and CV computed. - `plugin-video.gif` — Video panel with video synced to the fUSI acquisition. +- `plugin-registration.gif` — Registration panel during a rigid between-session angiography run. +- `plugin-registration-volumewise.gif` — Registration panel during within-scan motion correction. Notes ----- @@ -56,6 +58,14 @@ _TASK = "spontaneous" _ACQ_SLICE = "slice04" +_REGISTRATION_SUBJECT = "CR022" +_REGISTRATION_FIXED_SESSION = "20201007" +_REGISTRATION_MOVING_SESSION = "20201011" + +_VOLUMEWISE_SUBJECT = "rat75" +_VOLUMEWISE_SESSION = "20220523" +_VOLUMEWISE_ACQ_SLICE = "slice32" + _SLICE_INDEX = int(_ACQ_SLICE.replace("slice", "")) _ROI_STRUCTURE_ID = 1089 @@ -233,8 +243,8 @@ def _best_matching_z_coordinate(reference_2d, volume_3d) -> float: _section("Load Data") console.print("Fetching Nunez-Elizalde 2022 dataset") bids_root = fetch_nunez_elizalde_2022( - subjects=[_SUBJECT], - sessions=[_SESSION], + subjects=[_SUBJECT, _REGISTRATION_SUBJECT], + sessions=[_SESSION, _REGISTRATION_FIXED_SESSION, _REGISTRATION_MOVING_SESSION], tasks=[_TASK], acqs=[_ACQ_SLICE], ) @@ -323,12 +333,12 @@ def _best_matching_z_coordinate(reference_2d, volume_3d) -> float: _VIDEO_SESSION = "20220525" _VIDEO_ACQ_SLICE = "slice37" -console.print("Fetching Cybis-Pereira 2026 dataset (for video GIF)") +console.print("Fetching Cybis-Pereira 2026 dataset (for video and registration GIFs)") video_bids_root = fetch_cybis_pereira_2026( datasets=_VIDEO_DATASETS, - subjects=[_VIDEO_SUBJECT], - sessions=[_VIDEO_SESSION], - acqs=[_VIDEO_ACQ_SLICE], + subjects=[_VIDEO_SUBJECT, _VOLUMEWISE_SUBJECT], + sessions=[_VIDEO_SESSION, _VOLUMEWISE_SESSION], + acqs=[_VIDEO_ACQ_SLICE, _VOLUMEWISE_ACQ_SLICE], ) _VIDEO_FUSI_PATH = ( @@ -415,13 +425,17 @@ def _open_accordion(widget, idx: int) -> None: get_qapp().processEvents() -def _accordion_index(widget, title: str) -> int: - """Return the accordion position of the section with the given title. +def _open_accordion_panel(widget, title: str): + """Open accordion panel *title* and return its widget. - Resolving the index from the section title keeps the screenshot code robust - to accordion reordering or the insertion of new sections. + This avoids hard-coding panel indices in the screenshot script, which is + brittle when the plugin adds or reorders sections. """ - return list(widget._accordion_panels).index(title) + for idx, (btn, _) in enumerate(widget._accordion_btns): + if btn.text() == title: + _open_accordion(widget, idx) + return widget._accordion_panels[title] + raise KeyError(f"Accordion panel not found: {title}") # --------------------------------------------------------------------------- @@ -472,12 +486,7 @@ def _accordion_index(widget, title: str) -> int: viewer2.window.add_dock_widget(widget2, name="ConfUSIus", area="right") _qt_sleep(200) - # Open Signals panel (index 2). - _open_accordion(widget2, 2) - - # Retrieve the Signals panel from the accordion container layout. - _container2 = widget2._accordion_btns[0][0].parent() - ts_panel = _container2.layout().itemAt(2 * 2 + 1).widget() + ts_panel = _open_accordion_panel(widget2, "Signals") # Open the bottom dock with the signals plotter. plotter = ts_panel._ensure_plotter() @@ -521,10 +530,7 @@ def _accordion_index(widget, title: str) -> int: viewer3.window.add_dock_widget(widget3, name="ConfUSIus", area="right") _qt_sleep(200) - # Open the Quality Control panel, resolving its position by name so the - # screenshot does not break when sections are reordered or inserted. - _open_accordion(widget3, _accordion_index(widget3, "Quality Control")) - qc_panel = widget3._accordion_panels["Quality Control"] + qc_panel = _open_accordion_panel(widget3, "Quality Control") # Select the layer in the QC panel. idx = qc_panel._layer_combo.findText(layer_name) @@ -574,10 +580,7 @@ def _accordion_index(widget, title: str) -> int: viewer4.window.add_dock_widget(widget4, name="ConfUSIus", area="right") _qt_sleep(200) - # Open Signals panel (index 2). - _open_accordion(widget4, 2) - _container4 = widget4._accordion_btns[0][0].parent() - ts_panel4 = _container4.layout().itemAt(2 * 2 + 1).widget() + ts_panel4 = _open_accordion_panel(widget4, "Signals") layer4 = viewer4.layers[0] shape4 = layer4.data.shape[1:] # (z, y, x) @@ -640,10 +643,7 @@ def _accordion_index(widget, title: str) -> int: viewer5.window.add_dock_widget(widget5, name="ConfUSIus", area="right") _qt_sleep(200) - # Open Signals panel (index 2). - _open_accordion(widget5, 2) - _container5 = widget5._accordion_btns[0][0].parent() - ts_panel5 = _container5.layout().itemAt(2 * 2 + 1).widget() + ts_panel5 = _open_accordion_panel(widget5, "Signals") layer5 = viewer5.layers[0] shape5 = layer5.data.shape[1:] # (z, y, x) @@ -729,8 +729,13 @@ def _accordion_index(widget, title: str) -> int: video_panel._load_from_path() _qt_sleep(200) - # Open the Video accordion section (index 1). - _open_accordion(widget6, 1) + video_layer = next( + layer + for layer in viewer6.layers + if layer is not fusi_layer and layer.name.startswith("Video:") + ) + + _open_accordion_panel(widget6, "Video") # Size the window, then refit camera to layers (napari "home" button). win6 = viewer6.window._qt_window @@ -747,18 +752,26 @@ def _accordion_index(widget, title: str) -> int: N_GIF_FRAMES = 60 GIF_FPS = 12 GIF_WIDTH = 1100 - # Scrub from 2 s to 17 s of scan world time. Use `set_point` (world - # coordinate) instead of `set_current_step` (index), because fUSI and - # video layers have different time scales in the shared grid. - GIF_T_START_S, GIF_T_STOP_S = 2.0, 17.0 - step_times = np.linspace(GIF_T_START_S, GIF_T_STOP_S, N_GIF_FRAMES) + # Scrub inside the actual time overlap between the fUSI and video layers. + # Use `set_point` (world coordinate) instead of `set_current_step` (index), + # because the layers have different time scales in the shared grid. + fusi_min, fusi_max = fusi_layer.extent.world[:, VIDEO_TIME_AXIS] + video_min, video_max = video_layer.extent.world[:, VIDEO_TIME_AXIS] + gif_t_start_s = max(float(fusi_min), float(video_min)) + 0.5 + gif_t_stop_s = min(float(fusi_max), float(video_max)) - 0.5 + if gif_t_stop_s <= gif_t_start_s: + raise RuntimeError("No overlapping fUSI/video time range available for GIF") + step_times = np.linspace(gif_t_start_s, gif_t_stop_s, N_GIF_FRAMES) frames_pil: list = [] for t in step_times: viewer6.dims.set_point(VIDEO_TIME_AXIS, float(t)) get_qapp().processEvents() get_qapp().processEvents() - raw = viewer6.screenshot(canvas_only=False)[..., :3] + raw = viewer6.screenshot(canvas_only=False) + if raw.size == 0: + raise RuntimeError("napari returned an empty screenshot frame") + raw = raw[..., :3] h, w = raw.shape[:2] scale = GIF_WIDTH / w frames_pil.append( @@ -784,7 +797,345 @@ def _accordion_index(widget, title: str) -> int: _warn(f"plugin-video.gif failed: {exc}") # --------------------------------------------------------------------------- -# 7. Events panel — GIF of creating events with the Start/End workflow +# 7. Registration panel — rigid between-session GIF via the plugin +# --------------------------------------------------------------------------- + +try: + from PIL import Image as _PILImage + + from confusius._napari._registration._panel import RegistrationPanel + from confusius.plotting.napari import plot_napari + + fixed_path = ( + bids_root + / f"sub-{_REGISTRATION_SUBJECT}/ses-{_REGISTRATION_FIXED_SESSION}/angio" + / f"sub-{_REGISTRATION_SUBJECT}_ses-{_REGISTRATION_FIXED_SESSION}_pwd.nii.gz" + ) + moving_path = ( + bids_root + / f"sub-{_REGISTRATION_SUBJECT}/ses-{_REGISTRATION_MOVING_SESSION}/angio" + / f"sub-{_REGISTRATION_SUBJECT}_ses-{_REGISTRATION_MOVING_SESSION}_pwd.nii.gz" + ) + + fixed_da = cf.load(fixed_path) + moving_da = cf.load(moving_path) + registration_contrast = ( + min(float(fixed_da.min()), float(moving_da.min())), + max(float(fixed_da.quantile(0.9995)), float(moving_da.quantile(0.9995))), + ) + + viewer7 = napari.Viewer(show=False) + _viewer7, fixed_layer7 = plot_napari( + fixed_da, + viewer=viewer7, + name=f"Fixed angio ({_REGISTRATION_FIXED_SESSION})", + gamma=0.4, + colormap="red", + contrast_limits=registration_contrast, + ) + _viewer7, moving_layer7 = plot_napari( + moving_da, + viewer=viewer7, + name=f"Moving angio ({_REGISTRATION_MOVING_SESSION})", + gamma=0.4, + colormap="cyan", + blending="additive", + contrast_limits=registration_contrast, + ) + + widget7 = ConfUSIusWidget(viewer7) + viewer7.window.add_dock_widget(widget7, name="ConfUSIus", area="right") + _qt_sleep(250) + + registration_panel = widget7.findChild(RegistrationPanel) + if registration_panel is None: + raise RuntimeError("RegistrationPanel not found in ConfUSIusWidget") + + _open_accordion_panel(widget7, "Registration") + + moving_idx = registration_panel._moving_combo.findText(moving_layer7.name) + if moving_idx >= 0: + registration_panel._moving_combo.setCurrentIndex(moving_idx) + fixed_idx = registration_panel._fixed_combo.findText(fixed_layer7.name) + if fixed_idx >= 0: + registration_panel._fixed_combo.setCurrentIndex(fixed_idx) + registration_panel._single_volume_radio.setChecked(True) + registration_panel._transform_combo.setCurrentText("rigid") + registration_panel._metric_combo.setCurrentText("correlation") + scale_idx = registration_panel._scale_combo.findData("off") + if scale_idx >= 0: + registration_panel._scale_combo.setCurrentIndex(scale_idx) + registration_panel._learning_rate_auto_check.setChecked(False) + registration_panel._learning_rate_edit.setValue(0.01) + registration_panel._iterations_spin.setValue(500) + registration_panel._validate_registration_selection() + _qt_sleep(100) + + win7 = viewer7.window._qt_window + win7.setAttribute(Qt.WA_DontShowOnScreen) + win7.show() + win7.resize(1500, 950) + get_qapp().processEvents() + viewer7.reset_view() + if fixed_da.sizes.get("z", 1) > 1: + viewer7.dims.set_point(0, fixed_da.sizes["z"] // 2) + _qt_sleep(100) + + frames_pil: list = [] + + def _capture_registration_frame() -> None: + raw = viewer7.screenshot(canvas_only=False) + if raw.size == 0: + raise RuntimeError("napari returned an empty registration GIF frame") + raw = raw[..., :3] + h, w = raw.shape[:2] + gif_width = 1200 + scale = gif_width / w + frames_pil.append( + _PILImage.fromarray(raw).resize( + (gif_width, int(h * scale)), _PILImage.Resampling.LANCZOS + ) + ) + + def _capture_z_sweep(n_frames: int = 20) -> None: + if fixed_da.sizes.get("z", 1) <= 1: + return + z_mid = fixed_da.sizes["z"] // 2 + z_max = fixed_da.sizes["z"] - 1 + segments = [ + np.linspace(z_mid, z_max, n_frames), + np.linspace(z_max, z_mid, n_frames)[1:], + np.linspace(z_mid, 0, n_frames)[1:], + np.linspace(0, z_mid, n_frames)[1:], + ] + for z in np.concatenate(segments): + viewer7.dims.set_current_step(0, int(round(float(z)))) + get_qapp().processEvents() + _qt_sleep(180) + _capture_registration_frame() + + _capture_z_sweep() + if fixed_da.sizes.get("z", 1) > 1: + viewer7.dims.set_current_step(0, fixed_da.sizes["z"] // 2) + for _ in range(6): + get_qapp().processEvents() + _qt_sleep(120) + _capture_registration_frame() + + registration_panel._run_registration() + if registration_panel._worker is None: + raise RuntimeError("Registration worker did not start") + + while registration_panel._worker is not None and len(frames_pil) < 96: + if fixed_da.sizes.get("z", 1) > 1: + viewer7.dims.set_current_step(0, fixed_da.sizes["z"] // 2) + get_qapp().processEvents() + _qt_sleep(90) + _capture_registration_frame() + + while registration_panel._worker is not None: + get_qapp().processEvents() + _qt_sleep(120) + + fixed_layer7.visible = False + moving_layer7.visible = False + get_qapp().processEvents() + + if fixed_da.sizes.get("z", 1) > 1: + viewer7.dims.set_current_step(0, fixed_da.sizes["z"] // 2) + for _ in range(8): + get_qapp().processEvents() + _qt_sleep(120) + _capture_registration_frame() + + _capture_z_sweep() + if fixed_da.sizes.get("z", 1) > 1: + viewer7.dims.set_current_step(0, fixed_da.sizes["z"] // 2) + for _ in range(6): + get_qapp().processEvents() + _qt_sleep(120) + _capture_registration_frame() + + palette_src = frames_pil[0].quantize(colors=256, dither=0) + quantized = [frame.quantize(palette=palette_src, dither=0) for frame in frames_pil] + + gif_path = str(HERE / "plugin-registration.gif") + quantized[0].save( + gif_path, + save_all=True, + append_images=quantized[1:], + duration=1000 // 16, + loop=0, + ) + viewer7.close() + _ok("Saved plugin-registration.gif") +except Exception as exc: + _warn(f"plugin-registration.gif failed: {exc}") + +# --------------------------------------------------------------------------- +# 8. Registration panel — within-scan motion-correction GIF via the plugin +# --------------------------------------------------------------------------- + +try: + from PIL import Image as _PILImage + + from confusius._napari._registration._panel import RegistrationPanel + from confusius.plotting.napari import plot_napari + + volumewise_path = ( + video_bids_root + / f"sub-{_VOLUMEWISE_SUBJECT}/ses-{_VOLUMEWISE_SESSION}/fusi" + / f"sub-{_VOLUMEWISE_SUBJECT}_ses-{_VOLUMEWISE_SESSION}_task-openfield_acq-{_VOLUMEWISE_ACQ_SLICE}_pwd.nii.gz" + ) + + volumewise_da = cf.load(volumewise_path).isel(time=slice(220, 340)).compute() + volumewise_contrast = ( + float(volumewise_da.min()), + float(volumewise_da.quantile(0.9995)), + ) + n_frames = volumewise_da.sizes["time"] + + viewer8 = napari.Viewer(show=False) + _viewer8, moving_layer8 = plot_napari( + volumewise_da, + viewer=viewer8, + name=f"Open field ({_VOLUMEWISE_SESSION})", + gamma=0.4, + contrast_limits=volumewise_contrast, + ) + + widget8 = ConfUSIusWidget(viewer8) + viewer8.window.add_dock_widget(widget8, name="ConfUSIus", area="right") + _qt_sleep(250) + + registration_panel8 = widget8.findChild(RegistrationPanel) + if registration_panel8 is None: + raise RuntimeError("RegistrationPanel not found in ConfUSIusWidget") + + _open_accordion_panel(widget8, "Registration") + + moving_idx = registration_panel8._moving_combo.findText(moving_layer8.name) + if moving_idx >= 0: + registration_panel8._moving_combo.setCurrentIndex(moving_idx) + registration_panel8._time_series_radio.setChecked(True) + registration_panel8._transform_combo.setCurrentText("rigid") + registration_panel8._metric_combo.setCurrentText("correlation") + scale_idx = registration_panel8._scale_combo.findData("off") + if scale_idx >= 0: + registration_panel8._scale_combo.setCurrentIndex(scale_idx) + registration_panel8._reference_time_spin.setValue(n_frames // 2) + registration_panel8._learning_rate_auto_check.setChecked(False) + registration_panel8._learning_rate_edit.setValue(1.0) + registration_panel8._n_jobs_spin.setValue(-1) + registration_panel8._keep_diagnostics_check.setChecked(False) + registration_panel8._validate_registration_selection() + _qt_sleep(100) + + win8 = viewer8.window._qt_window + win8.setAttribute(Qt.WA_DontShowOnScreen) + win8.show() + win8.resize(1500, 950) + get_qapp().processEvents() + viewer8.reset_view() + viewer8.dims.set_current_step(0, 0) + _qt_sleep(100) + + frames_pil: list = [] + + def _capture_volumewise_frame() -> None: + raw = viewer8.screenshot(canvas_only=False) + if raw.size == 0: + raise RuntimeError( + "napari returned an empty volumewise registration GIF frame" + ) + raw = raw[..., :3] + h, w = raw.shape[:2] + gif_width = 1200 + scale = gif_width / w + frames_pil.append( + _PILImage.fromarray(raw).resize( + (gif_width, int(h * scale)), _PILImage.Resampling.LANCZOS + ) + ) + + def _capture_time_sweep(n_frames_sweep: int = 18) -> None: + forward = np.linspace(0, n_frames - 1, n_frames_sweep) + backward = np.linspace(n_frames - 1, 0, n_frames_sweep)[1:] + for t in np.concatenate([forward, backward]): + viewer8.dims.set_current_step(0, int(round(float(t)))) + get_qapp().processEvents() + _qt_sleep(70) + _capture_volumewise_frame() + + _capture_time_sweep() + viewer8.dims.set_current_step(0, 0) + for _ in range(4): + get_qapp().processEvents() + _qt_sleep(90) + _capture_volumewise_frame() + + registration_panel8._run_registration() + if registration_panel8._worker is None: + raise RuntimeError("Volumewise registration worker did not start") + + registration_start_frames = len(frames_pil) + while ( + registration_panel8._worker is not None + and len(frames_pil) - registration_start_frames < 140 + ): + viewer8.dims.set_current_step(0, 0) + get_qapp().processEvents() + _qt_sleep(100) + _capture_volumewise_frame() + + while registration_panel8._worker is not None: + get_qapp().processEvents() + _qt_sleep(100) + + moving_layer8.visible = False + try: + viewer8.layers["Moving"].visible = False + except KeyError: + pass + try: + motion_corrected_layer = viewer8.layers["Motion corrected"] + except KeyError: + pass + else: + motion_corrected_layer.colormap = "gray" + motion_corrected_layer.blending = "translucent_no_depth" + get_qapp().processEvents() + + viewer8.dims.set_current_step(0, n_frames // 2) + for _ in range(4): + get_qapp().processEvents() + _qt_sleep(90) + _capture_volumewise_frame() + + _capture_time_sweep() + viewer8.dims.set_current_step(0, n_frames // 2) + for _ in range(4): + get_qapp().processEvents() + _qt_sleep(90) + _capture_volumewise_frame() + + palette_src = frames_pil[0].quantize(colors=256, dither=0) + quantized = [frame.quantize(palette=palette_src, dither=0) for frame in frames_pil] + + gif_path = str(HERE / "plugin-registration-volumewise.gif") + quantized[0].save( + gif_path, + save_all=True, + append_images=quantized[1:], + duration=1000 // 16, + loop=0, + ) + viewer8.close() + _ok("Saved plugin-registration-volumewise.gif") +except Exception as exc: + _warn(f"plugin-registration-volumewise.gif failed: {exc}") +# --------------------------------------------------------------------------- +# 9. Events panel — GIF of creating events with the Start/End workflow # --------------------------------------------------------------------------- try: @@ -798,15 +1149,15 @@ def _accordion_index(widget, title: str) -> int: da_gif = cf.load(_VIDEO_FUSI_PATH) - viewer8 = napari.Viewer(show=False) - _viewer8, fusi8 = plot_napari( + viewer9 = napari.Viewer(show=False) + _viewer9, fusi9 = plot_napari( da_gif, - viewer=viewer8, + viewer=viewer9, gamma=DISPLAY_GAMMA, contrast_limits=VIDEO_DISPLAY_CONTRAST, ) - widget8 = ConfUSIusWidget(viewer8) - viewer8.window.add_dock_widget(widget8, name="ConfUSIus", area="right") + widget9 = ConfUSIusWidget(viewer9) + viewer9.window.add_dock_widget(widget9, name="ConfUSIus", area="right") _qt_sleep(200) # Two behavioural events to annotate, with absolute world-time onsets and @@ -818,93 +1169,93 @@ def _accordion_index(widget, title: str) -> int: settle_t = grooming_onset + grooming_duration # --- Labels layer aligned to the fUSI spatial axes, two painted regions. --- - da_meta8 = fusi8.metadata["xarray"] - spatial8 = [i for i, d in enumerate(da_meta8.dims) if d in ("pose", "z", "y", "x")] - spatial_shape8 = tuple(da_meta8.shape[i] for i in spatial8) - spatial_scale8 = tuple(float(fusi8.scale[i]) for i in spatial8) - spatial_translate8 = tuple(float(fusi8.translate[i]) for i in spatial8) - _ny8, _nx8 = spatial_shape8[-2], spatial_shape8[-1] - _yy8, _xx8 = np.ogrid[:_ny8, :_nx8] - _r8 = 0.05 * min(_ny8, _nx8) - _blob1 = ((_yy8 - 0.10 * _ny8) ** 2 + (_xx8 - 0.63 * _nx8) ** 2) < _r8**2 - _blob2 = ((_yy8 - 0.35 * _ny8) ** 2 + (_xx8 - 0.73 * _nx8) ** 2) < _r8**2 - _label_data8 = np.zeros(spatial_shape8, dtype=np.int32) - _label_data8[0][_blob1] = 1 - _label_data8[0][_blob2] = 2 - labels8 = viewer8.add_labels( - _label_data8, + da_meta9 = fusi9.metadata["xarray"] + spatial9 = [i for i, d in enumerate(da_meta9.dims) if d in ("pose", "z", "y", "x")] + spatial_shape9 = tuple(da_meta9.shape[i] for i in spatial9) + spatial_scale9 = tuple(float(fusi9.scale[i]) for i in spatial9) + spatial_translate9 = tuple(float(fusi9.translate[i]) for i in spatial9) + _ny9, _nx9 = spatial_shape9[-2], spatial_shape9[-1] + _yy9, _xx9 = np.ogrid[:_ny9, :_nx9] + _r9 = 0.05 * min(_ny9, _nx9) + _blob1 = ((_yy9 - 0.10 * _ny9) ** 2 + (_xx9 - 0.63 * _nx9) ** 2) < _r9**2 + _blob2 = ((_yy9 - 0.35 * _ny9) ** 2 + (_xx9 - 0.73 * _nx9) ** 2) < _r9**2 + _label_data9 = np.zeros(spatial_shape9, dtype=np.int32) + _label_data9[0][_blob1] = 1 + _label_data9[0][_blob2] = 2 + labels9 = viewer9.add_labels( + _label_data9, name="Labels (3D)", - scale=spatial_scale8, - translate=spatial_translate8, + scale=spatial_scale9, + translate=spatial_translate9, opacity=0.7, ) # --- Behavioural video via the Video panel; group fUSI + labels in one cell. --- - video_panel8 = widget8.findChild(VideoPanel) - _ref_idx8 = video_panel8._ref_combo.findText(fusi8.name) - if _ref_idx8 >= 0: - video_panel8._ref_combo.setCurrentIndex(_ref_idx8) - video_panel8._path_edit.setText(str(_VIDEO_MP4_PATH)) - video_panel8._load_from_path() + video_panel9 = widget9.findChild(VideoPanel) + _ref_idx9 = video_panel9._ref_combo.findText(fusi9.name) + if _ref_idx9 >= 0: + video_panel9._ref_combo.setCurrentIndex(_ref_idx9) + video_panel9._path_edit.setText(str(_VIDEO_MP4_PATH)) + video_panel9._load_from_path() _qt_sleep(300) # stride=2 keeps [fUSI, Labels] overlaid in one grid cell; the video gets its own. - viewer8.grid.stride = 2 + viewer9.grid.stride = 2 get_qapp().processEvents() - events_panel8 = widget8._accordion_panels["Events"] + events_panel9 = widget9._accordion_panels["Events"] # --- Signals plotter in Labels mode (mean signal per region) with the cursor. --- - signals_panel8 = widget8._accordion_panels["Signals"] - plotter8 = signals_panel8._ensure_plotter() + signals_panel9 = widget9._accordion_panels["Signals"] + plotter9 = signals_panel9._ensure_plotter() _qt_sleep(350) - plotter8.set_source_mode("labels") - plotter8.set_labels_layer(labels8) - plotter8.set_ref_layers([fusi8]) - plotter8._cursor_world = rearing_onset - plotter8.set_show_cursor(True) - plotter8._update_plot_from_labels() + plotter9.set_source_mode("labels") + plotter9.set_labels_layer(labels9) + plotter9.set_ref_layers([fusi9]) + plotter9._cursor_world = rearing_onset + plotter9.set_show_cursor(True) + plotter9._update_plot_from_labels() # Select the fUSI so events and the overlay read its true time coordinate, then # put the slider at the first onset and activate the overlay. - viewer8.layers.selection = {fusi8} - viewer8.dims.set_point(VIDEO_TIME_AXIS, rearing_onset) - widget8._time_overlay.check() + viewer9.layers.selection = {fusi9} + viewer9.dims.set_point(VIDEO_TIME_AXIS, rearing_onset) + widget9._time_overlay.check() # Open the Events accordion and show the window so the geometry is final. - _open_accordion(widget8, _accordion_index(widget8, "Events")) - win8 = viewer8.window._qt_window - win8.setAttribute(Qt.WA_DontShowOnScreen) - win8.show() - win8.resize(1400, 1050) + _open_accordion_panel(widget9, "Events") + win9 = viewer9.window._qt_window + win9.setAttribute(Qt.WA_DontShowOnScreen) + win9.show() + win9.resize(1400, 1050) get_qapp().processEvents() - viewer8.reset_view() + viewer9.reset_view() get_qapp().processEvents() # Scroll the sidebar down so the Events panel is fully visible. - scroll8 = widget8.findChild(QScrollArea) - if scroll8 is not None and scroll8.widget() is not None: - first_btn8 = widget8._accordion_btns[0][0] - scroll8.verticalScrollBar().setValue( - first_btn8.mapTo(scroll8.widget(), QPoint(0, 0)).y() + scroll9 = widget9.findChild(QScrollArea) + if scroll9 is not None and scroll9.widget() is not None: + first_btn9 = widget9._accordion_btns[0][0] + scroll9.verticalScrollBar().setValue( + first_btn9.mapTo(scroll9.widget(), QPoint(0, 0)).y() ) get_qapp().processEvents() # --- GIF frame capture -------------------------------------------------- - GIF_WIDTH8 = 1100 - GIF_FPS8 = 12 - frames8: list = [] + GIF_WIDTH9 = 1100 + GIF_FPS9 = 12 + frames9: list = [] try: badge_font = ImageFont.truetype(font_manager.findfont("DejaVu Sans:bold"), 30) except (OSError, ValueError): badge_font = ImageFont.load_default() - def _grab8(badge_text: str | None = None, repeat: int = 1) -> None: - raw = viewer8.screenshot(canvas_only=False)[..., :3] + def _grab9(badge_text: str | None = None, repeat: int = 1) -> None: + raw = viewer9.screenshot(canvas_only=False)[..., :3] h, w = raw.shape[:2] - scale = GIF_WIDTH8 / w + scale = GIF_WIDTH9 / w frame = _PILImage.fromarray(raw).resize( - (GIF_WIDTH8, int(h * scale)), _PILImage.Resampling.LANCZOS + (GIF_WIDTH9, int(h * scale)), _PILImage.Resampling.LANCZOS ) if badge_text is not None: draw = ImageDraw.Draw(frame, "RGBA") @@ -925,75 +1276,75 @@ def _grab8(badge_text: str | None = None, repeat: int = 1) -> None: fill=(233, 75, 95, 255), ) for _ in range(repeat): - frames8.append(frame) + frames9.append(frame) - def _set_cursor8(t: float) -> None: - viewer8.dims.set_point(VIDEO_TIME_AXIS, t) - plotter8._cursor_world = t + def _set_cursor9(t: float) -> None: + viewer9.dims.set_point(VIDEO_TIME_AXIS, t) + plotter9._cursor_world = t # Re-render the labels-mode plot so the cursor line moves to the new time # (the video frame in the grid also updates to this time step). - plotter8.set_show_cursor(True) - plotter8._update_plot_from_labels() - widget8._time_overlay.update() + plotter9.set_show_cursor(True) + plotter9._update_plot_from_labels() + widget9._time_overlay.update() get_qapp().processEvents() get_qapp().processEvents() - def _type_name8(name: str) -> None: + def _type_name9(name: str) -> None: """Type *name* into the event-name field one character at a time.""" - events_panel8._name_edit.setText("") + events_panel9._name_edit.setText("") get_qapp().processEvents() - _grab8(repeat=2) + _grab9(repeat=2) for i in range(1, len(name) + 1): - events_panel8._name_edit.setText(name[:i]) + events_panel9._name_edit.setText(name[:i]) get_qapp().processEvents() - _grab8() - _grab8(repeat=2) + _grab9() + _grab9(repeat=2) - def _annotate8(name: str, onset: float, duration: float) -> None: + def _annotate9(name: str, onset: float, duration: float) -> None: """Drive the full type → Start → scrub → End workflow for one event.""" - _set_cursor8(onset) - _type_name8(name) + _set_cursor9(onset) + _type_name9(name) # Start (S) marks the onset at the current time. - events_panel8._on_start() + events_panel9._on_start() get_qapp().processEvents() - _grab8(badge_text="S · Start", repeat=7) + _grab9(badge_text="S · Start", repeat=7) # Scrub the time slider forward to the offset. for t in np.linspace(onset, onset + duration, 12): - _set_cursor8(float(t)) - _grab8() + _set_cursor9(float(t)) + _grab9() # End (E) creates the event, shading the plot and filling the table. - events_panel8._on_end() + events_panel9._on_end() get_qapp().processEvents() get_qapp().processEvents() - _grab8(badge_text="E · End", repeat=7) + _grab9(badge_text="E · End", repeat=7) # 1. Annotate the "rearing" event. - _annotate8("rearing", rearing_onset, rearing_duration) + _annotate9("rearing", rearing_onset, rearing_duration) # 2. Travel forward (no recording) to the "grooming" onset. for t in np.linspace(rearing_onset + rearing_duration, grooming_onset, 14): - _set_cursor8(float(t)) - _grab8() + _set_cursor9(float(t)) + _grab9() # 3. Annotate the "grooming" event. - _annotate8("grooming", grooming_onset, grooming_duration) + _annotate9("grooming", grooming_onset, grooming_duration) # 4. Settle inside the rearing event so the overlay names the active event. - _set_cursor8(settle_t) - _grab8(repeat=12) + _set_cursor9(settle_t) + _grab9(repeat=12) # --- Assemble the GIF (shared-palette quantize, like the video GIF). --- - palette_src8 = frames8[0].quantize(colors=256, dither=0) - quantized8 = [frame.quantize(palette=palette_src8, dither=0) for frame in frames8] - gif_path8 = str(HERE / "plugin-events-create.gif") - quantized8[0].save( - gif_path8, + palette_src9 = frames9[0].quantize(colors=256, dither=0) + quantized9 = [frame.quantize(palette=palette_src9, dither=0) for frame in frames9] + gif_path9 = str(HERE / "plugin-events-create.gif") + quantized9[0].save( + gif_path9, save_all=True, - append_images=quantized8[1:], - duration=1000 // GIF_FPS8, + append_images=quantized9[1:], + duration=1000 // GIF_FPS9, loop=0, ) - viewer8.close() + viewer9.close() _ok("Saved plugin-events-create.gif") except Exception as exc: _warn(f"plugin-events-create.gif failed: {exc}") diff --git a/src/confusius/_napari/_data/_save_panel.py b/src/confusius/_napari/_data/_save_panel.py index 82c58236..b35ff653 100644 --- a/src/confusius/_napari/_data/_save_panel.py +++ b/src/confusius/_napari/_data/_save_panel.py @@ -110,8 +110,8 @@ def _setup_ui(self) -> None: "via the ConfUSIus reader." ) - layer_form.addRow("Save layer:", self._layer_combo) - layer_form.addRow("Coordinates from:", self._template_combo) + layer_form.addRow("Save layer", self._layer_combo) + layer_form.addRow("Coordinates from", self._template_combo) save_layout.addLayout(layer_form) self._path_edit = QLineEdit() diff --git a/src/confusius/_napari/_qc/_panel.py b/src/confusius/_napari/_qc/_panel.py index bdbdb51f..df589ab9 100644 --- a/src/confusius/_napari/_qc/_panel.py +++ b/src/confusius/_napari/_qc/_panel.py @@ -292,8 +292,8 @@ def _on_theme_changed(self) -> None: def _time_dim_index(self) -> int: """Return the viewer dimension index for the time dimension. - Searches all layers for xarray metadata containing a ``time`` - dimension and returns its index. Falls back to ``0`` when no + Searches all layers for xarray metadata containing a `time` + dimension and returns its index. Falls back to `0` when no suitable layer is found (same convention as the signals panel). """ for layer in self.viewer.layers: @@ -311,7 +311,7 @@ def _current_time_world(self) -> float: def _on_time_step_changed(self) -> None: """Forward the current napari time world coordinate to the cursor. - Reads the world coordinate directly from ``self.viewer.dims.point`` + Reads the world coordinate directly from `self.viewer.dims.point` for the time dimension, which is correct regardless of which layer is selected (mirrors the approach used in the signals panel). """ diff --git a/src/confusius/_napari/_qt.py b/src/confusius/_napari/_qt.py new file mode 100644 index 00000000..da4e5543 --- /dev/null +++ b/src/confusius/_napari/_qt.py @@ -0,0 +1,33 @@ +"""Shared Qt helpers for internal napari panels.""" + +from __future__ import annotations + +from qtpy.QtWidgets import QMainWindow, QWidget + + +def find_main_window(widget: QWidget) -> QMainWindow | None: + """Return the ancestor `QMainWindow` for a widget, if present. + + Parameters + ---------- + widget : QWidget + Starting widget to search from. + + Returns + ------- + QMainWindow or None + The containing main window, or `None` if no ancestor main window is + found or the Qt object was already deleted. + """ + try: + parent = widget.parent() + except RuntimeError: + return None + while parent is not None: + if isinstance(parent, QMainWindow): + return parent + try: + parent = parent.parent() + except RuntimeError: + return None + return None diff --git a/src/confusius/_napari/_registration/__init__.py b/src/confusius/_napari/_registration/__init__.py new file mode 100644 index 00000000..c34f67fa --- /dev/null +++ b/src/confusius/_napari/_registration/__init__.py @@ -0,0 +1,5 @@ +"""Registration panel for the ConfUSIus napari plugin.""" + +from confusius._napari._registration._panel import RegistrationPanel + +__all__ = ["RegistrationPanel"] diff --git a/src/confusius/_napari/_registration/_metric_plotter.py b/src/confusius/_napari/_registration/_metric_plotter.py new file mode 100644 index 00000000..a82eac76 --- /dev/null +++ b/src/confusius/_napari/_registration/_metric_plotter.py @@ -0,0 +1,165 @@ +"""Bottom-dock widget that plots the registration optimizer metric. + +Mirrors the [`SignalPlotter`][confusius._napari._signals._plotter.SignalPlotter] +layout—a small matplotlib figure in the bottom dock—but stays deliberately simple: a +single line chart of the per-iteration metric value. The widget is created lazily by +`RegistrationPanel` when a registration starts, and torn down on completion so the dock +returns to its pre-run layout. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +from matplotlib.backends.backend_qtagg import ( + FigureCanvasQTAgg as FigureCanvas, +) +from matplotlib.backends.backend_qtagg import ( + NavigationToolbar2QT as NavigationToolbar, +) +from matplotlib.figure import Figure +from qtpy.QtCore import QSize, QTimer +from qtpy.QtWidgets import QSizePolicy, QVBoxLayout, QWidget + +from confusius._napari._theme import get_napari_colors, style_plot_toolbar + +if TYPE_CHECKING: + from napari import Viewer + + +class RegistrationMetricPlotter(QWidget): + """Bottom-dock widget that plots the per-iteration optimizer metric. + + The widget is intentionally minimal: a single matplotlib axes, a navigation toolbar, + and a thin status footer. Layout decisions (e.g. y-axis limits, line width) follow + the same conventions as + [`SignalPlotter`][confusius._napari._signals._plotter.SignalPlotter] for visual + consistency between the two bottom-dock tabs. + + Parameters + ---------- + viewer : napari.Viewer + Active napari viewer, used to detect theme changes. + """ + + def __init__(self, viewer: "Viewer") -> None: + super().__init__() + self._viewer = viewer + self._metric_values: list[float] = [] + self._metric_line = None + # Throttle redraws to ~60 fps so rapid iteration events (and the + # arrival of their queued Qt signals) don't flood the GUI thread. + self._redraw_timer = QTimer(self) + self._redraw_timer.setSingleShot(True) + self._redraw_timer.setInterval(16) + self._redraw_timer.timeout.connect(self._render) + self.setSizePolicy( + QSizePolicy.Policy.Expanding, + QSizePolicy.Policy.Expanding, + ) + self.setMinimumHeight(300) + self._setup_ui() + self._apply_theme() + self._viewer.events.theme.connect(lambda *_: self._apply_theme()) + + def sizeHint(self) -> QSize: + """Return the preferred initial size of the widget. + + Returns + ------- + QSize + Preferred initial size of 800 x 370 pixels. + """ + return QSize(800, 370) + + def _setup_ui(self) -> None: + """Build the matplotlib canvas and toolbar.""" + layout = QVBoxLayout(self) + layout.setContentsMargins(4, 4, 4, 4) + layout.setSpacing(0) + + self._figure = Figure(tight_layout=True) + self._canvas = FigureCanvas(self._figure) + self._canvas.setSizePolicy( + QSizePolicy.Policy.Expanding, + QSizePolicy.Policy.Expanding, + ) + self._toolbar = NavigationToolbar(self._canvas, self) + layout.addWidget(self._toolbar) + layout.addWidget(self._canvas) + + self._axes = self._figure.add_subplot(111) + self._metric_line = self._axes.plot([], [], color="#e94b5f", linewidth=1.4)[0] + self._axes.set_xlabel("Iteration") + self._axes.set_ylabel("Metric value") + self._axes.set_title("Registration metric") + self._axes.grid(True, alpha=0.3) + + def _apply_theme(self) -> None: + """Re-style the axes and toolbar to match the current napari theme.""" + colors = get_napari_colors(self._viewer.theme) + self._figure.patch.set_facecolor(colors["bg"]) + self._axes.set_facecolor(colors["bg"]) + for spine in self._axes.spines.values(): + spine.set_edgecolor(colors["fg"]) + self._axes.tick_params(colors=colors["fg"]) + self._axes.xaxis.label.set_color(colors["fg"]) + self._axes.yaxis.label.set_color(colors["fg"]) + self._axes.title.set_color(colors["fg"]) + if self._metric_line is not None: + self._metric_line.set_color(colors["accent"]) + style_plot_toolbar(self._toolbar, colors) + self._canvas.draw_idle() + + def add_metric(self, value: float) -> None: + """Append a metric value and schedule a redraw. + + Called from the GUI thread via the + `NapariRegistrationProgressPlotterBridge.metric_updated` signal. Rapid iteration + events are coalesced through a single-shot timer so the canvas is redrawn at + most once per ~16 ms regardless of the worker-side event rate. + + Parameters + ---------- + value : float + Optimizer metric value at the current iteration. + """ + self._metric_values.append(float(value)) + if not self._redraw_timer.isActive(): + self._redraw_timer.start() + + def reset(self) -> None: + """Clear the metric buffer and redraw an empty plot. + + Called before each new registration run so the curve starts from + scratch instead of overlaying the previous run's data. + """ + self._redraw_timer.stop() + self._metric_values.clear() + if self._metric_line is not None: + self._metric_line.set_data([], []) + self._axes.relim() + self._axes.autoscale_view() + self._canvas.draw_idle() + + def _render(self) -> None: + """Redraw the metric line with the buffered values.""" + if self._metric_line is None: + return + n = len(self._metric_values) + self._metric_line.set_data(np.arange(1, n + 1), self._metric_values) + self._axes.relim() + self._axes.autoscale_view() + self._canvas.draw_idle() + + @property + def metric_values(self) -> list[float]: + """Copy of the metric value buffer. + + Returns + ------- + list of float + Optimizer metric value recorded at each iteration. + """ + return list(self._metric_values) diff --git a/src/confusius/_napari/_registration/_panel.py b/src/confusius/_napari/_registration/_panel.py new file mode 100644 index 00000000..61fbb499 --- /dev/null +++ b/src/confusius/_napari/_registration/_panel.py @@ -0,0 +1,1831 @@ +"""Registration panel for the ConfUSIus napari plugin.""" + +from __future__ import annotations + +from threading import Event +from typing import TYPE_CHECKING, Any, Literal, NotRequired, TypedDict, cast + +import numpy as np +from napari.qt.threading import thread_worker +from qtpy.QtCore import Qt +from qtpy.QtWidgets import ( + QApplication, + QButtonGroup, + QCheckBox, + QComboBox, + QDockWidget, + QDoubleSpinBox, + QFormLayout, + QGridLayout, + QGroupBox, + QHBoxLayout, + QLabel, + QLineEdit, + QProgressBar, + QPushButton, + QRadioButton, + QSizePolicy, + QSpinBox, + QToolButton, + QVBoxLayout, + QWidget, +) + +from confusius._dims import SPATIAL_DIMS_WITH_POSE, TIME_DIM +from confusius._napari._registration._metric_plotter import ( + RegistrationMetricPlotter, +) +from confusius._napari._registration._panel_parameters import ( + get_default_registration_parameters, + get_registration_parameters, + set_registration_parameters, +) +from confusius._napari._registration._panel_progress import ( + create_volume_progress_plotter, + setup_volumewise_progress, +) +from confusius._napari._registration._panel_results import ( + on_volume_registration_finished, + on_volumewise_registration_finished, +) +from confusius._napari._registration._panel_selection import ( + current_metric, + current_resample_interpolation, + current_scale_mode, + current_transform_model, + get_layer_by_name, + on_moving_layer_changed, + refresh_layers, + selected_layer, + set_layer_validation_style, + set_run_btn_enabled, + update_reference_time_bounds, + validate_registration_selection, +) +from confusius._napari._registration._panel_transforms import ( + apply_selected_inverse_transform, + apply_selected_transform, + get_available_transform_payloads, + get_selected_center_initialization, + get_selected_initial_transform, + load_transform, + refresh_transform_controls, + save_selected_transform, +) +from confusius._napari._registration._panel_utils import ( + ScientificDoubleSpinBox, + _apply_registration_scale, + _get_source_dataarray, + _is_registration_source_layer, + _parse_comma_separated_ints, + _prepare_between_scan_data, +) +from confusius._napari._registration._panel_worker_state import on_registration_failed +from confusius._napari._registration._progress import ( + NapariRegistrationProgressPlotterBridge, + NapariRegistrationProgressReporterBridge, +) +from confusius._napari._registration._transform_payloads import TransformPayload +from confusius.registration import register_volume, register_volumewise + +if TYPE_CHECKING: + import napari + import numpy.typing as npt + from napari.layers import Image, Layer + + +ScaleMode = Literal["off", "dB", "sqrt"] +"""Allowed registration intensity-scaling modes used by the panel.""" + +MetricName = Literal["correlation", "mattes_mi"] +"""Allowed registration metric names exposed by the panel.""" + +VolumeTransformType = Literal["translation", "rigid", "affine", "bspline"] +"""Allowed transform models for between-scan registration.""" + +VolumewiseTransformType = Literal["translation", "rigid", "affine"] +"""Allowed transform models for within-scan registration.""" + +ResampleInterpolation = Literal["linear", "bspline"] +"""Allowed interpolation modes for resampling previews and outputs.""" + +CenterInitialization = Literal["center_geometry", "center_moments"] +"""Allowed built-in center-based initialization modes.""" + +RegistrationOperation = Literal["register_volume", "register_volumewise"] +"""Allowed registration workflows handled by the panel.""" + +TransformSourceKind = Literal["loaded", "layer", "manual"] +"""Kinds of transform sources offered in the transforms UI.""" + +TransformSourceData = tuple[TransformSourceKind, str] +"""Validated transform-source selector payload `(kind, name)`.""" + +InitializationSelection = CenterInitialization | TransformSourceData | None +"""Validated initialization selection from the registration UI.""" + +RegistrationParameterMode = Literal["volume", "volumewise"] +"""Registration-parameter mode used for UI snapshot and restore helpers.""" + + +class ModeParameters(TypedDict): + """Session-scoped UI parameters for one registration mode.""" + + transform: str + metric: MetricName + scale: ScaleMode + initialization: InitializationSelection + learning_rate_auto: bool + learning_rate_value: float + number_of_iterations: int + number_of_histogram_bins: int + mesh_size: tuple[int, int, int] + convergence_minimum_value: float + convergence_window_size: int + use_multi_resolution: bool + shrink_factors: str + smoothing_sigmas: str + resample_interpolation: ResampleInterpolation + fill_value_auto: bool + fill_value: float + reference_time: int + n_jobs: int + sitk_threads: int + optimizer_weights_enabled: bool + optimizer_weights_values: list[float] + keep_diagnostics: bool + advanced_open: bool + + +class RegistrationRunPayloadBase(TypedDict): + """Shared UI snapshot fields captured before a registration worker starts.""" + + moving_layer_name: str + metric: MetricName + scale: ScaleMode + learning_rate: float | Literal["auto"] + number_of_iterations: int + use_multi_resolution: bool + resample_interpolation: ResampleInterpolation + number_of_histogram_bins: int + convergence_minimum_value: float + convergence_window_size: int + initialization: InitializationSelection + shrink_factors: tuple[int, ...] | None + smoothing_sigmas: tuple[int, ...] | None + optimizer_weights: list[float] | None + keep_diagnostics: bool + fill_value: float | None + + +class VolumeRegistrationRunPayload(RegistrationRunPayloadBase): + """UI snapshot for between-scan registration.""" + + operation: Literal["register_volume"] + transform: VolumeTransformType + mesh_size: tuple[int, int, int] + fixed_layer_name: str + fixed_mask_layer_name: str | None + moving_mask_layer_name: str | None + sitk_threads: int + initial_transform_source: NotRequired[str] + + +class VolumewiseRegistrationRunPayload(RegistrationRunPayloadBase): + """UI snapshot for within-scan registration.""" + + operation: Literal["register_volumewise"] + transform: VolumewiseTransformType + mesh_size: tuple[int, int, int] + reference_time: int + n_jobs: int + + +class ApplyTransformPayload(TypedDict): + """UI snapshot for applying an existing transform.""" + + moving_layer_name: str + target_layer_name: str + transform_source: str + direction: Literal["forward", "inverse"] + + +def _get_optimizer_weight_labels( + transform: VolumeTransformType | VolumewiseTransformType, ndim: int +) -> list[str]: + """Return user-facing optimizer-weight labels for one transform model.""" + if transform == "bspline": + return [] + if transform == "translation": + return ["tx", "ty"] if ndim == 2 else ["tx", "ty", "tz"] + if transform == "rigid": + return ( + ["angle", "tx", "ty"] + if ndim == 2 + else ["angleX", "angleY", "angleZ", "tx", "ty", "tz"] + ) + if transform == "affine": + if ndim == 2: + return ["a00", "a01", "a10", "a11", "tx", "ty"] + return [ + "a00", + "a01", + "a02", + "a10", + "a11", + "a12", + "a20", + "a21", + "a22", + "tx", + "ty", + "tz", + ] + raise ValueError(f"Unknown transform model: {transform!r}.") + + +class RegistrationPanel(QWidget): + """Right-side panel for running registration from napari. + + Parameters + ---------- + viewer : napari.Viewer + The active napari viewer instance. + """ + + def __init__(self, viewer: napari.Viewer) -> None: + super().__init__() + self.viewer = viewer + self._worker = None + self._abort_event: Event | None = None + self._loaded_transform_payload: TransformPayload | None = None + self._optimizer_weight_spins: list[QDoubleSpinBox] = [] + # Per-run progress state. Set on the GUI thread before the worker starts. + self._progress_bridge: NapariRegistrationProgressPlotterBridge | None = None + self._progress_layer: Image | None = None + self._progress_fixed_layer: Image | None = None + self._progress_moving_layer: Image | None = None + self._manual_transform_event_layers: list[Layer] = [] + self._volumewise_progress_bridge: ( + NapariRegistrationProgressReporterBridge | None + ) = None + self._volumewise_progress_layer: Image | None = None + self._volumewise_moving_preview_layer: Image | None = None + self._volumewise_progress_time_axis: int | None = None + self._volumewise_progress_total: int | None = None + # Bottom-dock metric curve. Created lazily on the first run, reused + # across subsequent runs, and torn down with the progress state. + self._metric_plotter: RegistrationMetricPlotter | None = None + self._metric_dock: QDockWidget | None = None + self._active_operation: Literal["register_volume", "register_volumewise"] = ( + "register_volume" + ) + self._registration_parameters_by_operation: dict[ + RegistrationOperation, ModeParameters + ] = {} + self._refresh_transform_controls_callback = lambda: refresh_transform_controls( + self + ) + self._save_transform_callback = lambda: save_selected_transform(self) + self._load_transform_callback = lambda: load_transform(self) + self._apply_transform_callback = lambda: apply_selected_transform(self) + self._apply_inverse_transform_callback = lambda: ( + apply_selected_inverse_transform(self) + ) + self._setup_ui() + self.viewer.layers.events.inserted.connect(self._refresh_layers) + self.viewer.layers.events.removed.connect(self._refresh_layers) + + def _make_form_label(self, text: str, *, tooltip: str | None = None) -> QLabel: + """Return a form label with an optional tooltip. + + Parameters + ---------- + text : str + Label text. + tooltip : str, optional + Tooltip shown when hovering the label. + + Returns + ------- + QLabel + Configured label widget. + """ + label = QLabel(text) + if tooltip is not None: + label.setToolTip(tooltip) + return label + + def _left_aligned_widget(self, widget: QWidget) -> QWidget: + """Wrap a compact widget so it stays left-aligned inside a form row. + + The trailing spacer has stretch 0, so an Expanding widget with a + maximum width still grows up to its cap before the spacer absorbs the + remainder; fixed-size widgets stay at their size hint. + + Parameters + ---------- + widget : QWidget + Widget to wrap. + + Returns + ------- + QWidget + Container holding `widget` left-aligned. + """ + container = QWidget() + layout = QHBoxLayout(container) + layout.setContentsMargins(0, 0, 0, 0) + layout.setSpacing(0) + layout.addWidget(widget, stretch=1) + layout.addStretch() + return container + + def _make_advanced_row( + self, + layout: QFormLayout, + label: str, + widget: QWidget, + *, + tooltip: str | None = None, + ) -> QWidget: + """Create a show/hide-able row container for one advanced parameter. + + Parameters + ---------- + layout : QFormLayout + Parent form layout receiving the row. + label : str + Row-label text. + widget : QWidget + Input widget shown on the row. + tooltip : str, optional + Tooltip shown on the row label. + + Returns + ------- + QWidget + Container widget added to `layout`. + """ + container = QWidget() + # A one-row QFormLayout (not QHBoxLayout) so the label can wrap above + # the field on narrow docks: an HBox row's minimum width is label + + # field, which forced horizontal overflow when the advanced section + # opened (see issue #183 for the same overflow in the signals panel). + row_layout = QFormLayout(container) + row_layout.setContentsMargins(0, 0, 0, 0) + row_layout.setRowWrapPolicy(QFormLayout.RowWrapPolicy.WrapLongRows) + row_layout.setFieldGrowthPolicy( + QFormLayout.FieldGrowthPolicy.ExpandingFieldsGrow + ) + lbl = self._make_form_label(label, tooltip=tooltip) + lbl.setSizePolicy(QSizePolicy.Policy.Preferred, QSizePolicy.Policy.Fixed) + row_layout.addRow(lbl, widget) + layout.addRow(container) + return container + + def _setup_ui(self) -> None: + layout = QVBoxLayout(self) + layout.setContentsMargins(10, 10, 10, 10) + layout.setSpacing(8) + + self._panel_group = QButtonGroup(self) + panel_row = QHBoxLayout() + panel_row.setSpacing(0) + self._register_panel_radio = QPushButton("Register") + self._register_panel_radio.setCheckable(True) + self._transforms_panel_radio = QPushButton("Transforms") + self._transforms_panel_radio.setCheckable(True) + self._register_panel_radio.setChecked(True) + self._panel_group.addButton(self._register_panel_radio) + self._panel_group.addButton(self._transforms_panel_radio) + self._register_panel_radio.setSizePolicy( + QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Fixed + ) + self._transforms_panel_radio.setSizePolicy( + QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Fixed + ) + _segment_btn_style = """ + QPushButton { + border-radius: 0; + } + QPushButton:checked { + background: #e94b5f; + color: white; + font-weight: bold; + } + """ + self._register_panel_radio.setStyleSheet( + _segment_btn_style + + "border-top-right-radius: 0; border-bottom-right-radius: 0;" + ) + self._transforms_panel_radio.setStyleSheet( + _segment_btn_style + + "border-top-left-radius: 0; border-bottom-left-radius: 0;" + ) + panel_row.addWidget(self._register_panel_radio) + panel_row.addWidget(self._transforms_panel_radio) + layout.addLayout(panel_row) + + operation_group = QGroupBox("Registration") + operation_layout = QFormLayout(operation_group) + operation_layout.setSpacing(6) + operation_layout.setRowWrapPolicy(QFormLayout.RowWrapPolicy.WrapLongRows) + operation_layout.setFieldGrowthPolicy( + QFormLayout.FieldGrowthPolicy.ExpandingFieldsGrow + ) + + self._mode_group = QButtonGroup(self) + mode_row = QHBoxLayout() + self._single_volume_radio = QRadioButton("Between scans") + self._time_series_radio = QRadioButton("Within-scan") + self._single_volume_radio.setChecked(True) + self._mode_group.addButton(self._single_volume_radio) + self._mode_group.addButton(self._time_series_radio) + mode_row.addWidget(self._single_volume_radio) + mode_row.addWidget(self._time_series_radio) + operation_layout.addRow( + self._make_form_label( + "Mode", + tooltip="Registration workflow. Use 'Between scans' for moving/fixed registration and 'Within-scan' for frame-to-reference motion correction.", + ), + mode_row, + ) + + self._moving_label = QLabel("Moving layer") + self._moving_combo = QComboBox() + self._moving_combo.setSizeAdjustPolicy( + QComboBox.SizeAdjustPolicy.AdjustToMinimumContentsLengthWithIcon + ) + self._moving_combo.setSizePolicy( + QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Fixed + ) + self._moving_combo.currentTextChanged.connect(self._on_moving_layer_changed) + self._moving_label.setToolTip( + "Layer containing the moving image or time series to register." + ) + operation_layout.addRow(self._moving_label, self._moving_combo) + + self._moving_mask_combo = QComboBox() + self._moving_mask_combo.setSizeAdjustPolicy( + QComboBox.SizeAdjustPolicy.AdjustToMinimumContentsLengthWithIcon + ) + self._moving_mask_combo.setSizePolicy( + QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Fixed + ) + self._moving_mask_combo.setToolTip( + "Optional Labels layer used as the moving-image metric mask. Nonzero labels are treated as True." + ) + self._new_moving_mask_btn = QPushButton("+") + self._new_moving_mask_btn.setStyleSheet("font-weight: bold; font-size: 14px;") + self._new_moving_mask_btn.setToolTip( + "Create a new 3D Labels layer (no time axis) aligned to the current image." + ) + self._new_moving_mask_btn.clicked.connect( + lambda: self._create_labels_layer(name="Moving mask") + ) + moving_mask_row = QHBoxLayout() + moving_mask_row.setContentsMargins(0, 0, 0, 0) + moving_mask_row.setSpacing(6) + moving_mask_row.addWidget(self._moving_mask_combo, stretch=1) + moving_mask_row.addWidget(self._new_moving_mask_btn) + moving_mask_container = QWidget() + moving_mask_container.setLayout(moving_mask_row) + self._moving_mask_row = moving_mask_container + self._moving_mask_label = self._make_form_label( + "Moving mask", + tooltip="Optional Labels layer used as the moving-image metric mask. Nonzero labels are treated as True.", + ) + operation_layout.addRow(self._moving_mask_label, moving_mask_container) + + self._fixed_label = QLabel("Fixed layer") + self._fixed_combo = QComboBox() + self._fixed_combo.setSizeAdjustPolicy( + QComboBox.SizeAdjustPolicy.AdjustToMinimumContentsLengthWithIcon + ) + self._fixed_combo.setSizePolicy( + QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Fixed + ) + self._fixed_label.setToolTip( + "Reference layer that defines the registration target grid." + ) + operation_layout.addRow(self._fixed_label, self._fixed_combo) + + self._fixed_mask_combo = QComboBox() + self._fixed_mask_combo.setSizeAdjustPolicy( + QComboBox.SizeAdjustPolicy.AdjustToMinimumContentsLengthWithIcon + ) + self._fixed_mask_combo.setSizePolicy( + QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Fixed + ) + self._fixed_mask_combo.setToolTip( + "Optional Labels layer used as the fixed-image metric mask. Nonzero labels are treated as True." + ) + self._new_fixed_mask_btn = QPushButton("+") + self._new_fixed_mask_btn.setStyleSheet("font-weight: bold; font-size: 14px;") + self._new_fixed_mask_btn.setToolTip( + "Create a new 3D Labels layer (no time axis) aligned to the current image." + ) + self._new_fixed_mask_btn.clicked.connect( + lambda: self._create_labels_layer(name="Fixed mask") + ) + fixed_mask_row = QHBoxLayout() + fixed_mask_row.setContentsMargins(0, 0, 0, 0) + fixed_mask_row.setSpacing(6) + fixed_mask_row.addWidget(self._fixed_mask_combo, stretch=1) + fixed_mask_row.addWidget(self._new_fixed_mask_btn) + fixed_mask_container = QWidget() + fixed_mask_container.setLayout(fixed_mask_row) + self._fixed_mask_row = fixed_mask_container + self._fixed_mask_label = self._make_form_label( + "Fixed mask", + tooltip="Optional Labels layer used as the fixed-image metric mask. Nonzero labels are treated as True.", + ) + operation_layout.addRow(self._fixed_mask_label, fixed_mask_container) + + self._reference_time_label = QLabel("Reference volume") + self._reference_time_spin = QSpinBox() + self._reference_time_spin.setMinimum(0) + self._reference_time_spin.setMaximumWidth(64) + self._reference_time_label.setToolTip( + "Volume index used as the registration target for within-scan motion correction." + ) + operation_layout.addRow(self._reference_time_label, self._reference_time_spin) + + self._n_jobs_spin = QSpinBox() + # Expanding: QFormLayout's ExpandingFieldsGrow only grows fields whose + # policy is Expanding/MinimumExpanding, and spinboxes are not by + # default — they would otherwise stay at hint width in advanced rows. + self._n_jobs_spin.setSizePolicy( + QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Fixed + ) + self._n_jobs_spin.setRange(-128, 128) + self._n_jobs_spin.setMaximumWidth(56) + self._n_jobs_spin.setSpecialValueText("auto") + self._n_jobs_spin.setToolTip( + "Number of workers for time-series registration. -1 uses all CPUs." + ) + + self._layer_validation = QLabel("") + self._layer_validation.setWordWrap(True) + self._layer_validation.setObjectName("status_err") + self._layer_validation.hide() + operation_layout.addRow(self._layer_validation) + + params_group = QGroupBox("Parameters") + params_layout = QFormLayout(params_group) + params_layout.setSpacing(6) + params_layout.setRowWrapPolicy(QFormLayout.RowWrapPolicy.WrapLongRows) + params_layout.setFieldGrowthPolicy( + QFormLayout.FieldGrowthPolicy.ExpandingFieldsGrow + ) + + self._transform_combo = QComboBox() + self._transform_combo.setMinimumContentsLength(14) + self._transform_combo.setSizeAdjustPolicy( + QComboBox.SizeAdjustPolicy.AdjustToMinimumContentsLengthWithIcon + ) + self._transform_combo.setSizePolicy( + QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Fixed + ) + params_layout.addRow( + self._make_form_label( + "Transform", + tooltip="Transform model optimized during registration: translation, rigid, affine, or bspline for between-scan registration.", + ), + self._transform_combo, + ) + + self._mesh_size_z_spin = QSpinBox() + self._mesh_size_z_spin.setRange(1, 512) + self._mesh_size_z_spin.setMaximumWidth(48) + self._mesh_size_z_spin.setToolTip("B-spline mesh size along z.") + self._mesh_size_y_spin = QSpinBox() + self._mesh_size_y_spin.setRange(1, 512) + self._mesh_size_y_spin.setMaximumWidth(48) + self._mesh_size_y_spin.setToolTip("B-spline mesh size along y.") + self._mesh_size_x_spin = QSpinBox() + self._mesh_size_x_spin.setRange(1, 512) + self._mesh_size_x_spin.setMaximumWidth(48) + self._mesh_size_x_spin.setToolTip("B-spline mesh size along x.") + self._mesh_size_row = QWidget() + mesh_size_layout = QVBoxLayout(self._mesh_size_row) + mesh_size_layout.setContentsMargins(0, 0, 0, 0) + mesh_size_layout.setSpacing(4) + mesh_size_label = self._make_form_label( + "Mesh size", + tooltip="B-spline mesh size used for B-spline registration.", + ) + mesh_size_layout.addWidget(mesh_size_label) + mesh_size_inputs = QHBoxLayout() + mesh_size_inputs.setContentsMargins(0, 0, 0, 0) + mesh_size_inputs.setSpacing(6) + mesh_size_inputs.addWidget(QLabel("Z")) + mesh_size_inputs.addWidget(self._mesh_size_z_spin) + mesh_size_inputs.addWidget(QLabel("Y")) + mesh_size_inputs.addWidget(self._mesh_size_y_spin) + mesh_size_inputs.addWidget(QLabel("X")) + mesh_size_inputs.addWidget(self._mesh_size_x_spin) + mesh_size_inputs.addStretch(1) + mesh_size_layout.addLayout(mesh_size_inputs) + params_layout.addRow(self._mesh_size_row) + + self._metric_combo = QComboBox() + self._metric_combo.setMinimumContentsLength(14) + self._metric_combo.setSizeAdjustPolicy( + QComboBox.SizeAdjustPolicy.AdjustToMinimumContentsLengthWithIcon + ) + self._metric_combo.setSizePolicy( + QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Fixed + ) + self._metric_combo.addItems(["correlation", "mattes_mi"]) + params_layout.addRow( + self._make_form_label( + "Metric", + tooltip="Similarity metric optimized during registration. 'correlation' suits same-modality data; 'mattes_mi' is more robust across intensity changes.", + ), + self._metric_combo, + ) + + self._scale_combo = QComboBox() + self._scale_combo.setMinimumContentsLength(10) + self._scale_combo.setSizeAdjustPolicy( + QComboBox.SizeAdjustPolicy.AdjustToMinimumContentsLengthWithIcon + ) + self._scale_combo.setSizePolicy( + QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Fixed + ) + self._scale_combo.addItem("decibel", "dB") + self._scale_combo.addItem("square root", "sqrt") + self._scale_combo.addItem("none", "off") + self._scale_combo.setToolTip( + "Optional intensity preprocessing applied before registration and used for registration preview layers." + ) + params_layout.addRow( + self._make_form_label( + "Scale", + tooltip="Optional intensity preprocessing applied before registration and used for registration preview layers.", + ), + self._scale_combo, + ) + + self._initialization_combo = QComboBox() + self._initialization_combo.setSizeAdjustPolicy( + QComboBox.SizeAdjustPolicy.AdjustToMinimumContentsLengthWithIcon + ) + self._initialization_combo.setSizePolicy( + QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Fixed + ) + self._initialization_combo.setToolTip( + "How to initialize registration before optimization: center geometry, center moments, identity, or an existing affine transform from the Transforms tab." + ) + params_layout.addRow( + self._make_form_label( + "Initialization", + tooltip="How to initialize registration before optimization: center geometry, center moments, identity, or an existing affine transform from the Transforms tab.", + ), + self._initialization_combo, + ) + + learning_rate_row = QHBoxLayout() + self._learning_rate_auto_check = QCheckBox("Auto") + self._learning_rate_edit = ScientificDoubleSpinBox() + self._learning_rate_edit.setRange(1e-10, 1e3) + self._learning_rate_edit.setSingleStep(0.1) + self._learning_rate_edit.setToolTip( + "Optimizer step size. Accepts decimal (0.1) or scientific notation (1e-5)." + ) + self._learning_rate_edit.setEnabled(False) + self._learning_rate_auto_check.toggled.connect( + lambda checked: self._learning_rate_edit.setEnabled(not checked) + ) + learning_rate_row.addWidget(self._learning_rate_auto_check) + learning_rate_row.addWidget(self._learning_rate_edit) + learning_rate_row.addStretch(1) + params_layout.addRow( + self._make_form_label( + "Learning rate", + tooltip="Optimizer step size. Auto re-estimates it each iteration; otherwise enter a fixed decimal or scientific-notation value.", + ), + learning_rate_row, + ) + + self._iterations_spin = QSpinBox() + self._iterations_spin.setRange(1, 100_000) + self._iterations_spin.setSingleStep(100) + self._iterations_spin.setMaximumWidth(96) + params_layout.addRow( + self._make_form_label( + "Iterations", + tooltip="Maximum number of optimizer iterations.", + ), + self._iterations_spin, + ) + + self._advanced_group = QWidget() + advanced_group_layout = QVBoxLayout(self._advanced_group) + advanced_group_layout.setContentsMargins(0, 6, 0, 6) + advanced_group_layout.setSpacing(6) + + advanced_header = QWidget() + advanced_header_layout = QHBoxLayout(advanced_header) + advanced_header_layout.setContentsMargins(0, 0, 0, 0) + advanced_header_layout.setSpacing(6) + + self._advanced_toggle = QToolButton() + self._advanced_toggle.setCheckable(True) + self._advanced_toggle.setAutoRaise(True) + self._advanced_toggle.setToolButtonStyle( + Qt.ToolButtonStyle.ToolButtonTextBesideIcon + ) + self._advanced_toggle.setText("Advanced") + self._advanced_toggle.setArrowType(Qt.ArrowType.RightArrow) + advanced_header_layout.addWidget(self._advanced_toggle) + advanced_header_layout.addStretch(1) + advanced_group_layout.addWidget(advanced_header) + + self._advanced_content = QWidget() + advanced_layout = QFormLayout(self._advanced_content) + # Left indent + top gap set the advanced rows apart from the main + # parameters. Keep the right margin at 0: side margins add 1:1 to the + # panel's minimum width, and the left indent is only "free" because + # the group layout below carries no right margin either. + advanced_layout.setContentsMargins(9, 6, 0, 0) + advanced_layout.setSpacing(6) + advanced_layout.setRowWrapPolicy(QFormLayout.RowWrapPolicy.WrapLongRows) + advanced_layout.setFieldGrowthPolicy( + QFormLayout.FieldGrowthPolicy.ExpandingFieldsGrow + ) + + self._histogram_bins_spin = QSpinBox() + self._histogram_bins_spin.setSizePolicy( + QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Fixed + ) + self._histogram_bins_spin.setRange(8, 512) + self._histogram_bins_spin.setToolTip( + "Number of histogram bins for Mattes mutual information metric." + ) + self._histogram_bins_row = self._make_advanced_row( + advanced_layout, + "Histogram bins", + self._histogram_bins_spin, + tooltip="Number of histogram bins used by the Mattes mutual information metric.", + ) + + self._convergence_min_edit = ScientificDoubleSpinBox() + self._convergence_min_edit.setSizePolicy( + QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Fixed + ) + self._convergence_min_edit.setRange(1e-10, 1.0) + self._convergence_min_edit.setSingleStep(1e-6) + # Cap the width instead of setting a hard minimum: a minimum this wide + # would raise the whole panel's minimum width when the advanced + # section opens (issue #183 pattern). Expanding + max lets the field + # fill up to 260px when space allows and shrink on narrow docks. + self._convergence_min_edit.setMaximumWidth(260) + self._convergence_min_edit.setToolTip( + "Convergence threshold. Accepts decimal (0.000001) or scientific notation (1e-6)." + ) + self._convergence_min_row = self._make_advanced_row( + advanced_layout, + "Convergence min", + self._left_aligned_widget(self._convergence_min_edit), + tooltip="Convergence threshold below which the optimizer stops early.", + ) + + self._convergence_window_spin = QSpinBox() + self._convergence_window_spin.setSizePolicy( + QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Fixed + ) + self._convergence_window_spin.setRange(1, 100) + self._convergence_window_spin.setToolTip( + "Number of recent metric values for convergence estimation." + ) + self._convergence_window_row = self._make_advanced_row( + advanced_layout, + "Convergence window", + self._left_aligned_widget(self._convergence_window_spin), + tooltip="Number of recent metric values used to estimate convergence.", + ) + + self._multi_resolution_check = QCheckBox("Enabled") + self._multi_resolution_check.setToolTip( + "Run registration from coarse to fine resolution levels." + ) + self._multi_resolution_row = self._make_advanced_row( + advanced_layout, + "Multi-resolution", + self._multi_resolution_check, + tooltip="Whether to optimize from coarse to fine resolution levels.", + ) + + self._shrink_factors_edit = QLineEdit() + self._shrink_factors_edit.setToolTip( + "Comma-separated shrink factors per resolution level for multi-resolution." + ) + self._shrink_factors_row = self._make_advanced_row( + advanced_layout, + "Shrink factors", + self._shrink_factors_edit, + tooltip="Comma-separated downsampling factors for each multi-resolution level.", + ) + + self._smoothing_sigmas_edit = QLineEdit() + self._smoothing_sigmas_edit.setToolTip( + "Comma-separated smoothing sigmas per resolution level for multi-resolution." + ) + self._smoothing_sigmas_row = self._make_advanced_row( + advanced_layout, + "Smoothing sigmas", + self._smoothing_sigmas_edit, + tooltip="Comma-separated Gaussian smoothing sigmas applied at each multi-resolution level.", + ) + + self._interpolation_combo = QComboBox() + self._interpolation_combo.setMinimumContentsLength(8) + self._interpolation_combo.setSizeAdjustPolicy( + QComboBox.SizeAdjustPolicy.AdjustToMinimumContentsLengthWithIcon + ) + self._interpolation_combo.setSizePolicy( + QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Fixed + ) + self._interpolation_combo.addItems(["linear", "bspline"]) + self._interpolation_combo.setToolTip( + "Interpolator used for the resampled output." + ) + self._interpolation_row = self._make_advanced_row( + advanced_layout, + "Resample interp.", + self._interpolation_combo, + tooltip="Interpolator used when resampling the registered output onto the target grid.", + ) + + self._fill_value_auto_check = QCheckBox("minimum") + self._fill_value_auto_check.setToolTip( + "Automatically use the minimum intensity of the fixed image as fill value." + ) + self._fill_value_spin = QDoubleSpinBox() + # Ignored horizontal policy: the spinbox's minimum size hint spans its + # widest possible text ("-1000000.000"), which would force the whole + # panel wider than the dock on narrow layouts. + self._fill_value_spin.setSizePolicy( + QSizePolicy.Policy.Ignored, QSizePolicy.Policy.Fixed + ) + self._fill_value_spin.setRange(-1e6, 1e6) + self._fill_value_spin.setDecimals(3) + self._fill_value_spin.setEnabled(False) + self._fill_value_spin.setToolTip( + "Fill value for resampled output outside the input domain." + ) + self._fill_value_auto_check.toggled.connect( + lambda checked: self._fill_value_spin.setEnabled(not checked) + ) + self._multi_resolution_check.toggled.connect( + self._update_multi_resolution_enabled + ) + fill_value_row = QHBoxLayout() + fill_value_row.setContentsMargins(0, 0, 0, 0) + fill_value_row.addWidget(self._fill_value_auto_check) + fill_value_row.addWidget(self._fill_value_spin, stretch=1) + fill_value_container = QWidget() + fill_value_container.setLayout(fill_value_row) + self._fill_value_row = self._make_advanced_row( + advanced_layout, + "Fill value", + fill_value_container, + tooltip="Value written outside the moving image field of view after resampling. Choose 'minimum' to use the image minimum automatically.", + ) + + self._keep_diagnostics_check = QCheckBox("Keep full traces") + self._keep_diagnostics_check.setToolTip( + "Whether to keep the full per-frame optimizer traces for within-scan registration." + ) + self._keep_diagnostics_row = self._make_advanced_row( + advanced_layout, + "Diagnostics", + self._keep_diagnostics_check, + tooltip="Whether to store the full per-frame optimizer traces for within-scan registration.", + ) + + self._n_jobs_row = self._make_advanced_row( + advanced_layout, + "Parallel jobs", + self._n_jobs_spin, + tooltip="Number of parallel workers used for within-scan registration. -1 uses all CPUs.", + ) + + self._sitk_threads_spin = QSpinBox() + self._sitk_threads_spin.setRange(-128, 128) + self._sitk_threads_spin.setMaximumWidth(56) + self._sitk_threads_spin.setSpecialValueText("auto") + self._sitk_threads_spin.setToolTip( + "Number of SimpleITK threads used for between-scan registration. -1 uses all CPUs." + ) + self._sitk_threads_row = self._make_advanced_row( + advanced_layout, + "ITK Threads", + self._left_aligned_widget(self._sitk_threads_spin), + tooltip="Number of SimpleITK threads used for between-scan registration. -1 uses all CPUs.", + ) + + self._optimizer_weights_check = QCheckBox("Custom") + self._optimizer_weights_check.setToolTip( + "Apply per-parameter optimizer weights. 0 freezes a parameter; 1 leaves it unchanged. Not available for B-spline transforms." + ) + self._optimizer_weights_widget = QWidget() + optimizer_layout = QVBoxLayout(self._optimizer_weights_widget) + optimizer_layout.setContentsMargins(0, 0, 0, 0) + optimizer_layout.setSpacing(4) + self._optimizer_weights_fields = QWidget() + self._optimizer_weights_fields_layout = QGridLayout( + self._optimizer_weights_fields + ) + self._optimizer_weights_fields_layout.setContentsMargins(0, 0, 0, 0) + self._optimizer_weights_fields_layout.setHorizontalSpacing(8) + self._optimizer_weights_fields_layout.setVerticalSpacing(6) + optimizer_layout.addWidget(self._optimizer_weights_fields) + self._optimizer_weights_row = QWidget() + optimizer_row_layout = QVBoxLayout(self._optimizer_weights_row) + optimizer_row_layout.setContentsMargins(0, 0, 0, 0) + optimizer_row_layout.setSpacing(4) + optimizer_header_row = QHBoxLayout() + optimizer_header_row.setContentsMargins(0, 0, 0, 0) + optimizer_header_row.setSpacing(8) + self._optimizer_weights_label = self._make_form_label( + "Optimizer weights", + tooltip="Per-parameter weights multiplied into the optimizer step size. 0 freezes a parameter; 1 leaves it unchanged. Not available for B-spline transforms.", + ) + optimizer_header_row.addWidget(self._optimizer_weights_label) + optimizer_header_row.addWidget(self._optimizer_weights_check) + optimizer_header_row.addStretch(1) + optimizer_row_layout.addLayout(optimizer_header_row) + optimizer_row_layout.addWidget(self._optimizer_weights_fields) + advanced_layout.addRow(self._optimizer_weights_row) + + advanced_group_layout.addWidget(self._advanced_content) + self._advanced_toggle.toggled.connect(self._on_advanced_toggled) + self._metric_combo.currentTextChanged.connect( + self._update_metric_dependent_visibility + ) + self._transform_combo.currentTextChanged.connect( + self._update_transform_dependent_visibility + ) + self._transform_combo.currentTextChanged.connect( + lambda _text: self._sync_optimizer_weight_editor() + ) + self._optimizer_weights_check.toggled.connect( + lambda checked: self._optimizer_weights_fields.setVisible(checked) + ) + self._on_advanced_toggled(False) + + self._register_panel = QWidget() + register_layout = QVBoxLayout(self._register_panel) + register_layout.setContentsMargins(0, 0, 0, 0) + register_layout.setSpacing(8) + params_layout.addRow(self._advanced_group) + + register_layout.addWidget(operation_group) + register_layout.addWidget(params_group) + + btn_row = QHBoxLayout() + self._run_btn = QPushButton("Run registration") + self._run_btn.setObjectName("primary_btn") + self._run_btn.clicked.connect(self._run_registration) + btn_row.addWidget(self._run_btn) + + self._abort_btn = QPushButton("Abort") + self._abort_btn.setObjectName("danger_btn") + self._abort_btn.setToolTip("Abort the running registration.") + self._abort_btn.clicked.connect(self._abort_registration) + self._abort_btn.hide() + btn_row.addWidget(self._abort_btn) + + register_layout.addLayout(btn_row) + + layout.addWidget(self._register_panel) + + transforms_group = QGroupBox("Transforms") + transforms_group.setToolTip( + "Save, load, and reapply transforms from registration results or manual napari layer transforms." + ) + transforms_layout = QFormLayout(transforms_group) + transforms_layout.setSpacing(6) + transforms_layout.setRowWrapPolicy(QFormLayout.RowWrapPolicy.WrapLongRows) + transforms_layout.setFieldGrowthPolicy( + QFormLayout.FieldGrowthPolicy.ExpandingFieldsGrow + ) + + self._transform_source_combo = QComboBox() + self._transform_source_combo.setMinimumContentsLength(18) + self._transform_source_combo.setSizeAdjustPolicy( + QComboBox.SizeAdjustPolicy.AdjustToMinimumContentsLengthWithIcon + ) + self._transform_source_combo.setSizePolicy( + QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Fixed + ) + transforms_layout.addRow("Transform", self._transform_source_combo) + + self._transform_target_combo = QComboBox() + self._transform_target_combo.setMinimumContentsLength(18) + self._transform_target_combo.setSizeAdjustPolicy( + QComboBox.SizeAdjustPolicy.AdjustToMinimumContentsLengthWithIcon + ) + self._transform_target_combo.setSizePolicy( + QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Fixed + ) + self._transform_target_combo.setToolTip( + "Layer to resample when applying the selected transform." + ) + transforms_layout.addRow("Apply to", self._transform_target_combo) + + transform_buttons = QHBoxLayout() + self._save_transform_btn = QPushButton("Save") + self._save_transform_btn.clicked.connect(self._save_transform_callback) + self._load_transform_btn = QPushButton("Load") + self._load_transform_btn.clicked.connect(self._load_transform_callback) + self._apply_transform_btn = QPushButton("Forward") + self._apply_transform_btn.clicked.connect(self._apply_transform_callback) + self._apply_inverse_transform_btn = QPushButton("Inverse") + self._apply_inverse_transform_btn.clicked.connect( + self._apply_inverse_transform_callback + ) + transform_buttons.addWidget(self._save_transform_btn) + transform_buttons.addWidget(self._load_transform_btn) + transform_buttons.addWidget(self._apply_transform_btn) + transform_buttons.addWidget(self._apply_inverse_transform_btn) + transforms_layout.addRow(transform_buttons) + + self._transforms_panel = transforms_group + layout.addWidget(self._transforms_panel) + + self._status = QLabel("") + self._status.setWordWrap(True) + self._status.setObjectName("status_err") + self._status.hide() + layout.addWidget(self._status) + + self._progress = QProgressBar() + self._progress.setRange(0, 0) + self._progress.setMinimumHeight(16) + self._progress.setTextVisible(True) + self._progress.hide() + layout.addWidget(self._progress) + + layout.addStretch() + + self._register_panel_radio.toggled.connect(self._on_panel_changed) + self._transforms_panel_radio.toggled.connect(self._on_panel_changed) + self._single_volume_radio.toggled.connect(self._on_mode_changed) + self._time_series_radio.toggled.connect(self._on_mode_changed) + self._fixed_combo.currentTextChanged.connect( + self._validate_registration_selection + ) + self._initialization_combo.currentTextChanged.connect( + self._validate_registration_selection + ) + self._transform_source_combo.currentTextChanged.connect( + self._validate_registration_selection + ) + self._learning_rate_auto_check.toggled.connect( + lambda checked: self._learning_rate_edit.setEnabled(not checked) + ) + + self._registration_parameters_by_operation = { + "register_volume": get_default_registration_parameters(mode="volume"), + "register_volumewise": get_default_registration_parameters( + mode="volumewise" + ), + } + set_registration_parameters( + self, + self._registration_parameters_by_operation["register_volume"], + mode="volume", + ) + + self._refresh_layers() + self._on_panel_changed() + self._on_mode_changed() + + def _sync_manual_transform_event_connections(self) -> None: + """Keep manual-transform refresh hooks in sync with viewer layers.""" + for layer in self._manual_transform_event_layers: + try: + layer.events.affine.disconnect( + self._refresh_transform_controls_callback + ) + except (TypeError, RuntimeError): + pass + self._manual_transform_event_layers = [] + + for layer in self.viewer.layers: + if not _is_registration_source_layer(layer): + continue + _get_source_dataarray(layer) + layer.events.affine.connect(self._refresh_transform_controls_callback) + self._manual_transform_event_layers.append(layer) + + _refresh_layers = refresh_layers + _get_layer_by_name = get_layer_by_name + _selected_layer = selected_layer + _current_scale_mode = current_scale_mode + _current_metric = current_metric + _current_resample_interpolation = current_resample_interpolation + _current_transform_model = current_transform_model + + def _set_image_layer_data(self, layer: Image, data: npt.NDArray[Any]) -> None: + """Assign image data despite the current napari stub mismatch. + + Parameters + ---------- + layer : napari.layers.Image + Image layer whose data should be replaced. + data : numpy.ndarray + Replacement array. + + Returns + ------- + None + Updates `layer` in place. + """ + cast("Any", layer).data = data + + def _make_unique_layer_name(self, base_name: str) -> str: + """Return a viewer-unique layer name based on `base_name`. + + Parameters + ---------- + base_name : str + Desired layer name. + + Returns + ------- + str + Unique layer name for the current viewer. + """ + existing_names = {layer.name for layer in self.viewer.layers} + if base_name not in existing_names: + return base_name + index = 1 + while True: + candidate = f"{base_name} [{index}]" + if candidate not in existing_names: + return candidate + index += 1 + + def _make_unique_transform_name(self, base_name: str) -> str: + """Return a viewer-unique transform payload name based on `base_name`. + + Parameters + ---------- + base_name : str + Desired transform name. + + Returns + ------- + str + Unique transform payload name for the current viewer. + """ + existing_names = { + payload["name"] for payload in get_available_transform_payloads(self) + } + if base_name not in existing_names: + return base_name + index = 1 + while True: + candidate = f"{base_name} [{index}]" + if candidate not in existing_names: + return candidate + index += 1 + + def _volume_result_layer_name( + self, + moving_name: str, + fixed_name: str, + *, + transform_model: str | None = None, + ) -> str: + """Return the napari layer name for between-scan registration output. + + Parameters + ---------- + moving_name : str + Moving-layer name. Unused, but kept for call-site clarity. + fixed_name : str + Fixed-layer name. Unused, but kept for call-site clarity. + transform_model : str, optional + Transform model to include in the result-layer label. + + Returns + ------- + str + Result-layer name. + """ + del moving_name, fixed_name + model = transform_model or self._transform_combo.currentText() + return f"Registered ({model})" + + def _volume_fixed_preview_layer_name(self) -> str: + """Return the napari layer name for the fixed preview layer.""" + return "Fixed" + + def _volume_moving_preview_layer_name(self) -> str: + """Return the napari layer name for the moving preview layer.""" + return "Moving" + + def _volumewise_result_layer_name(self, moving_name: str) -> str: + """Return the napari layer name for within-scan registration output. + + Parameters + ---------- + moving_name : str + Moving-layer name. Unused, but kept for call-site symmetry. + + Returns + ------- + str + Result-layer name. + """ + del moving_name + return "Motion corrected" + + def _volumewise_moving_preview_layer_name(self) -> str: + """Return the napari layer name for the within-scan moving preview.""" + return "Moving" + + _update_reference_time_bounds = update_reference_time_bounds + _set_layer_validation_style = set_layer_validation_style + _set_run_btn_enabled = set_run_btn_enabled + _validate_registration_selection = validate_registration_selection + _on_moving_layer_changed = on_moving_layer_changed + + def _optimizer_weight_values(self) -> list[float]: + """Return the currently visible optimizer-weight values.""" + return [spin.value() for spin in self._optimizer_weight_spins] + + def _infer_registration_spatial_ndim(self) -> int | None: + """Return the spatial dimensionality of the current registration input.""" + moving_layer = self._selected_layer(self._moving_combo) + if moving_layer is None: + return None + try: + data = _get_source_dataarray(moving_layer) + except Exception: + return None + if self._operation() == "register_volume": + data = _prepare_between_scan_data(data) + return len([dim for dim in data.dims if dim != TIME_DIM]) + + def _sync_optimizer_weight_editor( + self, + *, + values: list[float] | None = None, + enabled: bool | None = None, + ) -> None: + """Rebuild the optimizer-weight editor for the current transform.""" + while self._optimizer_weights_fields_layout.count() > 0: + item = self._optimizer_weights_fields_layout.takeAt(0) + if item is None: + continue + widget = item.widget() + if widget is not None: + widget.deleteLater() + self._optimizer_weight_spins = [] + + transform = self._current_transform_model() + ndim = self._infer_registration_spatial_ndim() or 3 + labels = _get_optimizer_weight_labels(transform, ndim) + is_bspline = transform == "bspline" + if enabled is not None: + self._optimizer_weights_check.setChecked(enabled and not is_bspline) + elif is_bspline: + self._optimizer_weights_check.setChecked(False) + + self._optimizer_weights_fields.setVisible( + self._optimizer_weights_check.isChecked() and not is_bspline + ) + + def _make_weight_cell(label: str, value: float) -> QWidget: + cell = QWidget() + cell_layout = QVBoxLayout(cell) + cell_layout.setContentsMargins(0, 0, 0, 0) + cell_layout.setSpacing(2) + title = QLabel(label) + title.setAlignment(Qt.AlignmentFlag.AlignCenter) + spin = QDoubleSpinBox() + spin.setRange(0.0, 1e6) + spin.setDecimals(3) + spin.setSingleStep(0.1) + spin.setValue(value) + spin.setToolTip( + "Per-parameter optimizer weight. 0 freezes the parameter; 1 leaves it unchanged." + ) + cell_layout.addWidget(title) + cell_layout.addWidget(spin) + self._optimizer_weight_spins.append(spin) + return cell + + if transform == "affine": + matrix_labels = labels[: ndim * ndim] + translation_labels = labels[ndim * ndim :] + index = 0 + for row in range(ndim): + for col in range(ndim): + value = ( + values[index] + if values is not None and index < len(values) + else 1.0 + ) + self._optimizer_weights_fields_layout.addWidget( + _make_weight_cell(matrix_labels[index], value), row, col + ) + index += 1 + for col, label in enumerate(translation_labels): + value = ( + values[index] if values is not None and index < len(values) else 1.0 + ) + self._optimizer_weights_fields_layout.addWidget( + _make_weight_cell(label, value), ndim, col + ) + index += 1 + return + + if transform == "rigid": + split_index = 1 if ndim == 2 else 3 + angle_labels = labels[:split_index] + translation_labels = labels[split_index:] + index = 0 + for col, label in enumerate(angle_labels): + value = ( + values[index] if values is not None and index < len(values) else 1.0 + ) + self._optimizer_weights_fields_layout.addWidget( + _make_weight_cell(label, value), 0, col + ) + index += 1 + for col, label in enumerate(translation_labels): + value = ( + values[index] if values is not None and index < len(values) else 1.0 + ) + self._optimizer_weights_fields_layout.addWidget( + _make_weight_cell(label, value), 1, col + ) + index += 1 + return + + for index, label in enumerate(labels): + value = values[index] if values is not None and index < len(values) else 1.0 + self._optimizer_weights_fields_layout.addWidget( + _make_weight_cell(label, value), 0, index + ) + + def _spatial_info( + self, + ) -> tuple[ + tuple[int, ...] | None, tuple[float, ...] | None, tuple[float, ...] | None + ]: + """Return shape, scale, and translation for the first spatial image layer.""" + for layer in self.viewer.layers: + if layer._type_string != "image": + continue + da = layer.metadata.get("xarray") + if da is not None: + spatial_indices = [ + i for i, dim in enumerate(da.dims) if dim in SPATIAL_DIMS_WITH_POSE + ] + if not spatial_indices: + continue + shape = tuple(da.shape[i] for i in spatial_indices) + scale = tuple(float(layer.scale[i]) for i in spatial_indices) + translate = tuple(float(layer.translate[i]) for i in spatial_indices) + return shape, scale, translate + if layer.data.ndim >= 4: + return ( + layer.data.shape[1:], + tuple(float(s) for s in layer.scale[1:]), + tuple(float(t) for t in layer.translate[1:]), + ) + return None, None, None + + def _create_labels_layer(self, *, name: str = "Labels (3D)") -> None: + """Add a new Labels layer aligned to the current spatial image grid.""" + import numpy as np + + shape, scale, translate = self._spatial_info() + shape = shape or (64, 64, 64) + kwargs: dict[str, object] = {} + if scale is not None: + kwargs["scale"] = scale + if translate is not None: + kwargs["translate"] = translate + self.viewer.add_labels( + np.zeros(shape, dtype=np.int32), + name=name, + **kwargs, + ) + + def _operation(self) -> Literal["register_volume", "register_volumewise"]: + """Return the currently selected registration workflow.""" + if self._time_series_radio.isChecked(): + return "register_volumewise" + return "register_volume" + + def _on_panel_changed(self) -> None: + """Switch between the register and transforms subpanels.""" + show_register = self._register_panel_radio.isChecked() + self._register_panel.setVisible(show_register) + self._transforms_panel.setVisible(not show_register) + + def _on_advanced_toggled(self, checked: bool) -> None: + """Expand or collapse the advanced-parameter group.""" + self._advanced_content.setVisible(checked) + self._advanced_toggle.setArrowType( + Qt.ArrowType.DownArrow if checked else Qt.ArrowType.RightArrow + ) + + def _update_metric_dependent_visibility(self, metric: str) -> None: + """Show metric-specific advanced parameters for the current metric.""" + self._histogram_bins_row.setVisible(metric == "mattes_mi") + + def _update_multi_resolution_enabled(self, checked: bool) -> None: + """Show or hide multi-resolution-only parameter inputs.""" + self._shrink_factors_row.setVisible(checked) + self._smoothing_sigmas_row.setVisible(checked) + + def _update_transform_dependent_visibility(self, transform: str) -> None: + """Show or hide transform-specific parameter inputs.""" + is_bspline = transform == "bspline" + self._mesh_size_row.setVisible( + self._operation() == "register_volume" and is_bspline + ) + self._optimizer_weights_check.setEnabled(not is_bspline) + self._optimizer_weights_row.setVisible( + self._operation() == "register_volume" or not is_bspline + ) + if is_bspline: + self._optimizer_weights_fields.hide() + + def _on_mode_changed(self) -> None: + """Update the panel when the registration mode changes.""" + new_mode = self._operation() + previous_mode = self._active_operation + is_volumewise = new_mode == "register_volumewise" + + if previous_mode in self._registration_parameters_by_operation: + self._registration_parameters_by_operation[previous_mode] = ( + get_registration_parameters(self) + ) + + self._fixed_label.setVisible(not is_volumewise) + self._fixed_combo.setVisible(not is_volumewise) + self._fixed_combo.setEnabled(not is_volumewise) + self._fixed_mask_label.setVisible(not is_volumewise) + self._fixed_mask_row.setVisible(not is_volumewise) + self._moving_mask_label.setVisible(not is_volumewise) + self._moving_mask_row.setVisible(not is_volumewise) + self._reference_time_label.setVisible(is_volumewise) + self._reference_time_spin.setVisible(is_volumewise) + self._n_jobs_row.setVisible(is_volumewise) + self._sitk_threads_row.setVisible(not is_volumewise) + + self._learning_rate_auto_check.setVisible(not is_volumewise) + self._fill_value_row.setVisible(not is_volumewise) + self._keep_diagnostics_row.setVisible(is_volumewise) + + set_registration_parameters( + self, + self._registration_parameters_by_operation[new_mode], + mode="volumewise" if is_volumewise else "volume", + ) + self._active_operation = new_mode + + self._update_reference_time_bounds() + self._validate_registration_selection() + + def _begin_work(self) -> None: + """Put the panel into its busy state.""" + self._run_btn.hide() + self._abort_btn.setEnabled(True) + self._abort_btn.setText("Abort") + self._abort_btn.show() + self._status.hide() + if self._volumewise_progress_total is None: + self._progress.setRange(0, 0) + else: + self._progress.setRange(0, self._volumewise_progress_total) + self._progress.setValue(0) + self._progress.show() + QApplication.processEvents() + + def _abort_registration(self) -> None: + """Request cooperative cancellation of the running registration.""" + if self._worker is None or self._abort_event is None: + return + self._abort_event.set() + self._abort_btn.setEnabled(False) + self._abort_btn.setText("Aborting…") + self._set_error("Aborting registration…") + + def _end_work(self) -> None: + """Restore the idle UI state after background work.""" + self._worker = None + self._abort_event = None + self._run_btn.show() + self._run_btn.setEnabled(True) + self._run_btn.setText("Run registration") + self._abort_btn.setEnabled(True) + self._abort_btn.setText("Abort") + self._abort_btn.hide() + self._progress.setRange(0, 0) + self._progress.setValue(0) + self._progress.hide() + + def _set_error(self, message: str) -> None: + """Show a validation or execution error in the panel. + + Parameters + ---------- + message : str + Error message to display. + """ + self._status.setText(message) + self._status.show() + + def _run_registration(self) -> None: + """Validate inputs and start the selected registration workflow.""" + operation = self._operation() + moving_layer = self._selected_layer(self._moving_combo) + fixed_layer = self._selected_layer(self._fixed_combo) + + if moving_layer is None: + self._set_error("Select a moving layer.") + return + if not self._validate_registration_selection(): + return + + try: + learning_rate: float | Literal["auto"] + if self._learning_rate_auto_check.isChecked(): + learning_rate = "auto" + else: + learning_rate = self._learning_rate_edit.value() + moving = _get_source_dataarray(moving_layer) + except Exception as exc: # noqa: BLE001 + self._set_error(str(exc)) + return + + convergence_minimum_value = self._convergence_min_edit.value() + + # Parse advanced parameters + shrink_factors = _parse_comma_separated_ints(self._shrink_factors_edit.text()) + smoothing_sigmas = _parse_comma_separated_ints( + self._smoothing_sigmas_edit.text() + ) + use_multi_res = self._multi_resolution_check.isChecked() + if not use_multi_res: + shrink_factors = None + smoothing_sigmas = None + + metric = self._current_metric() + scale_mode = self._current_scale_mode() + resample_interpolation = self._current_resample_interpolation() + transform = self._current_transform_model() + initialization = cast( + "InitializationSelection", self._initialization_combo.currentData() + ) + optimizer_weights = ( + self._optimizer_weight_values() + if self._optimizer_weights_check.isChecked() and transform != "bspline" + else None + ) + self._abort_event = Event() + + if operation == "register_volume": + if fixed_layer is None: + self._set_error("Select a fixed layer.") + return + if fixed_layer is moving_layer: + self._set_error("Moving and fixed layers must be different.") + return + + try: + fixed = _get_source_dataarray(fixed_layer) + except Exception as exc: # noqa: BLE001 + self._set_error(str(exc)) + return + + fixed_mask = None + fixed_mask_layer_name = None + fixed_mask_layer = self._selected_layer(self._fixed_mask_combo) + if fixed_mask_layer is not None: + fixed_mask = _get_source_dataarray(fixed_mask_layer) > 0 + fixed_mask_layer_name = fixed_mask_layer.name + + moving_mask = None + moving_mask_layer_name = None + moving_mask_layer = self._selected_layer(self._moving_mask_combo) + if moving_mask_layer is not None: + moving_mask = _get_source_dataarray(moving_mask_layer) > 0 + moving_mask_layer_name = moving_mask_layer.name + + moving = _prepare_between_scan_data(moving) + fixed = _prepare_between_scan_data(fixed) + moving = _apply_registration_scale(moving, scale_mode) + fixed = _apply_registration_scale(fixed, scale_mode) + + initial_transform: npt.NDArray[np.floating] | None = None + try: + initial_transform, initial_transform_source = ( + get_selected_initial_transform( + self, + moving, + moving_layer=moving_layer, + fixed_layer=fixed_layer, + ) + ) + except Exception as exc: # noqa: BLE001 + self._set_error(str(exc)) + return + + if transform not in {"translation", "rigid", "affine", "bspline"}: + self._set_error(f"Unknown transform model: {transform!r}.") + return + volume_payload: VolumeRegistrationRunPayload = { + "operation": "register_volume", + "moving_layer_name": moving_layer.name, + "transform": transform, + "metric": metric, + "scale": scale_mode, + "learning_rate": learning_rate, + "number_of_iterations": self._iterations_spin.value(), + "use_multi_resolution": use_multi_res, + "resample_interpolation": resample_interpolation, + "number_of_histogram_bins": self._histogram_bins_spin.value(), + "convergence_minimum_value": convergence_minimum_value, + "convergence_window_size": self._convergence_window_spin.value(), + "initialization": initialization, + "shrink_factors": shrink_factors, + "smoothing_sigmas": smoothing_sigmas, + "optimizer_weights": optimizer_weights, + "keep_diagnostics": self._keep_diagnostics_check.isChecked(), + "fill_value": None + if self._fill_value_auto_check.isChecked() + else self._fill_value_spin.value(), + "mesh_size": ( + self._mesh_size_z_spin.value(), + self._mesh_size_y_spin.value(), + self._mesh_size_x_spin.value(), + ), + "fixed_layer_name": fixed_layer.name, + "fixed_mask_layer_name": fixed_mask_layer_name, + "moving_mask_layer_name": moving_mask_layer_name, + "sitk_threads": self._sitk_threads_spin.value(), + } + if initial_transform_source is not None: + volume_payload["initial_transform_source"] = initial_transform_source + + initialization_arg = ( + initial_transform + if initial_transform is not None + else get_selected_center_initialization(self) + ) + + try: + progress_plotter = create_volume_progress_plotter( + self, + moving_layer=cast("Image", moving_layer), + fixed_layer=cast("Image", fixed_layer), + moving=moving, + fixed=fixed, + layer_name=self._make_unique_layer_name( + self._volume_result_layer_name( + volume_payload["moving_layer_name"], + volume_payload["fixed_layer_name"], + transform_model=volume_payload["transform"], + ) + ), + initial_transform=initial_transform, + scale_mode=volume_payload["scale"], + ) + except Exception: + return + + worker = thread_worker(register_volume)( + moving, + fixed, + fixed_mask=fixed_mask, + moving_mask=moving_mask, + transform_type=volume_payload["transform"], + metric=volume_payload["metric"], + learning_rate=learning_rate, + number_of_iterations=volume_payload["number_of_iterations"], + use_multi_resolution=volume_payload["use_multi_resolution"], + resample=True, + resample_interpolation=volume_payload["resample_interpolation"], + mesh_size=volume_payload["mesh_size"], + number_of_histogram_bins=volume_payload["number_of_histogram_bins"], + convergence_minimum_value=volume_payload["convergence_minimum_value"], + convergence_window_size=volume_payload["convergence_window_size"], + initialization=initialization_arg, + optimizer_weights=volume_payload["optimizer_weights"], + shrink_factors=volume_payload["shrink_factors"] or (6, 2, 1), + smoothing_sigmas=volume_payload["smoothing_sigmas"] or (6, 2, 1), + fill_value=volume_payload["fill_value"], + sitk_threads=volume_payload["sitk_threads"], + show_progress=True, + progress_plotter=progress_plotter, + abort_event=self._abort_event, + ) + self._worker = worker + self._begin_work() + worker.returned.connect( + lambda result: on_volume_registration_finished( + self, volume_payload, result + ) + ) + else: + if TIME_DIM not in moving.dims: + self._set_error( + "register_volumewise requires a layer with a time dimension." + ) + return + if transform == "bspline": + self._set_error(f"Unknown transform model: {transform!r}.") + return + + volumewise_payload: VolumewiseRegistrationRunPayload = { + "operation": "register_volumewise", + "moving_layer_name": moving_layer.name, + "transform": transform, + "metric": metric, + "scale": scale_mode, + "learning_rate": learning_rate, + "number_of_iterations": self._iterations_spin.value(), + "use_multi_resolution": use_multi_res, + "resample_interpolation": resample_interpolation, + "number_of_histogram_bins": self._histogram_bins_spin.value(), + "convergence_minimum_value": convergence_minimum_value, + "convergence_window_size": self._convergence_window_spin.value(), + "initialization": initialization, + "shrink_factors": shrink_factors, + "smoothing_sigmas": smoothing_sigmas, + "optimizer_weights": optimizer_weights, + "keep_diagnostics": self._keep_diagnostics_check.isChecked(), + "fill_value": None + if self._fill_value_auto_check.isChecked() + else self._fill_value_spin.value(), + "mesh_size": ( + self._mesh_size_z_spin.value(), + self._mesh_size_y_spin.value(), + self._mesh_size_x_spin.value(), + ), + "reference_time": self._reference_time_spin.value(), + "n_jobs": self._n_jobs_spin.value(), + } + moving = _apply_registration_scale(moving, volumewise_payload["scale"]) + + progress_reporter = setup_volumewise_progress( + self, + moving_layer=cast("Image", moving_layer), + moving=moving, + layer_name=self._make_unique_layer_name( + self._volumewise_result_layer_name( + volumewise_payload["moving_layer_name"] + ) + ), + scale_mode=volumewise_payload["scale"], + ) + + worker = thread_worker(register_volumewise)( + moving, + reference_time=volumewise_payload["reference_time"], + n_jobs=volumewise_payload["n_jobs"], + transform=volumewise_payload["transform"], + metric=volumewise_payload["metric"], + learning_rate=learning_rate, + number_of_iterations=volumewise_payload["number_of_iterations"], + use_multi_resolution=volumewise_payload["use_multi_resolution"], + resample_interpolation=volumewise_payload["resample_interpolation"], + number_of_histogram_bins=volumewise_payload["number_of_histogram_bins"], + convergence_minimum_value=volumewise_payload[ + "convergence_minimum_value" + ], + convergence_window_size=volumewise_payload["convergence_window_size"], + initialization=get_selected_center_initialization(self), + optimizer_weights=volumewise_payload["optimizer_weights"], + shrink_factors=volumewise_payload["shrink_factors"] or (6, 2, 1), + smoothing_sigmas=volumewise_payload["smoothing_sigmas"] or (6, 2, 1), + keep_diagnostics=volumewise_payload["keep_diagnostics"], + show_progress=False, + abort_event=self._abort_event, + progress_reporter=progress_reporter, + ) + self._worker = worker + self._begin_work() + worker.returned.connect( + lambda result: on_volumewise_registration_finished( + self, volumewise_payload, result + ) + ) + worker.errored.connect(lambda exc: on_registration_failed(self, exc)) + worker.finished.connect(self._end_work) + worker.start() diff --git a/src/confusius/_napari/_registration/_panel_parameters.py b/src/confusius/_napari/_registration/_panel_parameters.py new file mode 100644 index 00000000..0270ad62 --- /dev/null +++ b/src/confusius/_napari/_registration/_panel_parameters.py @@ -0,0 +1,176 @@ +"""Registration-parameter helpers for the napari registration panel.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from confusius._napari._registration._panel import ( + ModeParameters, + RegistrationPanel, + RegistrationParameterMode, + ) + + +def get_default_registration_parameters( + *, mode: "RegistrationParameterMode" +) -> "ModeParameters": + """Return the default parameter state for one registration mode. + + Parameters + ---------- + mode : {"volume", "volumewise"} + Registration workflow whose defaults should be returned. + + Returns + ------- + ModeParameters + Default parameter values for the requested workflow. + """ + is_volumewise = mode == "volumewise" + return { + "transform": "rigid", + "metric": "correlation", + "scale": "dB", + "initialization": "center_geometry", + "learning_rate_auto": not is_volumewise, + "learning_rate_value": 0.01 if is_volumewise else 0.1, + "number_of_iterations": 100, + "number_of_histogram_bins": 50, + "mesh_size": (10, 10, 10), + "convergence_minimum_value": 1e-6, + "convergence_window_size": 10, + "use_multi_resolution": False, + "shrink_factors": "6, 2, 1", + "smoothing_sigmas": "6, 2, 1", + "resample_interpolation": "linear", + "fill_value_auto": True, + "fill_value": 0.0, + "reference_time": 0, + "n_jobs": -1, + "sitk_threads": -1, + "optimizer_weights_enabled": False, + "optimizer_weights_values": [], + "keep_diagnostics": False, + "advanced_open": False, + } + + +def get_registration_parameters(panel: "RegistrationPanel") -> "ModeParameters": + """Return the current parameter state shown in the panel. + + Parameters + ---------- + panel : RegistrationPanel + Panel whose widgets should be read. + + Returns + ------- + ModeParameters + Current parameter values read from the visible widgets. + """ + return { + "transform": panel._transform_combo.currentText() or "rigid", + "metric": panel._current_metric(), + "scale": panel._current_scale_mode(), + "initialization": panel._initialization_combo.currentData(), + "learning_rate_auto": panel._learning_rate_auto_check.isChecked(), + "learning_rate_value": panel._learning_rate_edit.value(), + "number_of_iterations": panel._iterations_spin.value(), + "number_of_histogram_bins": panel._histogram_bins_spin.value(), + "mesh_size": ( + panel._mesh_size_z_spin.value(), + panel._mesh_size_y_spin.value(), + panel._mesh_size_x_spin.value(), + ), + "convergence_minimum_value": panel._convergence_min_edit.value(), + "convergence_window_size": panel._convergence_window_spin.value(), + "use_multi_resolution": panel._multi_resolution_check.isChecked(), + "shrink_factors": panel._shrink_factors_edit.text(), + "smoothing_sigmas": panel._smoothing_sigmas_edit.text(), + "resample_interpolation": panel._current_resample_interpolation(), + "fill_value_auto": panel._fill_value_auto_check.isChecked(), + "fill_value": panel._fill_value_spin.value(), + "reference_time": panel._reference_time_spin.value(), + "n_jobs": panel._n_jobs_spin.value(), + "sitk_threads": panel._sitk_threads_spin.value(), + "optimizer_weights_enabled": panel._optimizer_weights_check.isChecked(), + "optimizer_weights_values": panel._optimizer_weight_values(), + "keep_diagnostics": panel._keep_diagnostics_check.isChecked(), + "advanced_open": panel._advanced_toggle.isChecked(), + } + + +def set_registration_parameters( + panel: "RegistrationPanel", + params: "ModeParameters", + *, + mode: "RegistrationParameterMode", +) -> None: + """Restore the parameter state for one registration mode. + + Parameters + ---------- + panel : RegistrationPanel + Panel whose widgets should be updated. + params : ModeParameters + Parameter values to push back into the widgets. + mode : {"volume", "volumewise"} + Registration workflow whose UI should be restored. + """ + panel._transform_combo.blockSignals(True) + panel._transform_combo.clear() + is_volumewise = mode == "volumewise" + if is_volumewise: + panel._transform_combo.addItems(["translation", "rigid", "affine"]) + else: + panel._transform_combo.addItems(["translation", "rigid", "affine", "bspline"]) + transform = params["transform"] + transform_index = panel._transform_combo.findText(transform) + if transform_index < 0: + transform_index = panel._transform_combo.findText("rigid") + if transform_index >= 0: + panel._transform_combo.setCurrentIndex(transform_index) + panel._transform_combo.blockSignals(False) + + panel._metric_combo.setCurrentText(params["metric"]) + scale_mode = params["scale"] + scale_index = panel._scale_combo.findData(scale_mode) + if scale_index >= 0: + panel._scale_combo.setCurrentIndex(scale_index) + initialization_data = params.get("initialization") + for i in range(panel._initialization_combo.count()): + if panel._initialization_combo.itemData(i) == initialization_data: + panel._initialization_combo.setCurrentIndex(i) + break + panel._learning_rate_auto_check.setChecked( + False if is_volumewise else params["learning_rate_auto"] + ) + panel._learning_rate_edit.setValue(params["learning_rate_value"]) + panel._iterations_spin.setValue(params["number_of_iterations"]) + panel._histogram_bins_spin.setValue(params["number_of_histogram_bins"]) + mesh_size = params["mesh_size"] + panel._mesh_size_z_spin.setValue(mesh_size[0]) + panel._mesh_size_y_spin.setValue(mesh_size[1]) + panel._mesh_size_x_spin.setValue(mesh_size[2]) + panel._convergence_min_edit.setValue(params["convergence_minimum_value"]) + panel._convergence_window_spin.setValue(params["convergence_window_size"]) + panel._multi_resolution_check.setChecked(params["use_multi_resolution"]) + panel._shrink_factors_edit.setText(params["shrink_factors"]) + panel._smoothing_sigmas_edit.setText(params["smoothing_sigmas"]) + panel._interpolation_combo.setCurrentText(params["resample_interpolation"]) + panel._fill_value_auto_check.setChecked(params["fill_value_auto"]) + panel._fill_value_spin.setValue(params["fill_value"]) + panel._reference_time_spin.setValue(params["reference_time"]) + panel._n_jobs_spin.setValue(params["n_jobs"]) + panel._sitk_threads_spin.setValue(params["sitk_threads"]) + panel._keep_diagnostics_check.setChecked(params["keep_diagnostics"]) + panel._advanced_toggle.setChecked(params["advanced_open"]) + panel._on_advanced_toggled(panel._advanced_toggle.isChecked()) + panel._update_metric_dependent_visibility(panel._metric_combo.currentText()) + panel._update_multi_resolution_enabled(panel._multi_resolution_check.isChecked()) + panel._update_transform_dependent_visibility(panel._transform_combo.currentText()) + panel._sync_optimizer_weight_editor( + values=params.get("optimizer_weights_values"), + enabled=params.get("optimizer_weights_enabled", False), + ) diff --git a/src/confusius/_napari/_registration/_panel_progress.py b/src/confusius/_napari/_registration/_panel_progress.py new file mode 100644 index 00000000..7f853c94 --- /dev/null +++ b/src/confusius/_napari/_registration/_panel_progress.py @@ -0,0 +1,505 @@ +"""Progress-layer helpers for the napari registration panel.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, cast + +import numpy as np +import xarray as xr +from napari.layers.utils.layer_utils import calc_data_range +from qtpy.QtCore import Qt, QTimer +from qtpy.QtWidgets import QDockWidget, QWidget + +from confusius._dims import TIME_DIM +from confusius._napari._qt import find_main_window +from confusius._napari._registration._panel_utils import ( + _gamma_needs_reset, + _get_image_display_kwargs_from_layer, + _preserve_view, +) +from confusius._napari._registration._progress import ( + NapariRegistrationProgressPlotterBridge, + NapariRegistrationProgressReporter, + NapariRegistrationProgressReporterBridge, + make_napari_progress_factory, +) +from confusius.plotting.napari import plot_napari +from confusius.registration import resample_like + +if TYPE_CHECKING: + from collections.abc import Callable + + import numpy.typing as npt + from napari.layers import Image + + from confusius._napari._registration._metric_plotter import ( + RegistrationMetricPlotter, + ) + from confusius._napari._registration._panel import RegistrationPanel + from confusius.registration import RegistrationProgress + + +def setup_volumewise_progress( + panel: "RegistrationPanel", + *, + moving_layer: "Image", + moving: xr.DataArray, + layer_name: str, + scale_mode: str, +) -> NapariRegistrationProgressReporter: + """Create volumewise preview layers and a progress reporter. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose viewer and progress widgets are updated. + moving_layer : napari.layers.Image + Source moving layer shown in the viewer. + moving : xarray.DataArray + Moving data used to seed the preview layer. + layer_name : str + Name for the live output layer. + scale_mode : str + Registration scaling mode used to decide preview gamma handling. + + Returns + ------- + NapariRegistrationProgressReporter + Worker-side reporter that forwards completed-frame updates back to the + panel. + """ + teardown_volumewise_progress(panel, remove_layer=True) + + moving_display_kwargs = _get_image_display_kwargs_from_layer(moving_layer) + moving_display_kwargs["colormap"] = "red" + if _gamma_needs_reset(scale_mode): + moving_display_kwargs["gamma"] = 1.0 + + display_kwargs = dict(moving_display_kwargs) + display_kwargs["colormap"] = "cyan" + display_kwargs["blending"] = "additive" + contrast_limits = tuple(calc_data_range(moving.data)) + preview_data = np.full( + moving.shape, + fill_value=float(np.min(moving.data)), + dtype=np.asarray(moving.data).dtype, + ) + preview = xr.DataArray( + preview_data, + dims=moving.dims, + coords=moving.coords, + attrs=moving.attrs.copy(), + ) + + with _preserve_view(panel.viewer): + try: + moving_preview_layer = panel._get_layer_by_name( + panel._volumewise_moving_preview_layer_name() + ) + if moving_preview_layer is None: + _, moving_preview_layer = plot_napari( + moving, + viewer=panel.viewer, + name=panel._volumewise_moving_preview_layer_name(), + show_colorbar=False, + contrast_limits=contrast_limits, + **moving_display_kwargs, + ) + else: + moving_preview_layer = cast("Image", moving_preview_layer) + panel._set_image_layer_data( + moving_preview_layer, np.asarray(moving.data) + ) + moving_preview_layer.colormap = moving_display_kwargs["colormap"] + moving_preview_layer.gamma = float( + moving_display_kwargs.get("gamma", 1.0) + ) + moving_preview_layer.contrast_limits = contrast_limits + + fixed_preview_layer = panel._get_layer_by_name( + panel._volume_fixed_preview_layer_name() + ) + if fixed_preview_layer is not None: + fixed_preview_layer.visible = False + + _, layer = plot_napari( + preview, + viewer=panel.viewer, + name=layer_name, + show_colorbar=False, + contrast_limits=contrast_limits, + **display_kwargs, + ) + except Exception as exc: # noqa: BLE001 + panel._set_error(f"Could not create progress layer: {exc}") + raise + bridge = NapariRegistrationProgressReporterBridge() + bridge.frame_progress.connect( + lambda completed_frames, total_frames: update_volumewise_progress_bar( + panel, completed_frames, total_frames + ) + ) + bridge.frame_completed.connect( + lambda frame_index, frame_data: update_volumewise_progress_frame( + panel, frame_index, frame_data + ) + ) + + panel._volumewise_progress_bridge = bridge + panel._volumewise_progress_layer = cast("Image", layer) + panel._volumewise_moving_preview_layer = cast("Image", moving_preview_layer) + panel._volumewise_progress_time_axis = moving.dims.index(TIME_DIM) + panel._volumewise_progress_total = moving.sizes[TIME_DIM] + panel._progress.setRange(0, panel._volumewise_progress_total) + panel._progress.setValue(0) + return NapariRegistrationProgressReporter( + bridge, + n_frames=panel._volumewise_progress_total, + ) + + +def update_volumewise_progress_bar( + panel: "RegistrationPanel", completed_frames: int, total_frames: int +) -> None: + """Update the volumewise progress bar from completed-frame counts. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose progress bar is updated. + completed_frames : int + Number of frames completed so far. + total_frames : int + Total number of frames expected for the run. + + Returns + ------- + None + Updates the panel progress bar in place. + """ + panel._volumewise_progress_total = total_frames + panel._progress.setRange(0, total_frames) + panel._progress.setValue(min(completed_frames, total_frames)) + + +def update_volumewise_progress_frame( + panel: "RegistrationPanel", frame_index: int, frame_data: object +) -> None: + """Write a completed frame into the volumewise preview layer. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose volumewise preview layer is updated. + frame_index : int + Time index of the completed frame. + frame_data : object + Registered frame data emitted by the worker. + + Returns + ------- + None + Writes the completed frame into the preview layer when valid. + """ + layer = panel._volumewise_progress_layer + time_axis = panel._volumewise_progress_time_axis + if layer is None or time_axis is None: + return + if not isinstance(frame_data, np.ndarray): + return + if frame_index < 0 or frame_index >= layer.data.shape[time_axis]: + return + slicer: list[int | slice] = [slice(None) for _ in range(layer.data.ndim)] + slicer[time_axis] = frame_index + np.asarray(layer.data)[tuple(slicer)] = frame_data + layer.refresh() + + +def teardown_volumewise_progress( + panel: "RegistrationPanel", *, remove_layer: bool +) -> None: + """Drop volumewise progress-layer references and optionally remove the layer. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose volumewise progress state is cleared. + remove_layer : bool + Whether to also remove the live output layer from the viewer. + + Returns + ------- + None + Clears volumewise progress state stored on the panel. + """ + if remove_layer and panel._volumewise_progress_layer is not None: + try: + panel.viewer.layers.remove(panel._volumewise_progress_layer) + except (KeyError, ValueError): + pass + panel._volumewise_progress_bridge = None + panel._volumewise_progress_layer = None + panel._volumewise_progress_time_axis = None + panel._volumewise_progress_total = None + + +def create_volume_progress_plotter( + panel: "RegistrationPanel", + *, + moving_layer: "Image", + fixed_layer: "Image", + moving: xr.DataArray, + fixed: xr.DataArray, + layer_name: str, + initial_transform: "npt.NDArray[np.floating] | None" = None, + scale_mode: str, +) -> "Callable[..., RegistrationProgress]": + """Create between-scan preview layers and a progress-plotter factory. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose viewer and progress widgets are updated. + moving_layer : napari.layers.Image + Source moving layer shown in the viewer. + fixed_layer : napari.layers.Image + Source fixed layer shown in the viewer. + moving : xarray.DataArray + Moving data used to seed the live preview. + fixed : xarray.DataArray + Fixed data defining the output grid. + layer_name : str + Name for the live registered preview layer. + initial_transform : numpy.ndarray, optional + Initial affine used to seed the preview before optimization starts. + scale_mode : str + Registration scaling mode used to decide preview gamma handling. + + Returns + ------- + callable + Progress-plotter factory for `register_volume`. + + Raises + ------ + Exception + If the preview layer could not be created. + """ + teardown_volume_progress(panel) + + fixed_display_kwargs = _get_image_display_kwargs_from_layer(fixed_layer) + fixed_display_kwargs["colormap"] = "red" + + moving_display_kwargs = _get_image_display_kwargs_from_layer(moving_layer) + moving_display_kwargs["colormap"] = "cyan" + moving_display_kwargs["blending"] = "additive" + if _gamma_needs_reset(scale_mode): + fixed_display_kwargs["gamma"] = 1.0 + moving_display_kwargs["gamma"] = 1.0 + + display_kwargs = dict(moving_display_kwargs) + display_kwargs["colormap"] = "cyan" + display_kwargs["blending"] = "additive" + + try: + seed_transform = ( + np.asarray(initial_transform, dtype=float) + if initial_transform is not None + else np.eye(fixed.ndim + 1, dtype=float) + ) + preview = resample_like( + moving, + fixed, + seed_transform, + interpolation="linear", + ) + preview_contrast_limits = tuple(calc_data_range(preview.data)) + except Exception as exc: # noqa: BLE001 + panel._set_error(f"Could not seed progress layer: {exc}") + preview = xr.DataArray( + np.zeros(fixed.shape, dtype=np.float32), + coords=fixed.coords, + dims=fixed.dims, + attrs=fixed.attrs.copy(), + ) + preview_contrast_limits = tuple(calc_data_range(preview.data)) + + with _preserve_view(panel.viewer): + try: + fixed_preview_layer = panel._get_layer_by_name( + panel._volume_fixed_preview_layer_name() + ) + if fixed_preview_layer is None: + _, fixed_preview_layer = plot_napari( + fixed, + viewer=panel.viewer, + name=panel._volume_fixed_preview_layer_name(), + show_colorbar=False, + **fixed_display_kwargs, + ) + else: + fixed_preview_layer = cast("Image", fixed_preview_layer) + panel._set_image_layer_data(fixed_preview_layer, np.asarray(fixed.data)) + fixed_preview_layer.colormap = fixed_display_kwargs["colormap"] + fixed_preview_layer.gamma = float( + fixed_display_kwargs.get("gamma", 1.0) + ) + fixed_preview_layer.visible = True + + moving_preview_layer = panel._get_layer_by_name( + panel._volume_moving_preview_layer_name() + ) + if moving_preview_layer is None: + _, moving_preview_layer = plot_napari( + preview, + viewer=panel.viewer, + name=panel._volume_moving_preview_layer_name(), + show_colorbar=False, + contrast_limits=preview_contrast_limits, + **moving_display_kwargs, + ) + else: + moving_preview_layer = cast("Image", moving_preview_layer) + panel._set_image_layer_data( + moving_preview_layer, np.asarray(preview.data) + ) + moving_preview_layer.colormap = moving_display_kwargs["colormap"] + moving_preview_layer.blending = moving_display_kwargs["blending"] + moving_preview_layer.gamma = float( + moving_display_kwargs.get("gamma", 1.0) + ) + moving_preview_layer.contrast_limits = preview_contrast_limits + moving_preview_layer.visible = False + + _, layer = plot_napari( + preview, + viewer=panel.viewer, + name=layer_name, + show_colorbar=False, + contrast_limits=preview_contrast_limits, + **display_kwargs, + ) + except Exception as exc: # noqa: BLE001 + panel._set_error(f"Could not create progress layer: {exc}") + raise + + bridge = NapariRegistrationProgressPlotterBridge() + bridge.iterated.connect(lambda arr: update_progress_layer(panel, arr)) + panel._progress_bridge = bridge + panel._progress_layer = cast("Image", layer) + panel._progress_fixed_layer = cast("Image", fixed_preview_layer) + panel._progress_moving_layer = cast("Image", moving_preview_layer) + panel._progress_moving_layer.visible = False + + ensure_metric_plotter(panel) + plotter = panel._metric_plotter + if plotter is not None: + plotter.reset() + bridge.metric_updated.connect(plotter.add_metric) + return make_napari_progress_factory(bridge) + + +def update_progress_layer(panel: "RegistrationPanel", arr: object) -> None: + """Write an intermediate resampled array into the volume preview layer. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose between-scan preview layer is updated. + arr : object + Intermediate resampled array emitted by the worker. + + Returns + ------- + None + Replaces the preview-layer data when the emitted array is valid. + """ + layer = panel._progress_layer + if layer is None: + return + if not isinstance(arr, np.ndarray): + return + if arr.shape != layer.data.shape: + return + panel._set_image_layer_data(layer, arr) + + +def teardown_volume_progress(panel: "RegistrationPanel") -> None: + """Remove the volume progress preview layer and bridge references. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose between-scan progress state is cleared. + + Returns + ------- + None + Clears between-scan progress state stored on the panel. + """ + if panel._progress_layer is not None: + try: + panel.viewer.layers.remove(panel._progress_layer) + except (KeyError, ValueError): + pass + panel._progress_layer = None + panel._progress_bridge = None + + +def ensure_metric_plotter( + panel: "RegistrationPanel", +) -> "RegistrationMetricPlotter | None": + """Return the right-dock metric plotter, creating it on first use. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose metric plotter should be available. + + Returns + ------- + RegistrationMetricPlotter or None + The docked metric plotter widget. + """ + if panel._metric_plotter is None: + from confusius._napari._registration._metric_plotter import ( + RegistrationMetricPlotter, + ) + + panel._metric_plotter = RegistrationMetricPlotter(panel.viewer) + + if panel._metric_dock is None or panel._metric_plotter.parent() is None: + dock = panel.viewer.window.add_dock_widget( + panel._metric_plotter, + name="Registration Metric", + area="right", + ) + panel._metric_dock = dock + + def _settle_layout() -> None: + main_win = find_main_window(dock) + if main_win is None: + return + from qtpy.QtCore import QSize + + central = main_win.centralWidget() + if central is None: + return + central.setMinimumSize(QSize(0, 0)) + for widget in central.findChildren(QWidget): + widget.setMinimumSize(QSize(0, 0)) + for side_dock in main_win.findChildren(QDockWidget): + if side_dock is dock: + continue + side_dock.setMinimumHeight(0) + widget = side_dock.widget() + if widget is not None: + widget.setMinimumSize(QSize(0, 0)) + current = main_win.size() + if current.height() < 800: + main_win.resize(current.width(), 800) + main_win.resizeDocks([dock], [220], Qt.Orientation.Vertical) + + QTimer.singleShot(200, _settle_layout) + + return panel._metric_plotter diff --git a/src/confusius/_napari/_registration/_panel_results.py b/src/confusius/_napari/_registration/_panel_results.py new file mode 100644 index 00000000..7d010d99 --- /dev/null +++ b/src/confusius/_napari/_registration/_panel_results.py @@ -0,0 +1,330 @@ +"""Result-handling helpers for the napari registration panel.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Literal, cast + +import numpy as np +import xarray as xr +from napari.layers.utils.layer_utils import calc_data_range +from napari.utils.notifications import show_info + +from confusius._napari._registration._panel_progress import teardown_volumewise_progress +from confusius._napari._registration._panel_transforms import refresh_transform_controls +from confusius._napari._registration._transform_payloads import ( + make_affine_transform_payload, + make_bspline_transform_payload, +) +from confusius._napari._registration._panel_utils import ( + _gamma_needs_reset, + _get_image_display_kwargs_from_layer, + _get_source_dataarray, + _prepare_between_scan_data, +) +from confusius.plotting.napari import plot_napari + +if TYPE_CHECKING: + import numpy.typing as npt + + from confusius._napari._registration._panel import ( + RegistrationPanel, + VolumeRegistrationRunPayload, + VolumewiseRegistrationRunPayload, + ) + from confusius.registration import RegistrationDiagnostics + + +def coerce_volume_registration_payload( + payload: dict[str, Any] | "VolumeRegistrationRunPayload", +) -> "VolumeRegistrationRunPayload": + """Return a typed between-scan registration payload. + + Parameters + ---------- + payload : dict[str, Any] or VolumeRegistrationRunPayload + Untyped or typed payload captured when the worker started. + + Returns + ------- + VolumeRegistrationRunPayload + Typed payload for a between-scan registration run. + + Raises + ------ + ValueError + If `payload["operation"]` is not `"register_volume"`. + """ + if payload.get("operation") != "register_volume": + raise ValueError("Expected a register_volume payload.") + return cast("VolumeRegistrationRunPayload", payload) + + +def coerce_volumewise_registration_payload( + payload: dict[str, Any] | "VolumewiseRegistrationRunPayload", +) -> "VolumewiseRegistrationRunPayload": + """Return a typed within-scan registration payload. + + Parameters + ---------- + payload : dict[str, Any] or VolumewiseRegistrationRunPayload + Untyped or typed payload captured when the worker started. + + Returns + ------- + VolumewiseRegistrationRunPayload + Typed payload for a within-scan registration run. + + Raises + ------ + ValueError + If `payload["operation"]` is not `"register_volumewise"`. + """ + if payload.get("operation") != "register_volumewise": + raise ValueError("Expected a register_volumewise payload.") + return cast("VolumewiseRegistrationRunPayload", payload) + + +def finalize_registration_layer( + panel: "RegistrationPanel", + *, + payload: "VolumeRegistrationRunPayload | VolumewiseRegistrationRunPayload", + registered: xr.DataArray, + layer_name: str, + metadata: dict[str, Any], + registration_status: Literal["completed", "aborted"], +) -> None: + """Attach registration metadata and add or update the result layer. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose viewer should receive the final layer. + payload : VolumeRegistrationRunPayload or VolumewiseRegistrationRunPayload + Typed UI snapshot captured before the worker started. + registered : xarray.DataArray + Registered output returned by the worker. + layer_name : str + Name to use when creating a new result layer. + metadata : dict[str, Any] + Extra layer metadata to attach. + registration_status : {"completed", "aborted"} + Final run status used for user feedback and layer naming. + """ + metadata["registration_operation"] = payload["operation"] + metadata["registration_parameters"] = payload.copy() + + source_layer = panel._get_layer_by_name(payload["moving_layer_name"]) + display_kwargs = ( + _get_image_display_kwargs_from_layer(source_layer) + if source_layer is not None + else {} + ) + if _gamma_needs_reset(payload.get("scale", "off")): + display_kwargs["gamma"] = 1.0 + if payload["operation"] == "register_volume": + display_kwargs["colormap"] = "cyan" + display_kwargs["blending"] = "additive" + contrast_limits = tuple(calc_data_range(registered.data)) + + if payload["operation"] == "register_volume" and panel._progress_layer is not None: + layer = panel._progress_layer + panel._set_image_layer_data(layer, np.asarray(registered.data)) + if hasattr(layer, "contrast_limits"): + layer.contrast_limits = contrast_limits + panel._progress_bridge = None + panel._progress_layer = None + elif ( + payload["operation"] == "register_volumewise" + and panel._volumewise_progress_layer is not None + ): + layer = panel._volumewise_progress_layer + panel._set_image_layer_data(layer, np.asarray(registered.data)) + if hasattr(layer, "contrast_limits"): + layer.contrast_limits = contrast_limits + teardown_volumewise_progress(panel, remove_layer=False) + else: + _, layer = plot_napari( + registered, + viewer=panel.viewer, + name=layer_name, + show_colorbar=False, + contrast_limits=contrast_limits, + **display_kwargs, + ) + layer.metadata.update(metadata) + layer.metadata["xarray"] = registered + panel.viewer.layers.selection.active = layer + refresh_transform_controls(panel) + + if payload["operation"] == "register_volumewise": + panel._progress.setValue(panel._progress.maximum()) + + if registration_status == "aborted": + layer.name = f"{layer.name} (aborted)" + panel._set_error("Registration aborted; added partial result.") + show_info(f"Registration aborted; added partial layer: {layer.name}") + else: + show_info(f"Added registered layer: {layer.name}") + + +def on_registration_finished( + panel: "RegistrationPanel", + payload: dict[str, Any], + result: object, +) -> None: + """Dispatch a finished registration callback to the typed handler. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel that started the worker. + payload : dict[str, Any] + Untyped compatibility payload captured when the worker started. + result : object + Worker result to forward to the operation-specific handler. + + Raises + ------ + ValueError + If `payload["operation"]` is not recognized. + """ + if payload.get("operation") == "register_volume": + on_volume_registration_finished( + panel, + coerce_volume_registration_payload(payload), + cast( + "tuple[xr.DataArray, npt.NDArray[np.floating] | xr.DataArray, RegistrationDiagnostics]", + result, + ), + ) + return + if payload.get("operation") == "register_volumewise": + on_volumewise_registration_finished( + panel, + coerce_volumewise_registration_payload(payload), + cast("xr.DataArray", result), + ) + return + raise ValueError(f"Unknown registration operation: {payload.get('operation')!r}.") + + +def on_volume_registration_finished( + panel: "RegistrationPanel", + payload: "VolumeRegistrationRunPayload", + result: tuple[ + xr.DataArray, + "npt.NDArray[np.floating] | xr.DataArray", + "RegistrationDiagnostics", + ], +) -> None: + """Add a between-scan registration result back to the viewer. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel that started the worker. + payload : VolumeRegistrationRunPayload + Typed UI snapshot captured before the worker started. + result : tuple + Registered volume, estimated transform, and diagnostics. + """ + registered, transform, diagnostics = result + registered = registered.copy(deep=False) + registered.attrs = registered.attrs.copy() + registered.attrs["registration_transform"] = transform + registered.attrs["registration_diagnostics"] = diagnostics + registered.attrs["registration_operation"] = payload["operation"] + registered.attrs["registration_status"] = diagnostics.status + metadata: dict[str, Any] = { + "registration_transform": transform, + "registration_diagnostics": diagnostics, + "registration_status": diagnostics.status, + } + transform_name = panel._make_unique_transform_name( + f"{payload['moving_layer_name']} → {payload['fixed_layer_name']} ({payload['transform']})" + ) + source_layer = panel._get_layer_by_name(payload["moving_layer_name"]) + source_data = ( + _prepare_between_scan_data(_get_source_dataarray(source_layer)) + if source_layer is not None + else None + ) + if isinstance(transform, np.ndarray): + metadata["confusius_transform"] = make_affine_transform_payload( + np.asarray(transform, dtype=float), + reference=registered, + source=source_data, + source_layer_name=payload["moving_layer_name"], + target_layer_name=payload["fixed_layer_name"], + operation=payload["operation"], + transform_model=payload["transform"], + metric=payload["metric"], + diagnostics=diagnostics, + name=transform_name, + ) + else: + metadata["confusius_transform"] = make_bspline_transform_payload( + transform, + reference=registered, + source=source_data, + source_layer_name=payload["moving_layer_name"], + target_layer_name=payload["fixed_layer_name"], + operation=payload["operation"], + transform_model=payload["transform"], + metric=payload["metric"], + diagnostics=diagnostics, + name=transform_name, + ) + finalize_registration_layer( + panel, + payload=payload, + registered=registered, + layer_name=panel._volume_result_layer_name( + payload["moving_layer_name"], + payload["fixed_layer_name"], + transform_model=payload["transform"], + ), + metadata=metadata, + registration_status=diagnostics.status, + ) + + +def on_volumewise_registration_finished( + panel: "RegistrationPanel", + payload: "VolumewiseRegistrationRunPayload", + result: xr.DataArray, +) -> None: + """Add a within-scan registration result back to the viewer. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel that started the worker. + payload : VolumewiseRegistrationRunPayload + Typed UI snapshot captured before the worker started. + result : xarray.DataArray + Motion-corrected time series returned by the worker. + """ + registered = result.copy(deep=False) + registered.attrs = registered.attrs.copy() + registered.attrs["registration_operation"] = payload["operation"] + motion_params = registered.attrs.get("motion_params") + registration_status = "completed" + if motion_params is not None: + try: + statuses = motion_params["status"] + except Exception: # noqa: BLE001 + statuses = None + if statuses is not None and bool((statuses == "aborted").any()): + registration_status = "aborted" + finalize_registration_layer( + panel, + payload=payload, + registered=registered, + layer_name=panel._volumewise_result_layer_name(payload["moving_layer_name"]), + metadata={ + "motion_params": motion_params, + "reference_time": payload["reference_time"], + }, + registration_status=registration_status, + ) diff --git a/src/confusius/_napari/_registration/_panel_selection.py b/src/confusius/_napari/_registration/_panel_selection.py new file mode 100644 index 00000000..03c1e919 --- /dev/null +++ b/src/confusius/_napari/_registration/_panel_selection.py @@ -0,0 +1,457 @@ +"""Layer-selection and validation helpers for the napari registration panel.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, cast + +from qtpy.QtWidgets import QComboBox + +from confusius._dims import TIME_DIM +from confusius._napari._registration._panel_transforms import ( + refresh_transform_controls, + validate_initial_transform_selection, +) +from confusius._napari._registration._panel_utils import ( + _get_source_dataarray, + _is_registration_source_layer, + _prepare_between_scan_data, +) + +if TYPE_CHECKING: + from napari.layers import Layer + + from confusius._napari._registration._panel import ( + MetricName, + RegistrationPanel, + ResampleInterpolation, + ScaleMode, + VolumeTransformType, + VolumewiseTransformType, + ) + + +def refresh_layers(panel: "RegistrationPanel") -> None: + """Repopulate the layer selectors from the viewer. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose layer selectors should be refreshed. + """ + moving_name = panel._moving_combo.currentText() + fixed_name = panel._fixed_combo.currentText() + fixed_mask_name = panel._fixed_mask_combo.currentText() + moving_mask_name = panel._moving_mask_combo.currentText() + + layer_names = [ + layer.name + for layer in panel.viewer.layers + if _is_registration_source_layer(layer) + ] + labels_layer_names = [ + layer.name for layer in panel.viewer.layers if layer._type_string == "labels" + ] + + panel._moving_combo.blockSignals(True) + panel._fixed_combo.blockSignals(True) + panel._fixed_mask_combo.blockSignals(True) + panel._moving_mask_combo.blockSignals(True) + panel._moving_combo.clear() + panel._fixed_combo.clear() + panel._fixed_mask_combo.clear() + panel._moving_mask_combo.clear() + panel._moving_combo.addItems(layer_names) + panel._fixed_combo.addItems(layer_names) + panel._fixed_mask_combo.addItem("") + panel._moving_mask_combo.addItem("") + panel._fixed_mask_combo.addItems(labels_layer_names) + panel._moving_mask_combo.addItems(labels_layer_names) + panel._moving_combo.blockSignals(False) + panel._fixed_combo.blockSignals(False) + panel._fixed_mask_combo.blockSignals(False) + panel._moving_mask_combo.blockSignals(False) + + moving_index = panel._moving_combo.findText(moving_name) + if moving_index >= 0: + panel._moving_combo.setCurrentIndex(moving_index) + + fixed_index = panel._fixed_combo.findText(fixed_name) + if fixed_index >= 0: + panel._fixed_combo.setCurrentIndex(fixed_index) + elif ( + panel._fixed_combo.count() > 1 + and panel._fixed_combo.currentText() == panel._moving_combo.currentText() + ): + panel._fixed_combo.setCurrentIndex(1) + + fixed_mask_index = panel._fixed_mask_combo.findText(fixed_mask_name) + if fixed_mask_index >= 0: + panel._fixed_mask_combo.setCurrentIndex(fixed_mask_index) + + moving_mask_index = panel._moving_mask_combo.findText(moving_mask_name) + if moving_mask_index >= 0: + panel._moving_mask_combo.setCurrentIndex(moving_mask_index) + + update_reference_time_bounds(panel) + panel._sync_manual_transform_event_connections() + refresh_transform_controls(panel) + panel._sync_optimizer_weight_editor() + validate_registration_selection(panel) + + +def get_layer_by_name(panel: "RegistrationPanel", name: str) -> "Layer | None": + """Return a viewer layer by name, if present. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose viewer should be searched. + name : str + Layer name to look up in the viewer. + + Returns + ------- + napari.layers.Layer or None + Matching layer when present, otherwise `None`. + """ + try: + return cast("Layer", panel.viewer.layers[name]) + except KeyError: + return None + + +def selected_layer(panel: "RegistrationPanel", combo: QComboBox) -> "Layer | None": + """Return the currently selected viewer layer for a combo box. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose viewer should be searched. + combo : QComboBox + Combo box containing layer names. + + Returns + ------- + napari.layers.Layer or None + Selected layer, or `None` when no valid selection exists. + """ + name = combo.currentText() + if not name: + return None + return get_layer_by_name(panel, name) + + +def current_scale_mode(panel: "RegistrationPanel") -> "ScaleMode": + """Return the validated registration scale mode from the combo box. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose scale selector should be read. + + Returns + ------- + {"off", "dB", "sqrt"} + Selected registration scale mode. + + Raises + ------ + ValueError + If the combo box contains an unexpected value. + """ + value = panel._scale_combo.currentData() + if value in {"off", "dB", "sqrt"}: + return value + raise ValueError(f"Unknown registration scale mode: {value!r}.") + + +def current_metric(panel: "RegistrationPanel") -> "MetricName": + """Return the validated registration metric from the combo box. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose metric selector should be read. + + Returns + ------- + {"correlation", "mattes_mi"} + Selected registration metric. + + Raises + ------ + ValueError + If the combo box contains an unexpected value. + """ + value = panel._metric_combo.currentText() + if value == "correlation": + return "correlation" + if value == "mattes_mi": + return "mattes_mi" + raise ValueError(f"Unknown registration metric: {value!r}.") + + +def current_resample_interpolation( + panel: "RegistrationPanel", +) -> "ResampleInterpolation": + """Return the validated resampling interpolation from the combo box. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose interpolation selector should be read. + + Returns + ------- + {"linear", "bspline"} + Selected resampling interpolation. + + Raises + ------ + ValueError + If the combo box contains an unexpected value. + """ + value = panel._interpolation_combo.currentText() + if value == "linear": + return "linear" + if value == "bspline": + return "bspline" + raise ValueError(f"Unknown resampling interpolation: {value!r}.") + + +def current_transform_model( + panel: "RegistrationPanel", +) -> "VolumeTransformType | VolumewiseTransformType": + """Return the validated transform model for the active mode. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose transform selector should be read. + + Returns + ------- + {"translation", "rigid", "affine", "bspline"} + Selected transform model, constrained by the active workflow. + + Raises + ------ + ValueError + If the combo box contains an unexpected value. + """ + value = panel._transform_combo.currentText() + if panel._operation() == "register_volume": + if value == "translation": + return "translation" + if value == "rigid": + return "rigid" + if value == "affine": + return "affine" + if value == "bspline": + return "bspline" + else: + if value == "translation": + return "translation" + if value == "rigid": + return "rigid" + if value == "affine": + return "affine" + raise ValueError(f"Unknown transform model: {value!r}.") + + +def update_reference_time_bounds(panel: "RegistrationPanel") -> None: + """Clamp the volumewise reference-volume widget to the moving layer. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose reference-volume bounds should be updated. + """ + moving_layer = selected_layer(panel, panel._moving_combo) + if moving_layer is None: + panel._reference_time_spin.setMaximum(0) + panel._reference_time_spin.setValue(0) + return + + data = _get_source_dataarray(moving_layer) + if TIME_DIM not in data.dims: + panel._reference_time_spin.setMaximum(0) + panel._reference_time_spin.setValue(0) + return + + panel._reference_time_spin.setMaximum(max(0, data.sizes[TIME_DIM] - 1)) + + +def set_layer_validation_style( + panel: "RegistrationPanel", + *, + moving_invalid: bool = False, + fixed_invalid: bool = False, + message: str | None = None, +) -> None: + """Update inline validation state for the layer selectors. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose validation widgets should be updated. + moving_invalid : bool, default: False + Whether to mark the moving-layer selector as invalid. + fixed_invalid : bool, default: False + Whether to mark the fixed-layer selector as invalid. + message : str, optional + Validation message to show below the layer selectors. + """ + error_style = "border: 1px solid #e05555;" + normal_style = "" + panel._moving_combo.setStyleSheet(error_style if moving_invalid else normal_style) + panel._fixed_combo.setStyleSheet(error_style if fixed_invalid else normal_style) + panel._moving_label.setStyleSheet("color: #e05555;" if moving_invalid else "") + panel._fixed_label.setStyleSheet("color: #e05555;" if fixed_invalid else "") + panel._reference_time_label.setStyleSheet("") + if message: + panel._layer_validation.setText(message) + panel._layer_validation.show() + else: + panel._layer_validation.hide() + panel._layer_validation.clear() + + +def set_run_btn_enabled(panel: "RegistrationPanel", enabled: bool) -> None: + """Enable or disable the Run button without changing its busy text. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose Run button should be updated. + enabled : bool + Whether to enable the idle-state Run button. + + Notes + ----- + The button is also disabled in `_begin_work` while a registration is + running; this helper only handles the idle-state gating driven by + layer-selection validation. + """ + if panel._run_btn.text() == "Registering…": + return + panel._run_btn.setEnabled(enabled) + + +def validate_registration_selection(panel: "RegistrationPanel") -> bool: + """Validate the current registration-layer selection and show inline feedback. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose selection state should be validated. + + Returns + ------- + bool + `True` when the selection is valid and a registration can be started, + `False` otherwise. As a side effect, the Run button is enabled or disabled + to match the validation result. + """ + moving_layer = selected_layer(panel, panel._moving_combo) + fixed_layer = selected_layer(panel, panel._fixed_combo) + operation = panel._operation() + + if moving_layer is None: + set_layer_validation_style(panel) + set_run_btn_enabled(panel, False) + return False + + try: + moving = _get_source_dataarray(moving_layer) + except Exception: + set_layer_validation_style( + panel, + moving_invalid=True, + message="Could not read the selected moving layer.", + ) + set_run_btn_enabled(panel, False) + return False + + if operation == "register_volumewise": + if TIME_DIM not in moving.dims: + set_layer_validation_style( + panel, + moving_invalid=True, + message="Within-scan registration requires a layer with a time dimension.", + ) + set_run_btn_enabled(panel, False) + return False + init_message = validate_initial_transform_selection( + panel, + operation=operation, + moving=moving, + ) + set_layer_validation_style(panel, message=init_message) + set_run_btn_enabled(panel, init_message is None) + return init_message is None + + moving_invalid = False + fixed_invalid = False + message: str | None = None + + if fixed_layer is None: + set_layer_validation_style( + panel, + moving_invalid=moving_invalid, + fixed_invalid=True, + message="Between-scans registration requires different moving and fixed layers.", + ) + set_run_btn_enabled(panel, False) + return False + + try: + fixed = _get_source_dataarray(fixed_layer) + except Exception: + set_layer_validation_style( + panel, + fixed_invalid=True, + message="Could not read the selected fixed layer.", + ) + set_run_btn_enabled(panel, False) + return False + + if fixed_layer is moving_layer: + moving_invalid = True + fixed_invalid = True + message = "Moving and fixed layers must be different." + + if message is None: + message = validate_initial_transform_selection( + panel, + operation=operation, + moving=_prepare_between_scan_data(moving), + fixed=_prepare_between_scan_data(fixed), + ) + + valid = not (moving_invalid or fixed_invalid or message is not None) + set_layer_validation_style( + panel, + moving_invalid=moving_invalid, + fixed_invalid=fixed_invalid, + message=message, + ) + set_run_btn_enabled(panel, valid) + return valid + + +def on_moving_layer_changed(panel: "RegistrationPanel", _name: str) -> None: + """Update dependent widgets when the moving layer changes. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose dependent widgets should be refreshed. + _name : str + Unused emitted layer name from the combo-box signal. + """ + del _name + update_reference_time_bounds(panel) + refresh_transform_controls(panel) + panel._sync_optimizer_weight_editor() + validate_registration_selection(panel) diff --git a/src/confusius/_napari/_registration/_panel_transforms.py b/src/confusius/_napari/_registration/_panel_transforms.py new file mode 100644 index 00000000..36ac3516 --- /dev/null +++ b/src/confusius/_napari/_registration/_panel_transforms.py @@ -0,0 +1,895 @@ +"""Transform payload and panel-specific transform helpers for the napari registration panel.""" + +from __future__ import annotations + +from collections.abc import Sequence +from pathlib import Path +from typing import TYPE_CHECKING, Literal, cast + +import numpy as np +import numpy.typing as npt +import xarray as xr +from napari.layers.utils.layer_utils import calc_data_range +from napari.qt.threading import thread_worker +from napari.utils.notifications import show_error, show_info +from qtpy.QtWidgets import QFileDialog + +from confusius._dims import SPATIAL_DIMS +from confusius._napari._registration._panel_utils import ( + _get_image_display_kwargs_from_layer, + _get_source_dataarray, + _prepare_between_scan_data, +) +from confusius._napari._registration._panel_worker_state import on_registration_failed +from confusius._napari._registration._transform_payloads import ( + AffineTransformPayload, + OutputGridPayload, + TransformPayload, + get_affine_transform_from_payload, + get_bspline_transform_from_payload, + get_input_grid_from_payload, + get_output_grid_from_payload, + load_transform_payload, + make_output_grid_payload, + save_transform_payload, +) +from confusius.plotting.napari import plot_napari +from confusius.registration import resample_volume + +if TYPE_CHECKING: + from napari.layers import Layer + + from confusius._napari._registration._panel import ( + ApplyTransformPayload, + RegistrationPanel, + TransformSourceData, + ) + + +def _get_affine_payload_from_layer(layer: "Layer") -> AffineTransformPayload | None: + """Return the stored affine transform payload for a napari layer. + + Parameters + ---------- + layer : napari.layers.Layer + Layer whose metadata should be inspected. + + Returns + ------- + AffineTransformPayload or None + Stored payload when present and affine, otherwise `None`. + """ + payload = layer.metadata.get("confusius_transform") + if not isinstance(payload, dict) or payload.get("kind") != "affine": + return None + get_affine_transform_from_payload(payload) + return cast("AffineTransformPayload", payload) + + +def _get_spatial_manual_affine_from_layer( + layer: "Layer", *, spatial_dims: Sequence[str] +) -> npt.NDArray[np.float64]: + """Return the spatial sub-affine from a napari layer's manual transform. + + Parameters + ---------- + layer : napari.layers.Layer + Layer whose manual napari affine should be extracted. + spatial_dims : sequence of str + Spatial dimension names, in the exact order expected by registration. + + Returns + ------- + (N+1, N+1) numpy.ndarray + Spatial homogeneous affine in world coordinates. + + Raises + ------ + ValueError + If the layer does not contain the requested spatial dimensions. + ValueError + If the layer affine has an unexpected shape. + ValueError + If the manual affine mixes selected spatial axes with ignored axes. + """ + data = _get_source_dataarray(layer) + layer_dims = [str(dim) for dim in data.dims] + missing_dims = [dim for dim in spatial_dims if dim not in layer_dims] + if missing_dims: + raise ValueError( + "Selected manual napari transform does not contain spatial dims " + f"{missing_dims}." + ) + + affine = np.asarray(layer.affine.affine_matrix, dtype=float) + expected_shape = (len(layer_dims) + 1, len(layer_dims) + 1) + if affine.shape != expected_shape: + raise ValueError( + f"Selected manual napari transform has shape {affine.shape}, " + f"but layer '{layer.name}' expects {expected_shape}." + ) + + spatial_indices = [layer_dims.index(dim) for dim in spatial_dims] + ignored_indices = [i for i in range(len(layer_dims)) if i not in spatial_indices] + linear = affine[:-1, :-1] + + if ignored_indices: + spatial_to_ignored = linear[np.ix_(spatial_indices, ignored_indices)] + ignored_to_spatial = linear[np.ix_(ignored_indices, spatial_indices)] + if not np.allclose(spatial_to_ignored, 0.0) or not np.allclose( + ignored_to_spatial, 0.0 + ): + raise ValueError( + "Selected manual napari transform mixes spatial axes with ignored " + "non-spatial axes, so it cannot be used as a registration " + "initialization." + ) + + spatial_affine = np.eye(len(spatial_dims) + 1, dtype=float) + spatial_affine[:-1, :-1] = linear[np.ix_(spatial_indices, spatial_indices)] + spatial_affine[:-1, -1] = affine[np.ix_(spatial_indices, [-1])].ravel() + return spatial_affine + + +def _make_manual_transform_payload(layer: "Layer") -> AffineTransformPayload: + """Build an affine payload from a layer's manual napari transform. + + Parameters + ---------- + layer : napari.layers.Layer + Layer whose current manual napari transform should be serialized. + + Returns + ------- + AffineTransformPayload + JSON-serializable affine payload representing the visible manual layer transform + on the layer's own spatial output grid. + """ + data = _get_source_dataarray(layer) + spatial_data = _prepare_between_scan_data(data) + spatial_dims = [str(dim) for dim in spatial_data.dims if dim in SPATIAL_DIMS] + manual_affine = _get_spatial_manual_affine_from_layer( + layer, spatial_dims=spatial_dims + ) + pull_affine = np.linalg.inv(manual_affine) + return { + "kind": "affine", + "name": f"{layer.name} (manual)", + "affine": pull_affine.tolist(), + "source_layer_name": layer.name, + "target_layer_name": layer.name, + "operation": "manual_napari_transform", + "transform_model": "affine", + "metric": "manual", + "output_grid": make_output_grid_payload(spatial_data), + "input_grid": make_output_grid_payload(spatial_data), + "diagnostics": { + "metric": "manual", + "final_metric_value": 0.0, + "n_iterations": 0, + "stop_condition": "Saved from manual napari layer transform.", + "status": "completed", + }, + } + + +def get_transform_source_data(value: object) -> "TransformSourceData | None": + """Return validated transform-source combo data. + + Parameters + ---------- + value : object + Raw combo-box payload to validate. + + Returns + ------- + tuple[str, str] or None + Validated `(kind, name)` pair, or `None` when the payload does not match the + expected transform-source schema. + """ + if not isinstance(value, tuple) or len(value) != 2: + return None + source_kind, source_name = value + if not isinstance(source_name, str): + return None + if source_kind == "loaded": + return ("loaded", source_name) + if source_kind == "layer": + return ("layer", source_name) + if source_kind == "manual": + return ("manual", source_name) + return None + + +def get_transform_payload_from_metadata(payload: object) -> TransformPayload | None: + """Return a validated transform payload stored in layer metadata. + + Parameters + ---------- + payload : object + Raw metadata payload to validate. + + Returns + ------- + TransformPayload or None + Validated transform payload, or `None` when the metadata does not contain a + supported transform payload. + """ + if not isinstance(payload, dict): + return None + payload_mapping = cast("dict[str, object]", payload) + kind = payload_mapping.get("kind") + if kind == "affine": + get_affine_transform_from_payload(payload_mapping) + return cast("TransformPayload", payload_mapping) + if kind == "bspline": + get_bspline_transform_from_payload(payload_mapping) + return cast("TransformPayload", payload_mapping) + return None + + +def get_transform_source_label( + payload: TransformPayload, *, suffix: str | None = None +) -> str: + """Return a user-facing label for a transform payload. + + Parameters + ---------- + payload : TransformPayload + Transform payload to label. + suffix : str, optional + Unused legacy suffix parameter kept to avoid wider churn. + + Returns + ------- + str + Label shown in transform selectors. + """ + del suffix + return payload["name"] + + +def get_available_transform_payloads( + panel: "RegistrationPanel", +) -> list[TransformPayload]: + """Return all transform payloads currently available in the UI. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose loaded payload and viewer layers are queried. + + Returns + ------- + list of TransformPayload + Loaded payload plus any validated payloads found on viewer layers. + """ + payloads: list[TransformPayload] = [] + if panel._loaded_transform_payload is not None: + payloads.append(panel._loaded_transform_payload) + for layer in panel.viewer.layers: + payload = get_transform_payload_from_metadata( + layer.metadata.get("confusius_transform") + ) + if payload is not None: + payloads.append(payload) + return payloads + + +def refresh_transform_controls(panel: "RegistrationPanel") -> None: + """Refresh the transform, initialization, and target selectors from the current viewer state. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose transform selectors are updated. + """ + source_data = panel._transform_source_combo.currentData() + initialization_data = panel._initialization_combo.currentData() + target_name = panel._transform_target_combo.currentText() + + transform_options: list[tuple[str, tuple[str, str]]] = [] + if panel._loaded_transform_payload is not None: + transform_options.append( + ( + get_transform_source_label( + panel._loaded_transform_payload, + suffix="loaded", + ), + ("loaded", ""), + ) + ) + for layer in panel.viewer.layers: + payload = get_transform_payload_from_metadata( + layer.metadata.get("confusius_transform") + ) + if payload is None: + continue + transform_options.append( + ( + get_transform_source_label(payload, suffix=layer.name), + ("layer", layer.name), + ) + ) + + manual_transform_options: list[tuple[str, tuple[str, str]]] = [] + manual_initialization_options: list[tuple[str, tuple[str, str]]] = [] + for layer in panel.viewer.layers: + try: + data = _get_source_dataarray(layer) + spatial_dims = [str(dim) for dim in data.dims if dim in SPATIAL_DIMS] + if not spatial_dims: + continue + manual_affine = _get_spatial_manual_affine_from_layer( + layer, + spatial_dims=spatial_dims, + ) + except Exception: # noqa: BLE001 + continue + if np.allclose(manual_affine, np.eye(len(spatial_dims) + 1)): + continue + manual_option = (f"{layer.name} (manual)", ("manual", layer.name)) + manual_transform_options.append(manual_option) + manual_initialization_options.append(manual_option) + + panel._transform_source_combo.blockSignals(True) + panel._transform_source_combo.clear() + for label, data in transform_options: + panel._transform_source_combo.addItem(label, data) + for label, data in manual_transform_options: + panel._transform_source_combo.addItem(label, data) + panel._transform_source_combo.blockSignals(False) + + panel._initialization_combo.blockSignals(True) + panel._initialization_combo.clear() + panel._initialization_combo.addItem("center_geometry", "center_geometry") + panel._initialization_combo.addItem("center_moments", "center_moments") + panel._initialization_combo.addItem("none", None) + for label, data in transform_options: + source_kind, source_name = data + if source_kind == "loaded": + if panel._loaded_transform_payload is None: + continue + if panel._loaded_transform_payload["kind"] != "affine": + continue + elif source_kind == "layer": + layer = panel._get_layer_by_name(source_name) + if layer is None or _get_affine_payload_from_layer(layer) is None: + continue + panel._initialization_combo.addItem(label, data) + for label, data in manual_initialization_options: + panel._initialization_combo.addItem(label, data) + panel._initialization_combo.blockSignals(False) + + panel._transform_target_combo.blockSignals(True) + panel._transform_target_combo.clear() + panel._transform_target_combo.addItems( + [layer.name for layer in panel.viewer.layers] + ) + panel._transform_target_combo.blockSignals(False) + + if source_data is not None: + for i in range(panel._transform_source_combo.count()): + if panel._transform_source_combo.itemData(i) == source_data: + panel._transform_source_combo.setCurrentIndex(i) + break + + if initialization_data is not None: + for i in range(panel._initialization_combo.count()): + if panel._initialization_combo.itemData(i) == initialization_data: + panel._initialization_combo.setCurrentIndex(i) + break + + target_index = panel._transform_target_combo.findText(target_name) + if target_index >= 0: + panel._transform_target_combo.setCurrentIndex(target_index) + + +def get_selected_transform_payload( + panel: "RegistrationPanel", +) -> TransformPayload | None: + """Return the currently selected transform payload. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose transform selection is read. + + Returns + ------- + TransformPayload or None + Selected transform payload, or `None` when no valid selection is available. + """ + source_data = get_transform_source_data(panel._transform_source_combo.currentData()) + if source_data is None: + return None + + source_kind, source_name = source_data + if source_kind == "loaded": + return panel._loaded_transform_payload + if not source_name: + return None + layer = panel._get_layer_by_name(source_name) + if layer is None: + return None + if source_kind == "layer": + return get_transform_payload_from_metadata( + layer.metadata.get("confusius_transform") + ) + if source_kind == "manual": + return _make_manual_transform_payload(layer) + return None + + +def get_selected_center_initialization( + panel: "RegistrationPanel", +) -> Literal["center_geometry", "center_moments"] | None: + """Return the selected built-in centering initialization, if any. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose initialization selector is read. + + Returns + ------- + {"center_geometry", "center_moments"} or None + Selected built-in centering initialization, or `None` when a different kind of + initialization is currently selected. + """ + value = panel._initialization_combo.currentData() + if value in {"center_geometry", "center_moments"}: + return value + return None + + +def get_selected_initial_transform_payload( + panel: "RegistrationPanel", +) -> AffineTransformPayload | None: + """Return the payload selected for registration initialization, if any. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose initialization selector is read. + + Returns + ------- + AffineTransformPayload or None + Selected affine transform payload, or `None` when no affine payload is currently + selected for initialization. + """ + source_data = get_transform_source_data(panel._initialization_combo.currentData()) + if source_data is None: + return None + + source_kind, source_name = source_data + if source_kind == "loaded": + if ( + panel._loaded_transform_payload is not None + and panel._loaded_transform_payload["kind"] == "affine" + ): + return panel._loaded_transform_payload + return None + if source_kind != "layer" or not source_name: + return None + layer = panel._get_layer_by_name(source_name) + if layer is None: + return None + return _get_affine_payload_from_layer(layer) + + +def get_selected_manual_initialization_layer( + panel: "RegistrationPanel", +) -> "Layer | None": + """Return the layer selected for manual napari initialization, if any. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose initialization selector is read. + + Returns + ------- + napari.layers.Layer or None + Layer selected as the manual napari initialization source, or `None` when no + manual initialization is currently selected. + """ + source_data = get_transform_source_data(panel._initialization_combo.currentData()) + if source_data is None: + return None + + source_kind, source_name = source_data + if source_kind != "manual" or not source_name: + return None + return panel._get_layer_by_name(source_name) + + +def get_selected_initial_transform( + panel: "RegistrationPanel", + moving: xr.DataArray, + *, + moving_layer: "Layer | None" = None, + fixed_layer: "Layer | None" = None, +) -> tuple[npt.NDArray[np.float64] | None, str | None]: + """Return the selected initialization affine and its source label. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose initialization selector is read. + moving : xarray.DataArray + Moving DataArray defining the spatial dimensions of the registration. + moving_layer : napari.layers.Layer, optional + Layer used as the moving input; required when a manual initialization is + selected. + fixed_layer : napari.layers.Layer, optional + Layer used as the fixed input; required when a manual initialization is + selected. + + Returns + ------- + affine : (N+1, N+1) numpy.ndarray or None + Selected initialization affine in homogeneous coordinates, or `None` when no + initialization is selected. + label : str or None + Human-readable label for the selected initialization source, or `None` when no + initialization is selected. + + Raises + ------ + ValueError + If a manual initialization is selected but the moving and fixed layers are not + provided, or the selected manual layer is not the current moving or fixed layer. + """ + payload = get_selected_initial_transform_payload(panel) + if payload is not None: + return get_affine_transform_from_payload(payload), payload["name"] + + layer = get_selected_manual_initialization_layer(panel) + if layer is None: + return None, None + if moving_layer is None or fixed_layer is None: + raise ValueError("Select moving and fixed layers.") + if layer not in {moving_layer, fixed_layer}: + raise ValueError( + "Selected manual initialization must come from the current moving " + "or fixed layer." + ) + + spatial_dims = [str(dim) for dim in moving.dims if dim in SPATIAL_DIMS] + moving_affine = _get_spatial_manual_affine_from_layer( + moving_layer, + spatial_dims=spatial_dims, + ) + fixed_affine = _get_spatial_manual_affine_from_layer( + fixed_layer, + spatial_dims=spatial_dims, + ) + affine = np.linalg.inv(moving_affine) @ fixed_affine + return affine, f"{layer.name} (manual)" + + +def validate_initial_transform_selection( + panel: "RegistrationPanel", + *, + operation: Literal["register_volume", "register_volumewise"], + moving: xr.DataArray, + fixed: xr.DataArray | None = None, +) -> str | None: + """Return an inline validation message for transform initialization. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose selection is validated. + operation : {"register_volume", "register_volumewise"} + Registration operation the panel is currently configured to run. + moving : xarray.DataArray + Moving DataArray used to check initialization transform shapes. + fixed : xarray.DataArray, optional + Fixed DataArray; required when an initialization is selected. + + Returns + ------- + str or None + Human-readable validation message, or `None` when the current selection is valid + (or no initialization is selected). + """ + if operation != "register_volume": + return None + if ( + get_selected_initial_transform_payload(panel) is None + and get_selected_manual_initialization_layer(panel) is None + ): + return None + if fixed is None: + return "Select a fixed layer." + + moving_layer = panel._selected_layer(panel._moving_combo) + fixed_layer = panel._selected_layer(panel._fixed_combo) + + try: + affine, _ = get_selected_initial_transform( + panel, + moving, + moving_layer=moving_layer, + fixed_layer=fixed_layer, + ) + except Exception as exc: # noqa: BLE001 + return str(exc) + + if affine is None: + return None + + expected_shape = (moving.ndim + 1, moving.ndim + 1) + if affine.shape != expected_shape: + return ( + f"Selected initialization transform has shape {affine.shape}, " + f"but this registration expects {expected_shape}." + ) + return None + + +def save_selected_transform(panel: "RegistrationPanel") -> None: + """Prompt for a destination path and save the currently selected transform payload. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose selected transform should be saved. + """ + payload = get_selected_transform_payload(panel) + if payload is None: + panel._set_error("Select a transform to save.") + return + + default_name = payload["name"].replace("/", "-") + suffix = ".json" if payload["kind"] == "affine" else ".zarr" + file_filter = ( + "JSON files (*.json)" if payload["kind"] == "affine" else "Zarr stores (*.zarr)" + ) + start = str(Path.home() / f"{default_name}{suffix}") + path_str, _ = QFileDialog.getSaveFileName( + panel, "Save transform", start, file_filter + ) + if not path_str: + return + + save_transform_payload(path_str, payload) + show_info(f"Saved transform: {path_str}") + + +def load_transform(panel: "RegistrationPanel") -> None: + """Prompt for a transform file, load it into the panel state, and refresh the selectors. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel that should receive the loaded transform. + """ + start = str(Path.home()) + path_str, _ = QFileDialog.getOpenFileName( + panel, + "Load transform", + start, + "Transform files (*.json *.zarr)", + ) + if not path_str: + return + + try: + panel._loaded_transform_payload = load_transform_payload(path_str) + except Exception as exc: # noqa: BLE001 + panel._set_error(str(exc)) + show_error(str(exc)) + return + + refresh_transform_controls(panel) + for i in range(panel._transform_source_combo.count()): + if panel._transform_source_combo.itemData(i) == ("loaded", ""): + panel._transform_source_combo.setCurrentIndex(i) + break + show_info(f"Loaded transform: {panel._loaded_transform_payload['name']}") + + +def _get_inverse_output_grid( + panel: "RegistrationPanel", payload: TransformPayload +) -> OutputGridPayload: + """Return the output grid to use when applying a transform inverse. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel that owns the transform selection. + payload : TransformPayload + Selected transform payload. + + Returns + ------- + OutputGridPayload + Grid of the original moving/source layer. + + Raises + ------ + ValueError + If the payload predates `input_grid` and the source layer is not available to + re-derive it. + """ + input_grid = get_input_grid_from_payload(payload) + if input_grid is not None: + return input_grid + + source_layer = panel._get_layer_by_name(payload["source_layer_name"]) + if source_layer is None: + raise ValueError( + "Transform payload does not contain an input grid. Reload the original " + "source layer or re-save the transform from a newer registration result." + ) + + source = _prepare_between_scan_data(_get_source_dataarray(source_layer)) + return make_output_grid_payload(source) + + +def apply_selected_transform(panel: "RegistrationPanel") -> None: + """Start a background resampling worker for the selected transform and target layer. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose selected transform and target layer should be used. + """ + payload = get_selected_transform_payload(panel) + if payload is None: + panel._set_error("Select a transform to apply.") + return + + moving_layer = panel._selected_layer(panel._transform_target_combo) + if moving_layer is None: + panel._set_error("Select an input layer to transform.") + return + + try: + moving = _get_source_dataarray(moving_layer) + if payload["kind"] == "affine": + transform = get_affine_transform_from_payload(payload) + else: + transform = get_bspline_transform_from_payload(payload) + output_grid = get_output_grid_from_payload(payload) + except Exception as exc: # noqa: BLE001 + panel._set_error(str(exc)) + return + + worker = thread_worker(resample_volume)( + moving, + transform, + shape=output_grid["shape"], + spacing=output_grid["spacing"], + origin=output_grid["origin"], + dims=output_grid["dims"], + interpolation=panel._current_resample_interpolation(), + ) + apply_payload: ApplyTransformPayload = { + "moving_layer_name": moving_layer.name, + "target_layer_name": payload["target_layer_name"], + "transform_source": payload["name"], + "direction": "forward", + } + panel._worker = worker + panel._begin_work() + + worker.returned.connect( + lambda result: on_apply_transform_finished(panel, apply_payload, result) + ) + worker.errored.connect(lambda exc: on_registration_failed(panel, exc)) + worker.finished.connect(panel._end_work) + worker.start() + + +def apply_selected_inverse_transform(panel: "RegistrationPanel") -> None: + """Start a background resampling worker for the inverse of the selected transform. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose selected transform and target layer should be used. + """ + payload = get_selected_transform_payload(panel) + if payload is None: + panel._set_error("Select a transform to apply.") + return + + moving_layer = panel._selected_layer(panel._transform_target_combo) + if moving_layer is None: + panel._set_error("Select an input layer to transform.") + return + + try: + moving = _get_source_dataarray(moving_layer) + if payload["kind"] == "affine": + transform = np.linalg.inv(get_affine_transform_from_payload(payload)) + else: + raise ValueError( + "Inverse apply for B-spline transforms is not available yet." + ) + output_grid = _get_inverse_output_grid(panel, payload) + except Exception as exc: # noqa: BLE001 + panel._set_error(str(exc)) + return + + worker = thread_worker(resample_volume)( + moving, + transform, + shape=output_grid["shape"], + spacing=output_grid["spacing"], + origin=output_grid["origin"], + dims=output_grid["dims"], + interpolation=panel._current_resample_interpolation(), + ) + apply_payload: ApplyTransformPayload = { + "moving_layer_name": moving_layer.name, + "target_layer_name": payload["source_layer_name"], + "transform_source": payload["name"], + "direction": "inverse", + } + panel._worker = worker + panel._begin_work() + + worker.returned.connect( + lambda result: on_apply_transform_finished(panel, apply_payload, result) + ) + worker.errored.connect(lambda exc: on_registration_failed(panel, exc)) + worker.finished.connect(panel._end_work) + worker.start() + + +def on_apply_transform_finished( + panel: "RegistrationPanel", payload: "ApplyTransformPayload", result: xr.DataArray +) -> None: + """Add the finished transformed layer to the viewer and attach apply-transform metadata. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel that initiated the resampling worker. + payload : ApplyTransformPayload + UI snapshot captured when the worker was started; carries the moving and target + layer names together with the source transform label. + result : xarray.DataArray + Resampled DataArray returned by the worker. + """ + registered = result.copy(deep=False) + registered.attrs = registered.attrs.copy() + registered.attrs["registration_operation"] = ( + "apply_inverse_transform" + if payload["direction"] == "inverse" + else "apply_transform" + ) + + name = panel._make_unique_layer_name( + f"{payload['moving_layer_name']} → {payload['target_layer_name']}" + ) + source_layer = panel._get_layer_by_name(payload["moving_layer_name"]) + display_kwargs = ( + _get_image_display_kwargs_from_layer(source_layer) + if source_layer is not None + else {} + ) + contrast_limits = tuple(calc_data_range(registered.data)) + _, layer = plot_napari( + registered, + viewer=panel.viewer, + name=name, + show_colorbar=False, + contrast_limits=contrast_limits, + **display_kwargs, + ) + layer.metadata["xarray"] = registered + layer.metadata["transform_source"] = payload["transform_source"] + layer.metadata["registration_operation"] = registered.attrs[ + "registration_operation" + ] + layer.metadata["registration_parameters"] = payload.copy() + panel.viewer.layers.selection.active = layer + panel._status.hide() + show_info(f"Added transformed layer: {layer.name}") diff --git a/src/confusius/_napari/_registration/_panel_utils.py b/src/confusius/_napari/_registration/_panel_utils.py new file mode 100644 index 00000000..3dd2e976 --- /dev/null +++ b/src/confusius/_napari/_registration/_panel_utils.py @@ -0,0 +1,443 @@ +"""Shared utility helpers for the napari registration panel.""" + +from __future__ import annotations + +from collections.abc import Iterator +from contextlib import contextmanager +from typing import TYPE_CHECKING, Any, Literal, cast + +import numpy as np +import xarray as xr +from qtpy.QtCore import QRegularExpression +from qtpy.QtGui import QValidator +from qtpy.QtWidgets import QDoubleSpinBox, QSizePolicy, QWidget + +from confusius._dims import SPATIAL_DIMS, TIME_DIM +from confusius.xarray.scale import db_scale, power_scale + +if TYPE_CHECKING: + import napari + from napari.layers import Layer + + +@contextmanager +def _preserve_view(viewer: "napari.Viewer") -> Iterator[None]: + """Keep the viewer camera and dims state across a block that adds layers. + + Adding image layers makes napari recompute `camera.center` and re-apply + `napari.imshow`'s default `ndisplay`/`order` to the dims, which yanks the canvas + back to a default framing. Wrapping the layer creation in this context manager + snapshots the current pan, zoom, rotation, and slider position and restores them + once the block exits, so the user keeps the view they were on when starting a + registration run. + + Parameters + ---------- + viewer : napari.Viewer + Viewer whose camera and dims state are snapshotted and restored. + + Yields + ------ + None + Control returns to the wrapped block; the saved state is restored when + it exits, including on early return or exception. + """ + camera = viewer.camera + dims = viewer.dims + center = tuple(camera.center) + zoom = camera.zoom + angles = tuple(camera.angles) + ndisplay = dims.ndisplay + order = tuple(dims.order) + current_step = tuple(dims.current_step) + try: + yield + finally: + dims.ndisplay = ndisplay + dims.order = order + dims.current_step = current_step + camera.center = center + camera.zoom = zoom + camera.angles = angles + + +def _get_default_dims_for_ndim(ndim: int) -> tuple[str, ...]: + """Return fallback dimension names for a raw napari layer. + + Parameters + ---------- + ndim : int + Number of array dimensions. + + Returns + ------- + tuple of str + Default dimension names compatible with ConfUSIus conventions when possible. + """ + defaults: dict[int, tuple[str, ...]] = { + 1: SPATIAL_DIMS[-1:], + 2: SPATIAL_DIMS[-2:], + 3: SPATIAL_DIMS, + 4: (TIME_DIM, *SPATIAL_DIMS), + } + return defaults.get(ndim, tuple(f"dim{i}" for i in range(ndim))) + + +def _normalize_layer_sequence(values: Any, ndim: int, fill: Any) -> list[Any]: + """Return a layer property as a list with length `ndim`. + + Parameters + ---------- + values : Any + Layer property such as `scale`, `translate`, `units`, or `axis_labels`. + ndim : int + Number of dimensions expected on the layer data. + fill : Any + Value used to pad missing entries. + + Returns + ------- + list of Any + Normalized sequence with exactly `ndim` elements. + """ + if values is None: + return [fill] * ndim + seq = list(values) + if len(seq) < ndim: + return ([fill] * (ndim - len(seq))) + seq + if len(seq) > ndim: + return seq[-ndim:] + return seq + + +def _reconstruct_layer_dataarray(layer: "Layer") -> xr.DataArray: + """Reconstruct a DataArray from the current napari layer state. + + Parameters + ---------- + layer : napari.layers.Layer + Napari layer to convert. + + Returns + ------- + xarray.DataArray + DataArray reconstructed from the layer's current axis labels, scale, translate, + and units. + """ + data = np.asarray(layer.data) + ndim = data.ndim + + raw_labels = _normalize_layer_sequence( + getattr(layer, "axis_labels", None), ndim, None + ) + axis_labels = tuple( + str(label) if label not in (None, "") else default + for label, default in zip( + raw_labels, _get_default_dims_for_ndim(ndim), strict=False + ) + ) + + scale = [ + float(v) + for v in _normalize_layer_sequence(getattr(layer, "scale", None), ndim, 1.0) + ] + translate = [ + float(v) + for v in _normalize_layer_sequence(getattr(layer, "translate", None), ndim, 0.0) + ] + raw_units = _normalize_layer_sequence(getattr(layer, "units", None), ndim, None) + units = [None if u is None or str(u) == "pixel" else str(u) for u in raw_units] + + coords: dict[str, xr.DataArray] = {} + for dim, n, spacing, origin, unit in zip( + axis_labels, data.shape, scale, translate, units, strict=False + ): + attrs: dict[str, Any] = {"voxdim": abs(spacing)} + if unit is not None: + attrs["units"] = unit + coords[dim] = xr.DataArray( + origin + np.arange(n) * spacing, dims=[dim], attrs=attrs + ) + + return xr.DataArray(data, dims=axis_labels, coords=coords) + + +def _is_registration_source_layer(layer: "Layer") -> bool: + """Return whether `layer` can be converted to a registration source. + + ConfUSIus-managed layers carry the original `xarray.DataArray` in metadata. For + plain napari image layers we can reconstruct one from eager NumPy data. Lazy + non-NumPy layers (for example the video panel's frame-on-demand array) are + intentionally excluded: forcing `np.asarray` on them can trigger expensive decoding + or backend errors while the registration panel is merely refreshing. + + Parameters + ---------- + layer : napari.layers.Layer + Layer whose registration-source eligibility should be checked. + + Returns + ------- + bool + Whether `layer` can be converted into a registration source. + """ + if layer.metadata.get("xarray") is not None: + return True + if layer.metadata.get("confusius_cached_registration_xarray") is not None: + return True + return isinstance(layer.data, np.ndarray) + + +def _get_source_dataarray(layer: "Layer") -> xr.DataArray: + """Return the stable source DataArray for a napari layer. + + Parameters + ---------- + layer : napari.layers.Layer + Napari layer to convert. + + Returns + ------- + xarray.DataArray + Original ConfUSIus DataArray when present in `layer.metadata`, otherwise a + cached reconstruction captured before later manual napari transforms mutate the + layer pose. + + Raises + ------ + TypeError + If the layer is backed by a lazy non-NumPy array that the registration + panel should ignore. + """ + existing = layer.metadata.get("xarray") + if existing is not None: + return cast("xr.DataArray", existing) + + cached = layer.metadata.get("confusius_cached_registration_xarray") + if cached is not None: + return cast("xr.DataArray", cached) + + if not isinstance(layer.data, np.ndarray): + raise TypeError( + f"Layer {layer.name!r} is not backed by eager NumPy data and cannot be used " + "for registration." + ) + + reconstructed = _reconstruct_layer_dataarray(layer) + layer.metadata["confusius_cached_registration_xarray"] = reconstructed + return reconstructed + + +def _prepare_between_scan_data(data: xr.DataArray) -> xr.DataArray: + """Return a spatial-only DataArray for between-scan registration. + + Parameters + ---------- + data : xarray.DataArray + Input layer data. + + Returns + ------- + xarray.DataArray + Spatial-only data. If the input has a time dimension, it is averaged over time + with attributes preserved. + """ + if TIME_DIM not in data.dims: + return data + averaged = data.mean(dim=TIME_DIM, keep_attrs=True) + averaged.attrs = data.attrs.copy() + return averaged + + +def _apply_registration_scale( + data: xr.DataArray, scale_mode: Literal["off", "dB", "sqrt"] +) -> xr.DataArray: + """Apply optional intensity preprocessing for registration. + + Parameters + ---------- + data : xarray.DataArray + Input data. + scale_mode : {"off", "dB", "sqrt"} + Intensity scaling mode used before registration. + + Returns + ------- + xarray.DataArray + Preprocessed data. + + Raises + ------ + ValueError + If `scale_mode` is not recognized. + """ + if scale_mode == "off": + return data + if scale_mode == "dB": + return db_scale(data) + if scale_mode == "sqrt": + return power_scale(data, exponent=0.5) + raise ValueError(f"Unknown registration scale mode: {scale_mode}.") + + +def _get_image_display_kwargs_from_layer(layer: "Layer") -> dict[str, Any]: + """Return image-display kwargs copied from an existing napari layer. + + Parameters + ---------- + layer : napari.layers.Layer + Source layer whose visual settings should be reused when possible. + + Returns + ------- + dict[str, Any] + Keyword arguments suitable for [`plot_napari`][confusius.plotting.plot_napari]. + """ + kwargs: dict[str, Any] = {} + for attr in ("colormap", "gamma", "opacity"): + if hasattr(layer, attr): + kwargs[attr] = getattr(layer, attr) + return kwargs + + +def _gamma_needs_reset(scale_mode: str) -> bool: + """Return whether registration preview/result gamma should be reset. + + When using intensity scaling, the gamma of the preview and result layers is forced + to 1.0 to avoid double scaling. When scaling is off, the original layer gamma is + preserved. + + Parameters + ---------- + scale_mode : str + Registration intensity scaling mode. + + Returns + ------- + bool + Whether preview/result layers should force `gamma=1.0`. + """ + return scale_mode != "off" + + +def _parse_comma_separated_ints(text: str, expected_len: int = 3) -> tuple[int, ...]: + """Parse comma-separated integers from a text field. + + Parameters + ---------- + text : str + Comma-separated text to parse, e.g. `"1, 2, 3"`. + expected_len : int, default: 3 + Required number of integers in the parsed result. + + Returns + ------- + tuple of int + Parsed integers of length `expected_len` on success, or an empty tuple + when the input is empty, contains a non-numeric value, or does not yield + exactly `expected_len` integers. + """ + parts = [p.strip() for p in text.split(",") if p.strip()] + if not parts: + return tuple() + try: + values = tuple(int(float(p)) for p in parts) + except ValueError: + return tuple() + if len(values) != expected_len: + return tuple() + return values + + +class ScientificDoubleSpinBox(QDoubleSpinBox): + """`QDoubleSpinBox` variant that accepts scientific notation. + + Parameters + ---------- + parent : QWidget, optional + Parent widget. + """ + + _ACCEPTABLE_RE = QRegularExpression( + r"^[+-]?(?:(?:\d+(?:\.\d*)?)|(?:\.\d+))(?:[eE][+-]?\d+)?$" + ) + _INTERMEDIATE_RE = QRegularExpression( + r"^[+-]?(?:(?:\d+(?:\.\d*)?)|(?:\.\d+))?(?:[eE][+-]?\d*)?$" + ) + + def __init__(self, parent: QWidget | None = None) -> None: + super().__init__(parent) + self.setDecimals(10) + self.setKeyboardTracking(False) + self.setAccelerated(True) + self.setSizePolicy(QSizePolicy.Policy.Ignored, QSizePolicy.Policy.Fixed) + + def validate( + self, input: str | None, pos: int + ) -> tuple[QValidator.State, str, int]: + """Validate decimals and scientific notation while the user types. + + Parameters + ---------- + input : str, optional + Current text being edited. + pos : int + Cursor position. + + Returns + ------- + state : QValidator.State + Validation state. + text : str + Normalized text. + pos : int + Cursor position. + """ + normalized = input or "" + if normalized in {"", "+", "-", ".", "+.", "-."}: + return (QValidator.State.Intermediate, normalized, pos) + if self._ACCEPTABLE_RE.match(normalized).hasMatch(): + return (QValidator.State.Acceptable, normalized, pos) + if self._INTERMEDIATE_RE.match(normalized).hasMatch(): + return (QValidator.State.Intermediate, normalized, pos) + return (QValidator.State.Invalid, normalized, pos) + + def valueFromText(self, text: str | None) -> float: + """Parse the current text into a float value. + + Parameters + ---------- + text : str, optional + Text to parse. + + Returns + ------- + float + Parsed numeric value. + """ + return float(text or 0.0) + + def textFromValue(self, v: float) -> str: + """Format values compactly, using scientific notation when helpful. + + Parameters + ---------- + v : float + Value to format. + + Returns + ------- + str + Formatted text. + """ + return f"{v:.12g}" + + def stepBy(self, steps: int) -> None: + """Apply additive stepping using the configured single-step size. + + Parameters + ---------- + steps : int + Number of steps to apply. + """ + self.setValue(self.value() + (steps * self.singleStep())) diff --git a/src/confusius/_napari/_registration/_panel_worker_state.py b/src/confusius/_napari/_registration/_panel_worker_state.py new file mode 100644 index 00000000..a87d6f30 --- /dev/null +++ b/src/confusius/_napari/_registration/_panel_worker_state.py @@ -0,0 +1,31 @@ +"""Worker-state cleanup helpers for the napari registration panel.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from napari.utils.notifications import show_error + +from confusius._napari._registration._panel_progress import ( + teardown_volume_progress, + teardown_volumewise_progress, +) + +if TYPE_CHECKING: + from confusius._napari._registration._panel import RegistrationPanel + + +def on_registration_failed(panel: "RegistrationPanel", exc: BaseException) -> None: + """Handle a failed worker execution. + + Parameters + ---------- + panel : RegistrationPanel + Registration panel whose in-flight state should be cleaned up. + exc : BaseException + Exception raised by the worker. + """ + teardown_volume_progress(panel) + teardown_volumewise_progress(panel, remove_layer=True) + panel._set_error(str(exc)) + show_error(str(exc)) diff --git a/src/confusius/_napari/_registration/_progress.py b/src/confusius/_napari/_registration/_progress.py new file mode 100644 index 00000000..5f1897ff --- /dev/null +++ b/src/confusius/_napari/_registration/_progress.py @@ -0,0 +1,308 @@ +"""Napari-layer-backed progress reporting for `register_volume`. + +This module provides a progress reporter that mirrors the matplotlib-based +[`MatplotlibRegistrationProgressPlotter`][confusius.registration.MatplotlibRegistrationProgressPlotter] but +streams the intermediate resampled volume into a napari Image layer instead of a +matplotlib figure. + +The reporter is intentionally split into two pieces so that napari layers can be +constructed and signal slots connected on the GUI thread before the registration worker +thread starts: + +- [`NapariRegistrationProgressPlotterBridge`][confusius._napari._registration._progress.NapariRegistrationProgressPlotterBridge] + is a lightweight `QObject` that lives on the GUI thread and exposes Qt signals. The + worker thread calls `emit` on it; Qt marshals the slot invocations back to the GUI + thread via an automatically-detected queued connection. +- [`NapariRegistrationProgressPlotter`][confusius._napari._registration._progress.NapariRegistrationProgressPlotter] + implements the [`RegistrationProgress`][confusius.registration.RegistrationProgress] + protocol. It is constructed inside `register_volume` (i.e. on the worker thread) and + resamples the moving image at every iteration using the current tentative transform, + forwarding the resulting array to the bridge. + +Connection lifecycle: + +1. The panel constructs a `NapariRegistrationProgressPlotterBridge` on the GUI thread and connects its + `iterated` signal to a slot that writes the array into the resampled napari layer. +2. The panel builds a factory (via + [`make_napari_progress_factory`][confusius._napari._registration._progress.make_napari_progress_factory]) + that closes over the bridge and returns a `NapariRegistrationProgressPlotter` instance when called + by `register_volume`. +3. `register_volume` instantiates the progress inside the worker thread and wires it to + SimpleITK's iteration and end events as usual. +""" + +from __future__ import annotations + +from threading import Lock +from typing import TYPE_CHECKING, Any, Callable, Literal, cast + +import numpy as np +from qtpy.QtCore import QObject, Signal + +from confusius.registration.progress import _resample_intermediate + +if TYPE_CHECKING: + import SimpleITK as sitk + import xarray as xr + + from confusius.registration import RegistrationDiagnostics, RegistrationProgress + + +class NapariRegistrationProgressPlotterBridge(QObject): + """Thread-boundary signal bridge for napari registration progress. + + Construct this on the GUI thread before starting the registration worker. Connect + `iterated` to a slot that mutates a napari layer (e.g. writes `layer.data = arr`); + the slot will be invoked on the GUI thread thanks to Qt's automatic cross-thread + connection. The bridge itself never touches the napari layer, keeping a clean + separation between the worker's data path and the GUI update path. + + See Also + -------- + NapariRegistrationProgressPlotter : Worker-side reporter that emits via this bridge. + """ + + iterated = Signal(object) + """Emitted at every optimizer iteration with the resampled moving image as a numpy + array in numpy axis order (matching `fixed`).""" + + metric_updated = Signal(float) + """Emitted at every optimizer iteration with the current optimizer metric value (a + float).""" + + finished = Signal() + """Emitted once when the registration end event fires.""" + + +class NapariRegistrationProgressPlotter: + """Napari-layer progress reporter for `register_volume`. + + Implements the [`RegistrationProgress`][confusius.registration.RegistrationProgress] + protocol. Stores the registration method and SimpleITK images it needs to resample + the moving image at each iteration. The resampled array is forwarded to the bridge + via a Qt signal, so this object is safe to call from the SimpleITK command callback + running on the worker thread. + + Parameters + ---------- + bridge : NapariRegistrationProgressPlotterBridge + GUI-thread signal bridge. Stored by reference; never accessed for GUI APIs from + this object. + registration_method : SimpleITK.ImageRegistrationMethod + Active registration method whose `GetInitialTransform` is used to resample the + moving image at each iteration. + fixed_img : SimpleITK.Image + Fixed image defining the resample grid. + moving_img : SimpleITK.Image + Moving image to resample. + plot_metric : bool, default: True + Whether to emit `metric_updated` on each iteration. Kept aligned with the + matplotlib plotter factory signature. + plot_composite : bool, default: True + Kept for signature compatibility with the matplotlib plotter factory. The + napari preview always shows the resampled moving image directly. + resample_kwargs : dict, optional + Extra keyword arguments for the intermediate resample. Supported keys are + `interpolation`, `fill_value`, and `sitk_threads`. + """ + + def __init__( + self, + bridge: NapariRegistrationProgressPlotterBridge, + registration_method: "sitk.ImageRegistrationMethod", + fixed_img: "sitk.Image", + moving_img: "sitk.Image", + *, + plot_metric: bool = True, + plot_composite: bool = True, + resample_kwargs: dict[str, Any] | None = None, + ) -> None: + self._bridge = bridge + self._method = registration_method + self._fixed_img = fixed_img + self._moving_img = moving_img + _kw = dict(resample_kwargs or {}) + self._interpolation = cast( + 'Literal["linear", "nearest", "bspline"]', + _kw.get("interpolation", "linear"), + ) + self._fill_value = float(_kw.get("fill_value", 0.0)) + self._sitk_threads = int(_kw.get("sitk_threads", -1)) + self._plot_metric = plot_metric + del plot_composite + + def update(self) -> None: + """Resample the moving image with the current transform and emit it. + + Called at every SimpleITK iteration event from the worker thread. The + resampled array is sent to the GUI thread via `bridge.iterated`; the + emit is thread-safe and does not require this object to live on the + GUI thread. The current optimizer metric value is also forwarded via + `bridge.metric_updated` so a metric-curve plotter can track + convergence. + """ + import SimpleITK as sitk + + if self._plot_metric: + self._bridge.metric_updated.emit(float(self._method.GetMetricValue())) + + resampled = _resample_intermediate( + self._method, + self._moving_img, + self._fixed_img, + interpolation=self._interpolation, + fill_value=self._fill_value, + sitk_threads=self._sitk_threads, + ) + # .T restores numpy axis order (inverse of the .T used when building + # the SITK image), matching what `register_volume` produces. + arr = np.asarray(sitk.GetArrayFromImage(resampled).T) + self._bridge.iterated.emit(arr) + + def close(self) -> None: + """Signal that the registration run has ended. + + Called at the SimpleITK end event. The final resampled state is + available on the bridge via the last `iterated` payload; the panel is + responsible for retrieving/refreshing the layer from `register_volume`'s + returned DataArray, so this signal is informational (e.g. to stop a + spinner or mark the layer as finalised). + """ + self._bridge.finished.emit() + + +class NapariRegistrationProgressReporterBridge(QObject): + """Thread-boundary signal bridge for volumewise registration progress.""" + + frame_progress = Signal(int, int) + """Emitted with `(completed_frames, total_frames)`.""" + + frame_completed = Signal(int, object) + """Emitted with `(frame_index, registered_frame_array)` when one frame + finishes.""" + + finished = Signal() + """Emitted once when the volumewise run ends.""" + + +class NapariRegistrationProgressReporter: + """Aggregate per-frame progress for `register_volumewise` on the GUI thread. + + Parameters + ---------- + bridge : NapariRegistrationProgressReporterBridge + GUI-thread signal bridge used to forward progress updates. + n_frames : int + Number of frames that will be registered. + """ + + def __init__( + self, + bridge: NapariRegistrationProgressReporterBridge, + *, + n_frames: int, + ) -> None: + self._bridge = bridge + self._n_frames = n_frames + self._completed_frames: set[int] = set() + self._lock = Lock() + + def frame_completed( + self, + frame_index: int, + registered_frame: "xr.DataArray", + diagnostics: "RegistrationDiagnostics", + ) -> None: + """Emit one completed frame for GUI-side layer updates. + + Parameters + ---------- + frame_index : int + Index of the completed frame. + registered_frame : xarray.DataArray + Registered frame data to write into the napari output layer. + diagnostics : confusius.registration.RegistrationDiagnostics + Diagnostics collected for the completed frame. + """ + with self._lock: + del diagnostics + self._completed_frames.add(frame_index) + completed = len(self._completed_frames) + total = self._n_frames + self._bridge.frame_progress.emit(completed, total) + self._bridge.frame_completed.emit( + frame_index, np.asarray(registered_frame.values) + ) + + def close(self) -> None: + """Signal the end of the volumewise run.""" + self._bridge.finished.emit() + + +def make_napari_progress_factory( + bridge: NapariRegistrationProgressPlotterBridge, +) -> "Callable[..., RegistrationProgress]": + """Return a progress-plotter factory bound to a bridge. + + The returned callable has the signature expected by `register_volume`'s + `progress_plotter` argument—it accepts `(registration_method, fixed_img, moving_img, + *, plot_metric, plot_composite, resample_kwargs)` and returns a + [`NapariRegistrationProgressPlotter`][confusius._napari._registration._progress.NapariRegistrationProgressPlotter] + instance wrapping `bridge`. + + Parameters + ---------- + bridge : NapariRegistrationProgressPlotterBridge + GUI-thread bridge the constructed reporter will emit through. + + Returns + ------- + callable + Factory suitable as the `progress_plotter` argument of + [`register_volume`][confusius.registration.register_volume]. + """ + + def factory( + registration_method: "sitk.ImageRegistrationMethod", + fixed_img: "sitk.Image", + moving_img: "sitk.Image", + *, + plot_metric: bool = True, + plot_composite: bool = True, + resample_kwargs: dict[str, Any] | None = None, + ) -> "RegistrationProgress": + """Build a NapariRegistrationProgressPlotter wrapping the captured bridge. + + Parameters + ---------- + registration_method : SimpleITK.ImageRegistrationMethod + Active registration method whose transform is sampled at every iteration. + fixed_img : SimpleITK.Image + Fixed reference image defining the resample grid. + moving_img : SimpleITK.Image + Moving image to resample. + plot_metric : bool, default: True + Whether to emit `metric_updated` on each iteration. + plot_composite : bool, default: True + Kept for signature compatibility with the matplotlib plotter factory. + resample_kwargs : dict, optional + Extra keyword arguments for the intermediate resample. Supported keys are + `interpolation`, `fill_value`, and `sitk_threads`. + + Returns + ------- + RegistrationProgress + Progress reporter ready to be wired to SimpleITK's iteration and + end events by `register_volume`. + """ + return NapariRegistrationProgressPlotter( + bridge, + registration_method, + fixed_img, + moving_img, + plot_metric=plot_metric, + plot_composite=plot_composite, + resample_kwargs=resample_kwargs, + ) + + return factory diff --git a/src/confusius/_napari/_registration/_transform_payloads.py b/src/confusius/_napari/_registration/_transform_payloads.py new file mode 100644 index 00000000..0ab5d772 --- /dev/null +++ b/src/confusius/_napari/_registration/_transform_payloads.py @@ -0,0 +1,583 @@ +"""Transform payload types and serialization helpers for napari registration.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import ( + TYPE_CHECKING, + Literal, + NotRequired, + SupportsFloat, + SupportsIndex, + TypedDict, + cast, +) + +import numpy as np +import numpy.typing as npt +import xarray as xr + +from confusius.registration.bspline import validate_bspline_dataarray + +if TYPE_CHECKING: + from collections.abc import Mapping + + from confusius.registration import RegistrationDiagnostics + + +class TransformDiagnosticsPayload(TypedDict): + """JSON-serializable registration diagnostics summary.""" + + metric: str + final_metric_value: float + n_iterations: int + stop_condition: str + status: str + + +class OutputGridPayload(TypedDict): + """JSON-serializable resampling grid description.""" + + dims: list[str] + shape: list[int] + spacing: list[float] + origin: list[float] + units: list[str | None] + + +class BSplineDataArrayPayload(TypedDict): + """JSON-serializable B-spline control-point DataArray.""" + + dims: list[str] + data: list[object] + coords: dict[str, list[float]] + attrs: dict[str, object] + + +class AffineTransformPayload(TypedDict): + """JSON-serializable affine transform payload used by the napari plugin.""" + + kind: Literal["affine"] + name: str + affine: list[list[float]] + source_layer_name: str + target_layer_name: str + operation: str + transform_model: str + metric: str + output_grid: OutputGridPayload + input_grid: NotRequired[OutputGridPayload] + diagnostics: TransformDiagnosticsPayload + + +class BSplineTransformPayload(TypedDict): + """B-spline transform payload used by the napari plugin.""" + + kind: Literal["bspline"] + name: str + bspline: BSplineDataArrayPayload + source_layer_name: str + target_layer_name: str + operation: str + transform_model: str + metric: str + output_grid: OutputGridPayload + input_grid: NotRequired[OutputGridPayload] + diagnostics: TransformDiagnosticsPayload + + +TransformPayload = AffineTransformPayload | BSplineTransformPayload +"""Union of affine and B-spline transform payloads.""" + + +def make_output_grid_payload(reference: "xr.DataArray") -> OutputGridPayload: + """Return the resampling grid defined by a reference DataArray. + + Parameters + ---------- + reference : xarray.DataArray + Spatial DataArray defining the output grid. + + Returns + ------- + OutputGridPayload + JSON-serializable output-grid description. + """ + dims = [str(dim) for dim in reference.dims] + return { + "dims": dims, + "shape": [int(reference.sizes[dim]) for dim in dims], + "spacing": [float(reference.fusi.spacing[dim]) for dim in dims], + "origin": [float(reference.fusi.origin[dim]) for dim in dims], + "units": [ + cast("str | None", reference.coords[dim].attrs.get("units")) + if dim in reference.coords + else None + for dim in dims + ], + } + + +def _make_diagnostics_payload( + diagnostics: "RegistrationDiagnostics", +) -> TransformDiagnosticsPayload: + """Return a JSON-serializable diagnostics summary. + + Parameters + ---------- + diagnostics : confusius.registration.RegistrationDiagnostics + Per-call registration diagnostics to serialize. + + Returns + ------- + TransformDiagnosticsPayload + JSON-serializable diagnostics summary. + """ + return { + "metric": diagnostics.metric, + "final_metric_value": float(diagnostics.final_metric_value), + "n_iterations": int(diagnostics.n_iterations), + "stop_condition": diagnostics.stop_condition, + "status": diagnostics.status, + } + + +def make_affine_transform_payload( + affine: npt.NDArray[np.floating], + *, + reference: "xr.DataArray", + source: "xr.DataArray | None" = None, + source_layer_name: str, + target_layer_name: str, + operation: str, + transform_model: str, + metric: str, + diagnostics: "RegistrationDiagnostics", + name: str | None = None, +) -> AffineTransformPayload: + """Build a JSON-serializable payload for a registered affine transform. + + Parameters + ---------- + affine : (N+1, N+1) numpy.ndarray + Affine transform in homogeneous coordinates. + reference : xarray.DataArray + Fixed/reference DataArray defining the output resampling grid. + source : xarray.DataArray, optional + Original moving/source DataArray defining the inverse-apply resampling grid. If + not provided, `input_grid` is omitted from the payload. + source_layer_name : str + Name of the moving/source layer used when estimating the transform. + target_layer_name : str + Name of the fixed/target layer used when estimating the transform. + operation : str + Registration operation that produced the transform. + transform_model : str + Transform model used during registration. + metric : str + Similarity metric used during registration. + diagnostics : confusius.registration.RegistrationDiagnostics + Per-call registration diagnostics. + name : str, optional + Human-friendly transform name. If not provided, a default name is generated. + + Returns + ------- + AffineTransformPayload + JSON-serializable affine transform payload. + """ + affine = np.asarray(affine, dtype=float) + payload_name = ( + name or f"{source_layer_name} → {target_layer_name} ({transform_model})" + ) + payload: AffineTransformPayload = { + "kind": "affine", + "name": payload_name, + "affine": affine.tolist(), + "source_layer_name": source_layer_name, + "target_layer_name": target_layer_name, + "operation": operation, + "transform_model": transform_model, + "metric": metric, + "output_grid": make_output_grid_payload(reference), + "diagnostics": _make_diagnostics_payload(diagnostics), + } + if source is not None: + payload["input_grid"] = make_output_grid_payload(source) + return payload + + +def _serialize_bspline_dataarray(transform: "xr.DataArray") -> BSplineDataArrayPayload: + """Return a JSON-serializable B-spline DataArray payload. + + Parameters + ---------- + transform : xarray.DataArray + B-spline control-point grid to serialize. + + Returns + ------- + BSplineDataArrayPayload + JSON-serializable B-spline DataArray payload. + """ + validate_bspline_dataarray(transform) + return { + "dims": [str(dim) for dim in transform.dims], + "data": np.asarray(transform, dtype=float).tolist(), + "coords": { + str(dim): np.asarray(transform.coords[dim], dtype=float).tolist() + for dim in transform.dims + if dim in transform.coords + }, + "attrs": json.loads(json.dumps(transform.attrs)), + } + + +def _deserialize_bspline_dataarray(payload: BSplineDataArrayPayload) -> xr.DataArray: + """Reconstruct a B-spline DataArray from its JSON payload. + + Parameters + ---------- + payload : BSplineDataArrayPayload + JSON payload describing a B-spline control-point grid. + + Returns + ------- + xarray.DataArray + Reconstructed B-spline control-point grid. + """ + dims = [str(dim) for dim in payload["dims"]] + coords = { + str(dim): xr.DataArray(np.asarray(values, dtype=float), dims=[str(dim)]) + for dim, values in payload["coords"].items() + } + transform = xr.DataArray( + np.asarray(payload["data"], dtype=float), + dims=dims, + coords=coords, + attrs=dict(payload["attrs"]), + ) + validate_bspline_dataarray(transform) + return transform + + +def make_bspline_transform_payload( + transform: "xr.DataArray", + *, + reference: "xr.DataArray", + source: "xr.DataArray | None" = None, + source_layer_name: str, + target_layer_name: str, + operation: str, + transform_model: str, + metric: str, + diagnostics: "RegistrationDiagnostics", + name: str | None = None, +) -> BSplineTransformPayload: + """Build a JSON-serializable payload for a registered B-spline transform. + + Parameters + ---------- + transform : xarray.DataArray + B-spline control-point grid. + reference : xarray.DataArray + Fixed/reference DataArray defining the output resampling grid. + source : xarray.DataArray, optional + Original moving/source DataArray defining the inverse-apply resampling grid. If + not provided, `input_grid` is omitted from the payload. + source_layer_name : str + Name of the moving/source layer used when estimating the transform. + target_layer_name : str + Name of the fixed/target layer used when estimating the transform. + operation : str + Registration operation that produced the transform. + transform_model : str + Transform model used during registration. + metric : str + Similarity metric used during registration. + diagnostics : confusius.registration.RegistrationDiagnostics + Per-call registration diagnostics. + name : str, optional + Human-friendly transform name. If not provided, a default name is generated. + + Returns + ------- + BSplineTransformPayload + JSON-serializable B-spline transform payload. + """ + payload_name = ( + name or f"{source_layer_name} → {target_layer_name} ({transform_model})" + ) + payload: BSplineTransformPayload = { + "kind": "bspline", + "name": payload_name, + "bspline": _serialize_bspline_dataarray(transform), + "source_layer_name": source_layer_name, + "target_layer_name": target_layer_name, + "operation": operation, + "transform_model": transform_model, + "metric": metric, + "output_grid": make_output_grid_payload(reference), + "diagnostics": _make_diagnostics_payload(diagnostics), + } + if source is not None: + payload["input_grid"] = make_output_grid_payload(source) + return payload + + +def get_affine_transform_from_payload( + payload: "Mapping[str, object]", +) -> npt.NDArray[np.float64]: + """Return the affine matrix stored in a payload. + + Parameters + ---------- + payload : mapping + Transform payload loaded from metadata or disk. + + Returns + ------- + (N+1, N+1) numpy.ndarray + Affine matrix. + """ + if payload.get("kind") != "affine": + raise ValueError("Transform payload is not an affine transform.") + + affine = np.asarray(payload.get("affine"), dtype=float) + if affine.ndim != 2 or affine.shape[0] != affine.shape[1] or affine.shape[0] < 3: + raise ValueError( + "Affine payload must contain a square homogeneous matrix of shape " + "(N+1, N+1)." + ) + return affine + + +def get_bspline_transform_from_payload(payload: "Mapping[str, object]") -> xr.DataArray: + """Return the B-spline transform stored in a payload. + + Parameters + ---------- + payload : mapping + Transform payload loaded from metadata or disk. + + Returns + ------- + xarray.DataArray + B-spline control-point grid. + """ + if payload.get("kind") != "bspline": + raise ValueError("Transform payload is not a B-spline transform.") + + bspline = payload.get("bspline") + if not isinstance(bspline, dict): + raise ValueError("B-spline payload must contain a serialized DataArray.") + return _deserialize_bspline_dataarray(cast("BSplineDataArrayPayload", bspline)) + + +def _coerce_grid_payload( + grid: object, *, field_name: str, missing_message: str +) -> OutputGridPayload: + """Return a validated grid payload from a raw mapping field.""" + if not isinstance(grid, dict): + raise ValueError(missing_message) + + grid_dict = cast("dict[str, object]", grid) + dims = grid_dict.get("dims") + shape = grid_dict.get("shape") + spacing = grid_dict.get("spacing") + origin = grid_dict.get("origin") + units = grid_dict.get("units") + if not all(isinstance(v, list) for v in (dims, shape, spacing, origin, units)): + raise ValueError(f"Transform payload {field_name} is malformed.") + + dims_list = cast("list[object]", dims) + shape_list = cast("list[SupportsIndex]", shape) + spacing_list = cast("list[SupportsFloat]", spacing) + origin_list = cast("list[SupportsFloat]", origin) + units_list = cast("list[object]", units) + + return { + "dims": [str(v) for v in dims_list], + "shape": [int(v) for v in shape_list], + "spacing": [float(v) for v in spacing_list], + "origin": [float(v) for v in origin_list], + "units": [None if v is None else str(v) for v in units_list], + } + + +def get_output_grid_from_payload(payload: "Mapping[str, object]") -> OutputGridPayload: + """Return the output grid stored in a transform payload. + + Parameters + ---------- + payload : mapping + Transform payload loaded from metadata or disk. + + Returns + ------- + OutputGridPayload + Output-grid description stored in the payload. + """ + return _coerce_grid_payload( + payload.get("output_grid"), + field_name="output grid", + missing_message="Transform payload does not contain an output grid.", + ) + + +def get_input_grid_from_payload( + payload: "Mapping[str, object]", +) -> OutputGridPayload | None: + """Return the input grid stored in a transform payload, if present. + + Parameters + ---------- + payload : mapping + Transform payload loaded from metadata or disk. + + Returns + ------- + OutputGridPayload or None + Input-grid description stored in the payload, or `None` when the payload does + not carry one. + """ + if "input_grid" not in payload: + return None + return _coerce_grid_payload( + payload.get("input_grid"), + field_name="input grid", + missing_message="Transform payload does not contain an input grid.", + ) + + +def _save_bspline_transform_payload( + path: str | Path, payload: BSplineTransformPayload +) -> None: + """Save a B-spline transform payload as Zarr. + + Parameters + ---------- + path : str or pathlib.Path + Output Zarr path. + payload : BSplineTransformPayload + Transform payload to save. + + Raises + ------ + ValueError + If `path` does not have a `.zarr` extension. + """ + path = Path(path) + if path.suffix != ".zarr": + raise ValueError("B-spline transform files must have .zarr extension.") + + transform = get_bspline_transform_from_payload(payload) + ds = transform.to_dataset(name="bspline_transform") + payload_metadata = { + key: value for key, value in payload.items() if key not in {"kind", "bspline"} + } + ds.attrs["confusius_transform_kind"] = "bspline" + ds.attrs["confusius_transform_payload_json"] = json.dumps(payload_metadata) + ds.to_zarr(path, mode="w") + + +def _load_bspline_transform_payload(path: str | Path) -> BSplineTransformPayload: + """Load a B-spline transform payload from Zarr. + + Parameters + ---------- + path : str or pathlib.Path + Input Zarr path. + + Returns + ------- + BSplineTransformPayload + Loaded B-spline transform payload. + """ + ds = xr.open_zarr(path) + try: + if ds.attrs.get("confusius_transform_kind") != "bspline": + raise ValueError( + "Zarr transform store does not contain a ConfUSIus B-spline transform." + ) + payload_metadata = json.loads( + cast("str", ds.attrs["confusius_transform_payload_json"]) + ) + if not isinstance(payload_metadata, dict): + raise ValueError("Stored transform payload metadata is malformed.") + transform = ds["bspline_transform"].load() + finally: + ds.close() + + validate_bspline_dataarray(transform) + payload: BSplineTransformPayload = { + "kind": "bspline", + "bspline": _serialize_bspline_dataarray(transform), + "name": str(payload_metadata["name"]), + "source_layer_name": str(payload_metadata["source_layer_name"]), + "target_layer_name": str(payload_metadata["target_layer_name"]), + "operation": str(payload_metadata["operation"]), + "transform_model": str(payload_metadata["transform_model"]), + "metric": str(payload_metadata["metric"]), + "output_grid": get_output_grid_from_payload(payload_metadata), + "diagnostics": cast( + "TransformDiagnosticsPayload", payload_metadata["diagnostics"] + ), + } + input_grid = get_input_grid_from_payload(payload_metadata) + if input_grid is not None: + payload["input_grid"] = input_grid + return payload + + +def save_transform_payload(path: str | Path, payload: TransformPayload) -> None: + """Save a transform payload to disk as JSON for affine payloads or Zarr for B-spline payloads. + + Parameters + ---------- + path : str or pathlib.Path + Output path. + payload : TransformPayload + Transform payload to save. + + Notes + ----- + Affine payloads are saved as JSON. B-spline payloads are saved as Zarr. + """ + if payload["kind"] == "affine": + Path(path).write_text(json.dumps(payload, indent=2) + "\n") + return + _save_bspline_transform_payload(path, payload) + + +def load_transform_payload(path: str | Path) -> TransformPayload: + """Load an affine or B-spline transform payload from disk. + + Parameters + ---------- + path : str or pathlib.Path + Input path. + + Returns + ------- + TransformPayload + Loaded transform payload. + """ + path = Path(path) + if path.suffix == ".zarr": + return _load_bspline_transform_payload(path) + + payload = json.loads(path.read_text()) + if not isinstance(payload, dict): + raise ValueError("Transform file must contain a JSON object.") + + kind = payload.get("kind") + if kind != "affine": + raise ValueError( + "JSON transform files currently support affine payloads only. " + "Use .zarr for B-spline transforms." + ) + get_affine_transform_from_payload(payload) + get_output_grid_from_payload(payload) + return cast("TransformPayload", payload) diff --git a/src/confusius/_napari/_signals/_panel.py b/src/confusius/_napari/_signals/_panel.py index a94ed65c..3a4e1b7d 100644 --- a/src/confusius/_napari/_signals/_panel.py +++ b/src/confusius/_napari/_signals/_panel.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any, cast import napari from qtpy.QtCore import Qt, QTimer @@ -12,17 +12,19 @@ QComboBox, QDockWidget, QDoubleSpinBox, + QGridLayout, QGroupBox, QHBoxLayout, QLabel, - QMainWindow, QPushButton, QRadioButton, + QSizePolicy, QVBoxLayout, QWidget, ) from confusius._dims import SPATIAL_DIMS_WITH_POSE, TIME_DIM +from confusius._napari._qt import find_main_window from confusius._napari._signals._manager import SignalsManagerDialog from confusius._napari._signals._plotter import SignalPlotter from confusius._napari._signals._store import SignalStore @@ -64,11 +66,14 @@ def _setup_ui(self) -> None: layout.setContentsMargins(10, 10, 10, 10) layout.setSpacing(8) - # Source group. + # Source group. A grid keeps the three combos left-aligned in one + # column regardless of the leading radio/label text widths. source_group = QGroupBox("Source") self._source_group = source_group - source_layout = QVBoxLayout(source_group) - source_layout.setSpacing(4) + source_layout = QGridLayout(source_group) + source_layout.setHorizontalSpacing(6) + source_layout.setVerticalSpacing(4) + source_layout.setColumnStretch(1, 1) self._source_btn_group = QButtonGroup(self) @@ -76,20 +81,19 @@ def _setup_ui(self) -> None: self._radio_mouse = QRadioButton("Mouse (Shift + hover)") self._radio_mouse.setChecked(True) self._source_btn_group.addButton(self._radio_mouse, 0) - source_layout.addWidget(self._radio_mouse) + source_layout.addWidget(self._radio_mouse, 0, 0, 1, 3) # Points row. Text is part of the radio button (same pattern as the Mouse row) # so the indicator and label are always flush with no gap. - points_row = QHBoxLayout() - self._radio_points = QRadioButton("Points:") + self._radio_points = QRadioButton("Points") self._source_btn_group.addButton(self._radio_points, 1) - points_row.addWidget(self._radio_points) + source_layout.addWidget(self._radio_points, 1, 0) self._points_combo = QComboBox() self._points_combo.setSizeAdjustPolicy( QComboBox.SizeAdjustPolicy.AdjustToMinimumContentsLengthWithIcon ) self._points_combo.setEnabled(False) - points_row.addWidget(self._points_combo, stretch=1) + source_layout.addWidget(self._points_combo, 1, 1) self._new_points_btn = QPushButton("+") self._new_points_btn.setStyleSheet("font-weight: bold; font-size: 14px;") self._new_points_btn.setToolTip( @@ -97,20 +101,18 @@ def _setup_ui(self) -> None: "Points will be visible at all time steps." ) self._new_points_btn.clicked.connect(self._create_points_layer) - points_row.addWidget(self._new_points_btn) - source_layout.addLayout(points_row) + source_layout.addWidget(self._new_points_btn, 1, 2) # Labels row. - labels_row = QHBoxLayout() - self._radio_labels = QRadioButton("Labels:") + self._radio_labels = QRadioButton("Labels") self._source_btn_group.addButton(self._radio_labels, 2) - labels_row.addWidget(self._radio_labels) + source_layout.addWidget(self._radio_labels, 2, 0) self._labels_combo = QComboBox() self._labels_combo.setSizeAdjustPolicy( QComboBox.SizeAdjustPolicy.AdjustToMinimumContentsLengthWithIcon ) self._labels_combo.setEnabled(False) - labels_row.addWidget(self._labels_combo, stretch=1) + source_layout.addWidget(self._labels_combo, 2, 1) self._new_labels_btn = QPushButton("+") self._new_labels_btn.setStyleSheet("font-weight: bold; font-size: 14px;") self._new_labels_btn.setToolTip( @@ -118,21 +120,19 @@ def _setup_ui(self) -> None: "Labels will be visible at all time steps." ) self._new_labels_btn.clicked.connect(self._create_labels_layer) - labels_row.addWidget(self._new_labels_btn) - source_layout.addLayout(labels_row) + source_layout.addWidget(self._new_labels_btn, 2, 2) - # Reference image (enabled in points/labels mode). - ref_row = QHBoxLayout() - self._ref_label = QLabel("Reference:") + # Reference image (enabled in points/labels mode). Spans the "+" + # column so its right edge stays flush with the buttons above. + self._ref_label = QLabel("Reference") self._ref_label.setEnabled(False) - ref_row.addWidget(self._ref_label) + source_layout.addWidget(self._ref_label, 3, 0) self._ref_combo = QComboBox() self._ref_combo.setSizeAdjustPolicy( QComboBox.SizeAdjustPolicy.AdjustToMinimumContentsLengthWithIcon ) self._ref_combo.setEnabled(False) - ref_row.addWidget(self._ref_combo, stretch=1) - source_layout.addLayout(ref_row) + source_layout.addWidget(self._ref_combo, 3, 1, 1, 2) layout.addWidget(source_group) @@ -156,7 +156,7 @@ def _setup_ui(self) -> None: # X-axis dimension selection. xaxis_row = QHBoxLayout() - xaxis_label = QLabel("x-axis:") + xaxis_label = QLabel("x-axis") xaxis_label.setTextFormat(Qt.TextFormat.RichText) xaxis_row.addWidget(xaxis_label) self._xaxis_combo = QComboBox() @@ -171,25 +171,6 @@ def _setup_ui(self) -> None: xaxis_row.addWidget(self._xaxis_combo, stretch=1) axis_layout.addLayout(xaxis_row) - spinbox: list[QDoubleSpinBox] = [] - for lim in ("min", "max"): - ylim_layout = QHBoxLayout() - ylim_label = QLabel(f"y {lim}:") - ylim_label.setTextFormat(Qt.TextFormat.RichText) - ylim_layout.addWidget(ylim_label) - spin = QDoubleSpinBox() - spin.setObjectName(f"y{lim}_spin") - spin.setAlignment(Qt.AlignmentFlag.AlignCenter) - spin.setRange(-1e9, 1e9) - spin.setValue(-1.0 if lim == "min" else 1.0) - spin.valueChanged.connect(self._apply_settings) - spinbox.append(spin) - ylim_layout.addWidget(spin) - - axis_layout.addLayout(ylim_layout) - - self._ymin_spin, self._ymax_spin = spinbox - # Autoscale checkbox. QCheckBox does not support rich text, so we pair a # text-less checkbox with a clickable QLabel to get the italic "y". autoscale_row = QHBoxLayout() @@ -198,12 +179,38 @@ def _setup_ui(self) -> None: self._autoscale_check.toggled.connect(self._on_autoscale_changed) autoscale_label = QLabel("Autoscale y-axis") autoscale_label.setTextFormat(Qt.TextFormat.RichText) - autoscale_label.mousePressEvent = lambda _e: self._autoscale_check.toggle() # type: ignore[method-assign] + setattr( + cast("Any", autoscale_label), + "mousePressEvent", + lambda _e: self._autoscale_check.toggle(), + ) autoscale_row.addWidget(self._autoscale_check) autoscale_row.addWidget(autoscale_label) autoscale_row.addStretch() axis_layout.addLayout(autoscale_row) + yminmax_row = QHBoxLayout() + spinbox: list[QDoubleSpinBox] = [] + for lim in ("min", "max"): + ylim_label = QLabel(f"y {lim}") + ylim_label.setTextFormat(Qt.TextFormat.RichText) + yminmax_row.addWidget(ylim_label) + spin = QDoubleSpinBox() + spin.setObjectName(f"y{lim}_spin") + spin.setAlignment(Qt.AlignmentFlag.AlignCenter) + # Ignored horizontal policy: the minimum size hint spans the + # widest possible value ("-1000000000.00"), which overflowed the + # panel in issue #183. The layout stretch shares the row width. + spin.setSizePolicy(QSizePolicy.Policy.Ignored, QSizePolicy.Policy.Fixed) + spin.setRange(-1e9, 1e9) + spin.setValue(-1.0 if lim == "min" else 1.0) + spin.valueChanged.connect(self._apply_settings) + spinbox.append(spin) + yminmax_row.addWidget(spin, stretch=1) + axis_layout.addLayout(yminmax_row) + + self._ymin_spin, self._ymax_spin = spinbox + # Apply initial autoscale state so spinboxes start disabled. self._on_autoscale_changed(True) @@ -293,7 +300,7 @@ def _ensure_plotter(self) -> SignalPlotter: # before the canvas first paints. This mirrors the pattern used in the QC # panel and prevents the HiDPI click-offset bug. def _settle_layout() -> None: - main_win = self._find_main_window(dock) + main_win = find_main_window(dock) if main_win is None: return # Zero minimum sizes on the central widget and all its children so the @@ -440,7 +447,7 @@ def _on_xaxis_step_changed(self, event) -> None: def _on_frame_clicked(self, frame: float) -> None: """Navigate the viewer to the clicked x-axis coordinate. - ``frame`` is the x-axis plot value (a world coordinate, e.g. time + `frame` is the x-axis plot value (a world coordinate, e.g. time in seconds). Using `dims.set_point` avoids the double-conversion bug that occurs when setting `current_step` directly — the step index depends on `dims.range.step`, which changes when a video @@ -457,26 +464,6 @@ def _on_theme_changed(self) -> None: if self._signals_manager is not None: self._signals_manager.apply_theme(self._viewer.theme) - def _find_main_window(self, widget: QWidget) -> QMainWindow | None: - """Traverse up the widget hierarchy to find the QMainWindow. - - Parameters - ---------- - widget : QWidget - Starting widget to search from. - - Returns - ------- - QMainWindow | None - The main window if found, None otherwise. - """ - parent = widget.parent() - while parent is not None: - if isinstance(parent, QMainWindow): - return parent - parent = parent.parent() - return None - # ------------------------------------------------------------------ # Source management # ------------------------------------------------------------------ diff --git a/src/confusius/_napari/_signals/_plotter.py b/src/confusius/_napari/_signals/_plotter.py index bb6f85e3..ae1c8d19 100644 --- a/src/confusius/_napari/_signals/_plotter.py +++ b/src/confusius/_napari/_signals/_plotter.py @@ -475,7 +475,7 @@ def _on_layer_change(self, event) -> None: """Handle layer insertion/removal or active-layer change events. When the newly active layer is not a valid signal source (e.g., a - video layer without xarray metadata), ``_current_layer`` is left + video layer without xarray metadata), `_current_layer` is left unchanged so that the plotter continues to reference the previous valid layer for signal extraction, cursor mapping, and click-to-navigate. @@ -843,14 +843,16 @@ def _extract_signals(self, layer, cursor_pos: np.ndarray) -> np.ndarray | None: Always uses the nearest voxel to the cursor position. """ data = layer.data - ind = list(int(round(x)) for x in layer.world_to_data(cursor_pos)) + ind: list[int | slice] = [ + int(round(x)) for x in layer.world_to_data(cursor_pos) + ] xaxis_index = self._xaxis_dim_index(layer) # Replace the x-axis index before bounds-checking: the injected x-axis world # coordinate (typically 0) may fall outside the data range (e.g. when the # coordinate starts at a non-zero offset), which would cause the check to # reject valid spatial positions. - ind[xaxis_index] = slice(None) # type: ignore[call-overload] + ind[xaxis_index] = slice(None) if not all( 0 <= i < max_i for i, max_i in zip(ind, data.shape) if isinstance(i, int) diff --git a/src/confusius/_napari/_signals/_store.py b/src/confusius/_napari/_signals/_store.py index 415e6ef0..2596dced 100644 --- a/src/confusius/_napari/_signals/_store.py +++ b/src/confusius/_napari/_signals/_store.py @@ -65,17 +65,17 @@ class LiveSignal: Attributes ---------- id : str - Stable identifier (e.g. ``"mouse-0"``, ``"point-3"``, ``"label-5"``). + Stable identifier (e.g. `"mouse-0"`, `"point-3"`, `"label-5"`). name : str Display name used in legends (editable by the user). color : str Hex color for the plot line. visible : bool Whether the signal should be plotted. - source_type : ``"mouse"`` | ``"point"`` | ``"label"`` + source_type : `"mouse"` | `"point"` | `"label"` Kind of napari source that produces this signal. source_id : int | None - ``None`` for mouse, point index for points, label integer for labels. + `None` for mouse, point index for points, label integer for labels. """ id: str @@ -320,7 +320,7 @@ def get_live_signal(self, signal_id: str) -> LiveSignal | None: Returns ------- LiveSignal | None - The signal, or ``None`` if not found. + The signal, or `None` if not found. """ return self._live_signals.get(signal_id) diff --git a/src/confusius/_napari/_tour.py b/src/confusius/_napari/_tour.py index f7bc8f9f..50bb2b4c 100644 --- a/src/confusius/_napari/_tour.py +++ b/src/confusius/_napari/_tour.py @@ -12,7 +12,9 @@ from qtpy.QtCore import QEvent, QObject, QPoint, QRect, Qt, QTimer, Signal from qtpy.QtGui import QColor, QFont, QPainter, QPen from qtpy.QtWidgets import ( + QAbstractButton, QDockWidget, + QGroupBox, QHBoxLayout, QLabel, QPushButton, @@ -546,6 +548,7 @@ def build_default_tour( from confusius._napari._data._save_panel import SavePanel from confusius._napari._events._panel import EventPanel from confusius._napari._qc._panel import QCPanel + from confusius._napari._registration._panel import RegistrationPanel from confusius._napari._signals._panel import SignalPanel from confusius._napari._video._video_panel import VideoPanel @@ -643,6 +646,36 @@ def _find() -> QRect | None: return _find + def _panel_label( + label: str, + widget_type: type[QWidget], + text: str, + ) -> Callable[[], QWidget | None]: + def _find() -> QWidget | None: + panel = _panel_descendant(label, widget_type)() + if panel is None: + return None + for child in panel.findChildren(QLabel): + if child.text() == text: + return child + return None + + return _find + + def _panel_attr_ancestor_rect( + label: str, + widget_type: type[QWidget], + attr: str, + ancestor_type: type[QWidget], + ) -> Callable[[], QRect | None]: + def _find() -> QRect | None: + widget = _panel_attr(label, widget_type, attr)() + while widget is not None and not isinstance(widget, ancestor_type): + widget = widget.parentWidget() + return _widget_rect(widget) + + return _find + def _expand_section(label: str) -> Callable[[], None]: def _action() -> None: for btn, _icon in getattr(plugin_widget, "_accordion_btns", []): @@ -652,6 +685,37 @@ def _action() -> None: return _action + def _run_actions(*actions: Callable[[], None]) -> Callable[[], None]: + def _action() -> None: + for step_action in actions: + step_action() + + return _action + + def _registration_panel() -> RegistrationPanel | None: + panel = _panel_descendant("Registration", RegistrationPanel)() + return panel if isinstance(panel, RegistrationPanel) else None + + def _set_panel_button( + label: str, + attr: str, + *, + checked: bool = True, + ) -> Callable[[], None]: + def _action() -> None: + _expand_section(label)() + panel = _registration_panel() + if panel is None: + return + button = getattr(panel, attr, None) + if isinstance(button, QAbstractButton) and button.isChecked() != checked: + button.click() + + return _action + + def _select_sub_panel(label: str, attr: str) -> Callable[[], None]: + return _set_panel_button(label, attr) + # Record the open panel before the tour starts so we can restore it when # the tour is closed or skipped. initial_open: str | None = next( @@ -663,7 +727,36 @@ def _action() -> None: None, ) + registration_panel = _registration_panel() + initial_registration_state = None + if registration_panel is not None: + initial_registration_state = { + "register_panel": registration_panel._register_panel_radio.isChecked(), + "single_volume": registration_panel._single_volume_radio.isChecked(), + "advanced_open": registration_panel._advanced_toggle.isChecked(), + } + def _restore_state() -> None: + if initial_registration_state is not None: + _expand_section("Registration")() + panel = _registration_panel() + if panel is not None: + if ( + panel._register_panel_radio.isChecked() + != initial_registration_state["register_panel"] + ): + panel._transforms_panel_radio.click() + if ( + panel._single_volume_radio.isChecked() + != initial_registration_state["single_volume"] + ): + panel._time_series_radio.click() + if ( + panel._advanced_toggle.isChecked() + != initial_registration_state["advanced_open"] + ): + panel._advanced_toggle.click() + if initial_open is None: return for btn, _ in getattr(plugin_widget, "_accordion_btns", []): @@ -859,6 +952,205 @@ def _restore_state() -> None: tooltip_target=_dock_widget, pre_action=_expand_section("Signals"), ), + TourStep( + target=_accordion_panel("Registration"), + title="Registration", + body=( + "Use this section to align fUSI scans. The Register sub-panel " + "runs new registrations; the Transforms sub-panel saves, " + "loads, and reapplies them." + ), + anchor="left", + spotlight_rect=_accordion_tab_rect("Registration"), + tooltip_target=_dock_widget, + pre_action=_expand_section("Registration"), + ), + TourStep( + target=_panel_attr( + "Registration", RegistrationPanel, "_register_panel_radio" + ), + title="Register and Transforms", + body=( + "Use Register to compute a new alignment. Switch to " + "Transforms to reuse, save, load, or apply transforms." + ), + anchor="left", + spotlight_rect=_panel_attr_rect( + "Registration", + RegistrationPanel, + "_register_panel_radio", + "_transforms_panel_radio", + ), + tooltip_target=_dock_widget, + pre_action=_select_sub_panel("Registration", "_register_panel_radio"), + ), + TourStep( + target=_panel_attr( + "Registration", RegistrationPanel, "_single_volume_radio" + ), + title="Between scans and Within-scan", + body=( + "Pick Between scans to align one layer to another, or " + "Within-scan to motion-correct a time series against one " + "reference volume." + ), + anchor="left", + spotlight_rect=lambda: _united_rect( + _panel_label("Registration", RegistrationPanel, "Mode")(), + _panel_attr( + "Registration", RegistrationPanel, "_single_volume_radio" + )(), + _panel_attr("Registration", RegistrationPanel, "_time_series_radio")(), + ), + tooltip_target=_dock_widget, + pre_action=_run_actions( + _select_sub_panel("Registration", "_register_panel_radio"), + _set_panel_button("Registration", "_single_volume_radio"), + _set_panel_button("Registration", "_advanced_toggle", checked=False), + ), + ), + TourStep( + target=_panel_attr("Registration", RegistrationPanel, "_moving_combo"), + title="Moving, Masks, and Fixed Layers", + body=( + "For Between scans, choose the Moving layer to align, " + "optional Moving/Fixed mask label layers to restrict the " + "metric, and the Fixed layer that defines the target space." + ), + anchor="left", + spotlight_rect=_panel_attr_rect( + "Registration", + RegistrationPanel, + "_moving_label", + "_moving_combo", + "_moving_mask_label", + "_moving_mask_row", + "_fixed_label", + "_fixed_combo", + "_fixed_mask_label", + "_fixed_mask_row", + ), + tooltip_target=_dock_widget, + pre_action=_run_actions( + _select_sub_panel("Registration", "_register_panel_radio"), + _set_panel_button("Registration", "_single_volume_radio"), + _set_panel_button("Registration", "_advanced_toggle", checked=False), + ), + ), + TourStep( + target=_panel_attr("Registration", RegistrationPanel, "_transform_combo"), + title="Parameters", + body=( + "This box contains the registration settings: transform, metric, " + "scale, initialization, and the main optimizer controls." + ), + anchor="left", + spotlight_rect=_panel_attr_ancestor_rect( + "Registration", + RegistrationPanel, + "_transform_combo", + QGroupBox, + ), + tooltip_target=_dock_widget, + pre_action=_run_actions( + _select_sub_panel("Registration", "_register_panel_radio"), + _set_panel_button("Registration", "_single_volume_radio"), + _set_panel_button("Registration", "_advanced_toggle", checked=False), + ), + ), + TourStep( + target=_panel_attr("Registration", RegistrationPanel, "_advanced_toggle"), + title="Advanced Parameters", + body=( + "If you need finer control, open Advanced for extra optimizer, " + "multi-resolution, metric, and resampling settings." + ), + anchor="left", + spotlight_rect=_panel_attr_rect( + "Registration", + RegistrationPanel, + "_advanced_toggle", + ), + tooltip_target=_dock_widget, + pre_action=_run_actions( + _select_sub_panel("Registration", "_register_panel_radio"), + _set_panel_button("Registration", "_single_volume_radio"), + _set_panel_button("Registration", "_advanced_toggle", checked=False), + ), + ), + TourStep( + target=_panel_attr("Registration", RegistrationPanel, "_run_btn"), + title="Run Registration", + body=( + "Click Run registration to start. ConfUSIus shows progress " + "while the alignment runs and adds the result as a new layer." + ), + anchor="left", + tooltip_target=_dock_widget, + pre_action=_run_actions( + _select_sub_panel("Registration", "_register_panel_radio"), + _set_panel_button("Registration", "_single_volume_radio"), + _set_panel_button("Registration", "_advanced_toggle", checked=False), + ), + ), + TourStep( + target=_panel_attr( + "Registration", RegistrationPanel, "_reference_time_spin" + ), + title="Within-scan Inputs", + body=( + "For Within-scan, choose one layer with a time dimension, then " + "pick the Reference volume that every frame should align to." + ), + anchor="left", + spotlight_rect=_panel_attr_rect( + "Registration", + RegistrationPanel, + "_time_series_radio", + "_moving_label", + "_moving_combo", + "_reference_time_label", + "_reference_time_spin", + ), + tooltip_target=_dock_widget, + pre_action=_run_actions( + _select_sub_panel("Registration", "_register_panel_radio"), + _set_panel_button("Registration", "_time_series_radio"), + _set_panel_button("Registration", "_advanced_toggle", checked=False), + ), + ), + TourStep( + target=_panel_attr("Registration", RegistrationPanel, "_transform_combo"), + title="Within-scan Parameters", + body=( + "The parameter box is similar in Within-scan mode, but some " + "defaults are adjusted for motion correction." + ), + anchor="left", + spotlight_rect=_panel_attr_ancestor_rect( + "Registration", + RegistrationPanel, + "_transform_combo", + QGroupBox, + ), + tooltip_target=_dock_widget, + pre_action=_run_actions( + _select_sub_panel("Registration", "_register_panel_radio"), + _set_panel_button("Registration", "_time_series_radio"), + _set_panel_button("Registration", "_advanced_toggle", checked=False), + ), + ), + TourStep( + target=_panel_attr("Registration", RegistrationPanel, "_transforms_panel"), + title="Transforms", + body=( + "The Transforms tab lets you inspect available transforms and " + "save, load, or apply them to other layers." + ), + anchor="left", + tooltip_target=_dock_widget, + pre_action=_select_sub_panel("Registration", "_transforms_panel_radio"), + ), TourStep( target=_accordion_panel("Events"), title="Events", @@ -970,8 +1262,9 @@ def _restore_state() -> None: title="You're Ready to Explore!", body=( "You're all set to start exploring. Load a dataset in Data I/O, " - "overlay a behavioral video, explore signals, run a few QC checks, " - "and have fun digging into some fUSI data!" + "overlay a behavioral video, explore signals, align scans in " + "Registration, run a few QC checks, and have fun digging into " + "some fUSI data!" ), anchor="left", spotlight_rect=lambda: _widget_rect(plugin_widget), diff --git a/src/confusius/_napari/_video/_video_panel.py b/src/confusius/_napari/_video/_video_panel.py index 2a06ddbb..4518923a 100644 --- a/src/confusius/_napari/_video/_video_panel.py +++ b/src/confusius/_napari/_video/_video_panel.py @@ -34,14 +34,14 @@ class _VideoArray: """Array-like wrapper around `VideoReaderNP` for napari Image layers. - Provides the ``shape``, ``dtype``, and ``__getitem__`` interface that + Provides the `shape`, `dtype`, and `__getitem__` interface that napari requires for lazy, frame-on-demand display. Handles singleton-dimension padding so the video matches the fUSI scan's dimensionality. - The positions of H and W in the shape are controlled by ``h_dim`` and - ``w_dim``. When ``h_dim > w_dim`` (H appears after W in the layout), - the raw ``(H, W)`` frame is transposed before reshaping so that the + The positions of H and W in the shape are controlled by `h_dim` and + `w_dim`. When `h_dim > w_dim` (H appears after W in the layout), + the raw `(H, W)` frame is transposed before reshaping so that the data matches the expected axis order. Parameters @@ -51,20 +51,20 @@ class _VideoArray: dtype : numpy.dtype Data type of a decoded frame. frame_shape : tuple[int, ...] - Shape of a single decoded frame --- ``(H, W)`` or ``(H, W, C)``. + Shape of a single decoded frame --- `(H, W)` or `(H, W, C)`. n_pad : int, default: 0 Number of size-1 dimensions inserted between the time axis and the spatial axes. step : int, default: 1 Show every *step*-th frame (temporal subsampling). Logical - frame ``t`` maps to physical frame ``t * step``. + frame `t` maps to physical frame `t * step`. time_dim : int, default: 0 Position of the time axis in the output shape. h_dim : int or None, optional - Position of the video height axis. Defaults to ``n_core - 2`` - where ``n_core = 1 + n_pad + 2``. + Position of the video height axis. Defaults to `n_core - 2` + where `n_core = 1 + n_pad + 2`. w_dim : int or None, optional - Position of the video width axis. Defaults to ``n_core - 1``. + Position of the video width axis. Defaults to `n_core - 1`. """ def __init__( @@ -221,7 +221,7 @@ class VideoPanel(QWidget): layer and its own grid cell. All videos share a single reference fUSI scan (selected at first load) so they align on the same time and spatial axes. Videos are passed to napari as lazy, array-like - objects backed by ``VideoReaderNP`` (OpenCV frame-on-demand + objects backed by `VideoReaderNP` (OpenCV frame-on-demand decoding). A thin wrapper (`_VideoArray`) handles singleton dimension padding and dimension reordering. @@ -229,9 +229,9 @@ class VideoPanel(QWidget): is rebuilt with a new `_VideoArray` whose shape places H and W at the currently displayed dim positions. - The video layers receive the same ``axis_labels`` as the fUSI scan + The video layers receive the same `axis_labels` as the fUSI scan so that napari handles dimension reordering identically for both. - The time scale is ``frame_step / fps``, shared across all videos. + The time scale is `frame_step / fps`, shared across all videos. Spatial dimensions use a per-video isotropic scale matching the fUSI scan height. @@ -366,7 +366,7 @@ def _setup_ui(self) -> None: playback_layout.setSpacing(4) step_row = QHBoxLayout() - step_row.addWidget(QLabel("Frame step:")) + step_row.addWidget(QLabel("Frame step")) self._step_spin = QSpinBox() self._step_spin.setRange(1, 100) self._step_spin.setValue(1) @@ -504,7 +504,7 @@ def _browse(self) -> None: self._load_from_path() def _load_from_path(self) -> None: - """Validate inputs and call ``_add_video``.""" + """Validate inputs and call `_add_video`.""" ref = self._get_ref_layer() if ref is None: show_error("Select a reference layer first.") @@ -616,8 +616,8 @@ def _rebuild_entry(self, entry: _VideoEntry) -> None: ) # Time scale = frame_step / fps. Each logical frame spans - # ``frame_step`` physical frames, so consecutive data points are - # ``frame_step / fps`` seconds apart. + # `frame_step` physical frames, so consecutive data points are + # `frame_step / fps` seconds apart. time_scale = frame_step / entry.fps if entry.fps > 0 else 1.0 # Isotropic spatial scale (video pixels are square). spatial_scale = self._compute_spatial_scale(displayed_v, entry.video_h) @@ -738,7 +738,7 @@ def _on_frame_step_changed(self, value: int) -> None: """Rebuild all video layers with a new frame step. The step is encoded in each `_VideoArray` shape (fewer logical - frames) and the layer's time scale (``value / fps``). napari + frames) and the layer's time scale (`value / fps`). napari auto-computes the correct slider range from shape and scale. The current world time is saved before the rebuild and restored @@ -769,7 +769,7 @@ def _on_frame_step_changed(self, value: int) -> None: def _lookup_coord(self, dim_idx: int) -> np.ndarray | None: """Return the reference xarray coordinate for *dim_idx*, or None. - Returns ``None`` when there is no reference layer, no xarray + Returns `None` when there is no reference layer, no xarray metadata, or the corresponding coordinate does not exist. """ ref = self._ref_layer @@ -795,7 +795,7 @@ def _compute_spatial_scale(self, vertical_dim: int, video_h: int) -> float: """Return the isotropic spatial scale for the video. The scale maps the video's height to the fUSI scan's extent - along ``vertical_dim`` and is then applied identically to both + along `vertical_dim` and is then applied identically to both displayed spatial axes so that video pixels remain square -- webcam pixels are isotropic and must not be stretched. @@ -827,8 +827,8 @@ def _compute_axis_center_translate( ) -> float: """Return the translation that centers the video on the fUSI. - The video's centre pixel along ``dim_idx`` (at index - ``(video_n - 1) / 2``) is placed at the midpoint of the fUSI + The video's centre pixel along `dim_idx` (at index + `(video_n - 1) / 2`) is placed at the midpoint of the fUSI coordinate range, so the video overlays the scan in both spatial axes. """ diff --git a/src/confusius/_napari/_widget.py b/src/confusius/_napari/_widget.py index 902083dd..85a65991 100644 --- a/src/confusius/_napari/_widget.py +++ b/src/confusius/_napari/_widget.py @@ -244,7 +244,7 @@ class ConfUSIusWidget(QWidget): def __init__(self, napari_viewer: napari.Viewer) -> None: super().__init__() self.viewer = napari_viewer - self.setMinimumWidth(350) + self.setMinimumWidth(430) self.setSizePolicy( QSizePolicy.Policy.MinimumExpanding, QSizePolicy.Policy.Expanding, @@ -377,8 +377,8 @@ def _make_header(self) -> QWidget: header.setObjectName("confusius_header") layout = QVBoxLayout(header) - layout.setContentsMargins(4, 6, 12, 14) - layout.setSpacing(2) + layout.setContentsMargins(4, 12, 12, 14) + layout.setSpacing(0) tour_btn = QPushButton("Take a Tour") tour_btn.setObjectName("tour_btn") @@ -388,11 +388,11 @@ def _make_header(self) -> QWidget: tour_btn.adjustSize() logo_widget = self._load_logo() - logo_row = QHBoxLayout() - logo_row.setContentsMargins(0, 0, 0, 6) - logo_row.setSpacing(10) + header_row = QHBoxLayout() + header_row.setContentsMargins(0, 0, 0, 6) + header_row.setSpacing(10) if logo_widget is not None: - logo_row.addWidget(logo_widget) + header_row.addWidget(logo_widget, alignment=Qt.AlignmentFlag.AlignTop) title = QLabel("ConfUSIus") title.setObjectName("confusius_title") @@ -402,24 +402,18 @@ def _make_header(self) -> QWidget: subtitle.setObjectName("confusius_subtitle") subtitle.setIndent(0) - tour_btn_title_and_subtitle = QWidget() - tour_btn_title_and_subtitle_layout = QVBoxLayout(tour_btn_title_and_subtitle) - tour_btn_title_and_subtitle_layout.setContentsMargins(0, 0, 0, 0) - tour_btn_title_and_subtitle_layout.setSpacing(0) - - tour_btn_title_and_subtitle_layout.addStretch() - tour_btn_title_and_subtitle_layout.addWidget( - tour_btn, alignment=Qt.AlignmentFlag.AlignTop | Qt.AlignmentFlag.AlignRight - ) - tour_btn_title_and_subtitle_layout.addWidget( - title, alignment=Qt.AlignmentFlag.AlignLeft - ) - tour_btn_title_and_subtitle_layout.addWidget( - subtitle, alignment=Qt.AlignmentFlag.AlignLeft - ) + title_block = QWidget() + title_block_layout = QVBoxLayout(title_block) + title_block_layout.setContentsMargins(0, 4, 0, 0) + title_block_layout.setSpacing(0) + title_block_layout.addWidget(title, alignment=Qt.AlignmentFlag.AlignLeft) + title_block_layout.addWidget(subtitle, alignment=Qt.AlignmentFlag.AlignLeft) + title_block_layout.addStretch() - logo_row.addWidget(tour_btn_title_and_subtitle) - layout.addLayout(logo_row) + header_row.addWidget(title_block, alignment=Qt.AlignmentFlag.AlignTop) + header_row.addStretch() + header_row.addWidget(tour_btn, alignment=Qt.AlignmentFlag.AlignTop) + layout.addLayout(header_row) return header @@ -496,6 +490,7 @@ def _make_accordion(self) -> QWidget: from confusius._napari._data._save_panel import SavePanel from confusius._napari._events._panel import EventPanel from confusius._napari._qc._panel import QCPanel + from confusius._napari._registration._panel import RegistrationPanel from confusius._napari._signals._panel import SignalPanel from confusius._napari._video._video_panel import VideoPanel @@ -521,6 +516,7 @@ def _make_accordion(self) -> QWidget: ("Data I/O", "file-input"), ("Video", "video"), ("Signals", "chart-line"), + ("Registration", "images"), ("Events", "calendar-clock"), ("Quality Control", "clipboard-check"), ] @@ -528,6 +524,7 @@ def _make_accordion(self) -> QWidget: data_panel, video_panel, SignalPanel(self.viewer, event_store=self._event_store), + RegistrationPanel(self.viewer), EventPanel(self.viewer, self._event_store), QCPanel(self.viewer), ] diff --git a/src/confusius/_napari/assets/images.svg b/src/confusius/_napari/assets/images.svg new file mode 100644 index 00000000..c6af1a80 --- /dev/null +++ b/src/confusius/_napari/assets/images.svg @@ -0,0 +1 @@ + diff --git a/src/confusius/datasets/_osf.py b/src/confusius/datasets/_osf.py index 8b3823dd..d0543070 100644 --- a/src/confusius/datasets/_osf.py +++ b/src/confusius/datasets/_osf.py @@ -236,7 +236,7 @@ def _validate_index(index: object, data_dir: Path) -> None: f"The dataset index in {data_dir} has an outdated or unrecognised structure " f"(every entry must define {sorted(_INDEX_ENTRY_KEYS)}). This local dataset was " f"likely downloaded with an older version of confusius. Delete the dataset " - f"directory and fetch it again to update it:\n\n {data_dir}\n" + f"directory and fetch it again to update it." ) diff --git a/src/confusius/registration/__init__.py b/src/confusius/registration/__init__.py index af02e9ee..52372159 100644 --- a/src/confusius/registration/__init__.py +++ b/src/confusius/registration/__init__.py @@ -1,11 +1,15 @@ """Registration module for fUSI data.""" -from confusius.registration._progress import RegistrationProgressPlotter +from confusius.registration.progress import ( + RegistrationProgress, + MatplotlibRegistrationProgressPlotter, +) from confusius.registration.affines import ( compose_affine, decompose_affine, ) from confusius.registration.diagnostics import RegistrationDiagnostics +from confusius.registration.exceptions import RegistrationAbortedError from confusius.registration.motion import ( compute_framewise_displacement, create_motion_dataframe, @@ -17,16 +21,20 @@ ) from confusius.registration.volume import register_volume from confusius.registration.volumewise import register_volumewise +from confusius.registration.volumewise_progress import VolumewiseProgressReporter __all__ = [ + "RegistrationAbortedError", "RegistrationDiagnostics", - "RegistrationProgressPlotter", + "RegistrationProgress", + "MatplotlibRegistrationProgressPlotter", "compose_affine", "decompose_affine", "register_volume", "resample_volume", "resample_like", "register_volumewise", + "VolumewiseProgressReporter", "extract_motion_parameters", "compute_framewise_displacement", "create_motion_dataframe", diff --git a/src/confusius/registration/_utils.py b/src/confusius/registration/_utils.py index 4339b34e..8b1aff4d 100644 --- a/src/confusius/registration/_utils.py +++ b/src/confusius/registration/_utils.py @@ -1,6 +1,8 @@ """Internal utilities shared by registration modules.""" import os +import signal +import threading from contextlib import contextmanager from copy import deepcopy from typing import TYPE_CHECKING, Generator @@ -9,6 +11,8 @@ import xarray as xr if TYPE_CHECKING: + from threading import Event + import SimpleITK as sitk @@ -70,6 +74,63 @@ def set_sitk_thread_count(n: int) -> Generator[None, None, None]: sitk.ProcessObject.SetGlobalDefaultNumberOfThreads(prev) +@contextmanager +def abort_on_sigint( + abort_event: "Event | None", +) -> Generator["Event", None, None]: + """Return an abort event that is set cooperatively on the first Ctrl+C. + + Parameters + ---------- + abort_event : threading.Event or None + Existing cooperative-cancellation event to reuse. If not provided, a + new event is created for the duration of the context. + + Yields + ------ + threading.Event + Event that is set when cooperative cancellation is requested, either + explicitly by the caller or via a Ctrl+C signal handled on the main + thread. + + Notes + ----- + On the main thread, the first `SIGINT`/Ctrl+C is converted into + `abort_event.set()` so long-running registrations can stop cleanly at the + next SimpleITK iteration boundary and return their current partial result. + A second Ctrl+C falls back to the previous signal handler so users can + still force an immediate interrupt if graceful cancellation stalls. + """ + shared_abort_event = abort_event or threading.Event() + + if threading.current_thread() is not threading.main_thread(): + yield shared_abort_event + return + + previous_handler = signal.getsignal(signal.SIGINT) + saw_sigint = False + + def _handle_sigint(signum: int, frame: object) -> None: + nonlocal saw_sigint + if not saw_sigint: + saw_sigint = True + shared_abort_event.set() + return + + if previous_handler in {signal.SIG_DFL, signal.default_int_handler}: + raise KeyboardInterrupt + if previous_handler == signal.SIG_IGN: + return + if callable(previous_handler): + previous_handler(signum, frame) + + signal.signal(signal.SIGINT, _handle_sigint) + try: + yield shared_abort_event + finally: + signal.signal(signal.SIGINT, previous_handler) + + def dataarray_to_sitk_image(da: xr.DataArray) -> "sitk.Image": """Convert a spatial or spatiotemporal DataArray to a SimpleITK image. diff --git a/src/confusius/registration/bspline.py b/src/confusius/registration/bspline.py index c2c8b65e..dded3e89 100644 --- a/src/confusius/registration/bspline.py +++ b/src/confusius/registration/bspline.py @@ -1,6 +1,6 @@ """B-spline transform helpers for fUSI registration. -A B-spline deformation field is represented as a DataArray with: +A B-spline control-point grid is represented as a DataArray with: - **dims**: `("component", )` — e.g. `("component", "z", "y", "x")`. - **coords**: physical mm positions of the control-point grid along each spatial axis. @@ -8,13 +8,13 @@ ```python { - "type": "bspline_transform", - "order": 3, # B-spline polynomial order - "direction": [[...], [...], [...]], # (ndim, ndim) direction cosine matrix - "affines": { - "bspline_initialization": [[...]] # optional (N+1, N+1) pre-affine; - # only present when register_volume - # was called with affine initialization. + "transform_type": "bspline_transform", + "order": 3, # B-spline polynomial order + "direction": [[...], [...], [...]], # (ndim, ndim) direction cosine matrix + "affines": { + "bspline_initialization": [[...]] # optional (N+1, N+1) pre-affine; + # only present when register_volume + # was called with affine initialization. } } ``` @@ -22,11 +22,15 @@ When a pre-affine is stored in `attrs["affines"]["bspline_initialization"]`, the full transform is a `CompositeTransform(pre_affine, bspline)` — i.e. the pre-affine is applied *first* (coarse global alignment) and the B-spline is applied *second* (local -deformation refinement). This mirrors the `inPlace=True` composite that SimpleITK +deformation refinement). This mirrors the `inPlace=True` composite that SimpleITK optimises during registration. + +This object is not a dense deformation field sampled on every voxel of the moving or +fixed image. Instead, it stores the sparse B-spline coefficient lattice that defines +the smooth deformation. """ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast import numpy as np import numpy.typing as npt @@ -58,7 +62,8 @@ def sitk_bspline_to_dataarray( Returns ------- xarray.DataArray - B-spline control-point DataArray with `attrs["type"] == "bspline_transform"`. + B-spline control-point DataArray with + `attrs["transform_type"] == "bspline_transform"`. Raises ------ @@ -102,7 +107,7 @@ def sitk_bspline_to_dataarray( coords[dim] = origin[i] + np.arange(grid_shape[i]) * spacing[i] attrs: dict[str, object] = { - "type": "bspline_transform", + "transform_type": "bspline_transform", "order": order, "direction": direction.tolist(), } @@ -142,7 +147,7 @@ def _dataarray_to_sitk_bspline(da: xr.DataArray) -> "sitk.Transform": """ import SimpleITK as sitk - _validate_bspline_dataarray(da) + validate_bspline_dataarray(da) ndim = da.ndim - 1 # subtract the component axis order = int(da.attrs["order"]) @@ -222,20 +227,23 @@ def _extract_bspline(transform: "sitk.Transform") -> "sitk.BSplineTransform": name = transform.GetName() if "BSpline" in name: - return transform # type: ignore[return-value] + return cast("sitk.BSplineTransform", transform) if name == "CompositeTransform": - n = transform.GetNumberOfTransforms() # type: ignore[attr-defined] - # The B-spline is the last sub-transform (it was added last and is optimised). - last = transform.GetNthTransform(n - 1) # type: ignore[attr-defined] - if "BSpline" in last.GetName(): - return last + get_count = getattr(transform, "GetNumberOfTransforms", None) + get_nth = getattr(transform, "GetNthTransform", None) + if callable(get_count) and callable(get_nth): + n = int(get_count()) + # The B-spline is the last sub-transform (it was added last and is optimised). + last = get_nth(n - 1) + if "BSpline" in last.GetName(): + return cast("sitk.BSplineTransform", last) raise TypeError( f"Expected a BSplineTransform or a CompositeTransform ending with a " f"BSplineTransform; got {transform.GetName()!r}." ) -def _validate_bspline_dataarray(da: xr.DataArray) -> None: +def validate_bspline_dataarray(da: xr.DataArray) -> None: """Raise ValueError if *da* does not look like a valid B-spline transform DataArray. Parameters @@ -246,12 +254,15 @@ def _validate_bspline_dataarray(da: xr.DataArray) -> None: Raises ------ ValueError - If `da.attrs["type"] != "bspline_transform"` or required attrs are missing. + If `da.attrs["transform_type"] != "bspline_transform"` or required attrs are + missing. """ - if da.attrs.get("type") != "bspline_transform": + transform_type = da.attrs.get("transform_type") + if transform_type != "bspline_transform": raise ValueError( - f"Expected a DataArray with attrs['type'] == 'bspline_transform'; " - f"got {da.attrs.get('type')!r}." + "Expected a DataArray with attrs['transform_type'] == " + "'bspline_transform'; " + f"got {transform_type!r}." ) for key in ("order", "direction"): if key not in da.attrs: diff --git a/src/confusius/registration/diagnostics.py b/src/confusius/registration/diagnostics.py index b251f98f..81538f6a 100644 --- a/src/confusius/registration/diagnostics.py +++ b/src/confusius/registration/diagnostics.py @@ -36,8 +36,13 @@ class RegistrationDiagnostics: smaller than `register_volume`'s `number_of_iterations` if the optimizer converged early. stop_condition : str - Human-readable description of the optimizer stop condition, as - returned by SimpleITK's `GetOptimizerStopConditionDescription`. + Human-readable description of the optimizer stop condition. When the + registration completes normally, this is SimpleITK's + `GetOptimizerStopConditionDescription`. When the registration is + aborted cooperatively, this is a short abort message. + status : {"completed", "aborted"} + Whether the registration ran to completion or returned an intermediate + result after cooperative cancellation. """ metric: Literal["correlation", "mattes_mi"] @@ -45,3 +50,4 @@ class RegistrationDiagnostics: final_metric_value: float n_iterations: int stop_condition: str + status: Literal["completed", "aborted"] diff --git a/src/confusius/registration/exceptions.py b/src/confusius/registration/exceptions.py new file mode 100644 index 00000000..2bb15203 --- /dev/null +++ b/src/confusius/registration/exceptions.py @@ -0,0 +1,5 @@ +"""Exceptions raised by registration workflows.""" + + +class RegistrationAbortedError(RuntimeError): + """Raised when a registration run is cancelled before completion.""" diff --git a/src/confusius/registration/_progress.py b/src/confusius/registration/progress.py similarity index 63% rename from src/confusius/registration/_progress.py rename to src/confusius/registration/progress.py index 868146f7..b1d11c5f 100644 --- a/src/confusius/registration/_progress.py +++ b/src/confusius/registration/progress.py @@ -1,7 +1,9 @@ """Registration progress visualization.""" +from __future__ import annotations + import warnings -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal, Protocol, cast import numpy as np @@ -12,6 +14,7 @@ import SimpleITK as sitk from matplotlib.figure import Figure + _INTERPOLATION_MAP = { "linear": "sitkLinear", "nearest": "sitkNearestNeighbor", @@ -19,7 +22,109 @@ } -class RegistrationProgressPlotter: +def _resolve_sitk_interpolation(interpolation: str | None) -> Any: + """Return the SimpleITK interpolator enum for a named interpolation. + + Parameters + ---------- + interpolation : str + One of `"linear"`, `"nearest"`, `"bspline"`. + + Returns + ------- + SimpleITK interpolator enum + The matching `sitk.sitk*` interpolator constant. + + Raises + ------ + ValueError + If `interpolation` is not one of the supported names. + """ + import SimpleITK as sitk + + if interpolation is None: + interpolation = "linear" + interp_name = _INTERPOLATION_MAP.get(interpolation) + if interp_name is None: + supported = ", ".join(sorted(_INTERPOLATION_MAP)) + msg = ( + f"Invalid `interpolation`: {interpolation!r}. Expected one of: {supported}." + ) + raise ValueError(msg) + return getattr(sitk, interp_name) + + +def _resample_intermediate( + registration_method: "sitk.ImageRegistrationMethod", + moving_img: "sitk.Image", + fixed_img: "sitk.Image", + *, + interpolation: Literal["linear", "nearest", "bspline"] = "linear", + fill_value: float = 0.0, + sitk_threads: int = -1, +) -> "sitk.Image": + """Resample the moving image onto the fixed grid using the current transform. + + Shared by the matplotlib and napari progress plotters so the per-iteration + resample logic stays in one place. + + Parameters + ---------- + registration_method : SimpleITK.ImageRegistrationMethod + The active registration method whose initial transform is used to + resample. + moving_img : SimpleITK.Image + Moving image to resample. + fixed_img : SimpleITK.Image + Reference image defining the output grid. + interpolation : {"linear", "nearest", "bspline"}, default: "linear" + Interpolator used for the intermediate resample. + fill_value : float, default: 0.0 + Fill value used outside the moving image field of view. + sitk_threads : int, default: -1 + Number of threads SimpleITK may use for the intermediate resample. + + Returns + ------- + SimpleITK.Image + Resampled image on the fixed grid. + """ + import SimpleITK as sitk + + from confusius.registration._utils import set_sitk_thread_count + + sitk_interp = _resolve_sitk_interpolation(interpolation) + + transform = registration_method.GetInitialTransform() + with set_sitk_thread_count(sitk_threads): + return sitk.Resample( + moving_img, + fixed_img, + transform, + sitk_interp, + fill_value, + moving_img.GetPixelID(), + ) + + +class RegistrationProgress(Protocol): + """Duck-typed contract for an iteration progress reporter. + + Implementations are called from the registration thread (SimpleITK's + iteration/end callbacks). They must be safe to call from a non-GUI thread; + any GUI side effects must be marshalled via Qt signals or similar. + """ + + def update(self) -> None: + """Called at every optimizer iteration event.""" + ... + + def close(self) -> None: + """Called once at the registration end event.""" + ... + + +class MatplotlibRegistrationProgressPlotter: """Plot registration progress in real time. Displays an optimizer metric curve, a composite fixed/moving overlay, or @@ -40,8 +145,8 @@ class RegistrationProgressPlotter: Whether to display a blended fixed/moving composite at each iteration. Requires an additional `sitk.Resample` call per iteration. resample_kwargs : dict, optional - Extra keyword arguments forwarded to the internal resample call at each - iteration. + Extra keyword arguments for the internal resample call at each iteration. + Supported keys are `interpolation`, `fill_value`, and `sitk_threads`. """ def __init__( @@ -65,11 +170,17 @@ def __init__( self._metric_values: list[float] = [] _kw: dict[str, Any] = dict(resample_kwargs or {}) - if "default_value" not in _kw: + self._interpolation = cast( + 'Literal["linear", "nearest", "bspline"]', + _kw.get("interpolation", "linear"), + ) + if "fill_value" in _kw: + self._fill_value = float(_kw["fill_value"]) + else: import SimpleITK as sitk - _kw["default_value"] = float(sitk.GetArrayFromImage(moving_img).min()) - self._resample_kwargs = _kw + self._fill_value = float(sitk.GetArrayFromImage(moving_img).min()) + self._sitk_threads = int(_kw.get("sitk_threads", -1)) # Detect Jupyter notebook environment. A plain IPython terminal shell # also has get_ipython() != None, so we check the kernel class name to @@ -153,36 +264,19 @@ def update(self) -> None: self._metric_ax.autoscale_view() if self._plot_composite: - import SimpleITK as sitk - - from confusius.registration._utils import set_sitk_thread_count + resampled = _resample_intermediate( + self._method, + self._moving_img, + self._fixed_img, + interpolation=self._interpolation, + fill_value=self._fill_value, + sitk_threads=self._sitk_threads, + ) - interpolation = self._resample_kwargs.get("interpolation", "linear") - interp_name = _INTERPOLATION_MAP.get(interpolation) - if interp_name is None: - supported = ", ".join(sorted(_INTERPOLATION_MAP)) - msg = ( - "Invalid `interpolation` in `resample_kwargs`: " - f"{interpolation!r}. Expected one of: {supported}." - ) - raise ValueError(msg) - sitk_interp = getattr(sitk, interp_name) - fill_value = self._resample_kwargs["default_value"] - sitk_threads = self._resample_kwargs.get("sitk_threads", -1) - - transform = self._method.GetInitialTransform() - with set_sitk_thread_count(sitk_threads): - resampled = sitk.Resample( - self._moving_img, - self._fixed_img, - transform, - sitk_interp, - fill_value, - self._moving_img.GetPixelID(), - ) + import SimpleITK as sitk - fixed_arr = sitk.GetArrayFromImage(self._fixed_img).T - moving_arr = sitk.GetArrayFromImage(resampled).T + fixed_arr = np.asarray(sitk.GetArrayFromImage(self._fixed_img).T) + moving_arr = np.asarray(sitk.GetArrayFromImage(resampled).T) if fixed_arr.ndim == 3: rgb = make_mosaic( diff --git a/src/confusius/registration/volume.py b/src/confusius/registration/volume.py index f9802e36..e31d4c49 100644 --- a/src/confusius/registration/volume.py +++ b/src/confusius/registration/volume.py @@ -1,6 +1,6 @@ """Volume-to-volume registration for fUSI data.""" -from collections.abc import Sequence +from collections.abc import Callable, Sequence from typing import TYPE_CHECKING, Literal, overload import numpy as np @@ -8,6 +8,7 @@ import xarray as xr from confusius.registration._utils import ( + abort_on_sigint, dataarray_to_sitk_image, replace_affines_attr, set_sitk_thread_count, @@ -19,8 +20,12 @@ from confusius.registration.diagnostics import RegistrationDiagnostics if TYPE_CHECKING: + from threading import Event + import SimpleITK as sitk + from confusius.registration.progress import RegistrationProgress + def _validate_register_volume_inputs( moving: xr.DataArray, @@ -240,6 +245,56 @@ def _expand_thin_dims(img: "sitk.Image", min_size: int = 4) -> "sitk.Image": return sitk.Expand(img, factors.tolist()) +def _translate_registration_runtime_error( + exc: RuntimeError, + *, + transform_type: Literal["translation", "rigid", "affine", "bspline"], + learning_rate: float | Literal["auto"], +) -> RuntimeError: + """Return a clearer registration error for known SimpleITK failures. + + Parameters + ---------- + exc : RuntimeError + Exception raised by SimpleITK during optimizer execution. + transform_type : {"translation", "rigid", "affine", "bspline"} + Registration model used for the failed run. + learning_rate : float or "auto" + User-requested learning rate mode. + + Returns + ------- + RuntimeError + Translated exception when the failure mode is recognized, otherwise `exc`. + """ + message = str(exc) + if "m_Scales values must be > epsilon" not in message: + return exc + + parts = [ + "SimpleITK could not compute valid optimizer scales for this registration.", + "Some transform parameters have near-zero physical effect, so the gradient-descent optimizer cannot choose a stable step size.", + ] + if transform_type == "bspline": + parts.append( + "This is most common for `transform_type='bspline'`, especially when the control-point grid is too fine for the image extent or overlap." + ) + if learning_rate == "auto": + parts.append( + 'Retry with a fixed `learning_rate` such as `0.1` or `0.01` instead of `"auto"`.' + ) + else: + parts.append( + "Changing `learning_rate` alone may not help because this failure happens before optimisation starts." + ) + if transform_type == "bspline": + parts.append( + "If that still fails, use a coarser `mesh_size` or run affine/rigid registration first and pass the result as `initialization`." + ) + + return RuntimeError(" ".join(parts)) + + @overload def register_volume( # numpydoc ignore=GL08,PR01,RT01 moving: xr.DataArray, @@ -264,11 +319,13 @@ def register_volume( # numpydoc ignore=GL08,PR01,RT01 smoothing_sigmas: Sequence[int] = ..., resample: bool = ..., resample_interpolation: Literal["linear", "bspline"] = ..., + fill_value: float | None = ..., sitk_threads: int = ..., show_progress: bool = ..., plot_metric: bool = ..., plot_composite: bool = ..., - fill_value: float | None = ..., + progress_plotter: "Callable[..., RegistrationProgress] | None" = None, + abort_event: "Event | None" = ..., ) -> "tuple[xr.DataArray, npt.NDArray[np.floating], RegistrationDiagnostics]": """Overload for linear transforms (translation/rigid/affine).""" ... @@ -298,11 +355,13 @@ def register_volume( # numpydoc ignore=GL08,PR01,RT01 smoothing_sigmas: Sequence[int] = ..., resample: bool = ..., resample_interpolation: Literal["linear", "bspline"] = ..., + fill_value: float | None = ..., sitk_threads: int = ..., show_progress: bool = ..., plot_metric: bool = ..., plot_composite: bool = ..., - fill_value: float | None = ..., + progress_plotter: "Callable[..., RegistrationProgress] | None" = None, + abort_event: "Event | None" = ..., ) -> "tuple[xr.DataArray, xr.DataArray, RegistrationDiagnostics]": """Overload for bspline transform (returns DataArray transform).""" ... @@ -331,11 +390,13 @@ def register_volume( # numpydoc ignore=GL08,PR01,RT01 smoothing_sigmas: Sequence[int] = ..., resample: bool = ..., resample_interpolation: Literal["linear", "bspline"] = ..., + fill_value: float | None = ..., sitk_threads: int = ..., show_progress: bool = ..., plot_metric: bool = ..., plot_composite: bool = ..., - fill_value: float | None = ..., + progress_plotter: "Callable[..., RegistrationProgress] | None" = None, + abort_event: "Event | None" = ..., ) -> "tuple[xr.DataArray, npt.NDArray[np.floating], RegistrationDiagnostics]": """Overload for default transform (rigid, returns affine).""" ... @@ -364,11 +425,13 @@ def register_volume( smoothing_sigmas: Sequence[int] = (6, 2, 1), resample: bool = True, resample_interpolation: Literal["linear", "bspline"] = "linear", + fill_value: float | None = None, sitk_threads: int = -1, show_progress: bool = False, plot_metric: bool = True, plot_composite: bool = True, - fill_value: float | None = None, + progress_plotter: "Callable[..., RegistrationProgress] | None" = None, + abort_event: "Event | None" = None, ) -> "tuple[xr.DataArray, npt.NDArray[np.floating] | xr.DataArray, RegistrationDiagnostics]": # noqa: E501 """Register a single 2D or 3D volume to a fixed reference. @@ -467,6 +530,13 @@ def register_volume( `"linear"` is fast and appropriate for most cases. `"bspline"` (3rd-order B-spline) produces smoother results and reduces ringing, useful for atlas registration. Only used when `resample=True`. + fill_value : float, optional + Fill value for voxels that fall outside the moving image's field of view after + resampling. Applied to both the final registered output (when `resample=True`) + and the progress composite overlay (when `show_progress=True` and + `plot_composite=True`). If not provided, defaults to the minimum value of + `moving`, which renders out-of-FOV regions as background regardless of intensity + scale (important for dB data where 0 is maximum intensity). sitk_threads : int, default: -1 Number of threads SimpleITK may use internally. Negative values resolve to `max(1, os.cpu_count() + 1 + sitk_threads)`, so `-1` means all CPUs, `-2` @@ -484,13 +554,22 @@ def register_volume( Whether to include a fixed/moving composite overlay in the progress plot. Requires resampling the moving image at every iteration. Ignored when `show_progress=False`. - fill_value : float, optional - Fill value for voxels that fall outside the moving image's field of view after - resampling. Applied to both the final registered output (when `resample=True`) - and the progress composite overlay (when `show_progress=True` and - `plot_composite=True`). If not provided, defaults to the minimum value of - `moving`, which renders out-of-FOV regions as background regardless of intensity - scale (important for dB data where 0 is maximum intensity). + progress_plotter : callable, optional + Factory that builds the progress reporter, called inside `register_volume` as + `progress_plotter(registration_method, fixed_img, moving_img, *, plot_metric, + plot_composite, resample_kwargs)`. Here `resample_kwargs` carries + `interpolation`, `fill_value`, and `sitk_threads`. The returned object must + implement the + [`RegistrationProgress`][confusius.registration.RegistrationProgress] protocol + (`update()` / `close()`). If not provided, the default + [`MatplotlibRegistrationProgressPlotter`][confusius.registration.MatplotlibRegistrationProgressPlotter] + is used. Ignored when `show_progress=False`. Custom factories are expected to + be safe to call from a non-GUI thread; GUI side effects must be marshalled via + thread-safe primitives such as Qt signals. + abort_event : threading.Event, optional + Cooperative cancellation flag. If set before or during optimisation, the + registration stops at the next SimpleITK iteration boundary and returns + the current intermediate result with `diagnostics.status="aborted"`. Returns ------- @@ -505,11 +584,16 @@ def register_volume( N+1)` in physical space, where `N` is the spatial dimensionality (2 or 3). Follows SimpleITK's pull/inverse convention: the matrix maps fixed-space coordinates to moving-space coordinates. For `transform_type="bspline"`, - returns a DataArray encoding the B-spline control-point grid (see - [`confusius.registration.bspline`][confusius.registration.bspline] for the - DataArray schema). When an affine `initialization` was also supplied, the DataArray - includes `attrs["affines"]["bspline_initialization"]` so that the full composite - (pre-affine + B-spline) can be reconstructed for resampling. + returns an `xarray.DataArray` containing the B-spline control-point grid, not a + dense deformation field. The first dimension is `component` with length `N`, + followed by spatial dimensions in ConfUSIus order (`("y", "x")` in 2D or + `("z", "y", "x")` in 3D). The coordinate values along each spatial axis are + the physical positions of the control points. Attributes include `type = + "bspline_transform"`, the spline `order`, and the control-grid `direction` + matrix. When an affine `initialization` was also supplied, the DataArray also + includes `attrs["affines"]["bspline_initialization"]` so that the full + composite transform (pre-affine + B-spline) can be reconstructed for later + resampling. diagnostics : confusius.registration.RegistrationDiagnostics Per-iteration metric values, final metric value, iteration count, and the optimizer stop condition. Useful for plotting convergence curves, comparing @@ -589,6 +673,8 @@ def register_volume( f"image dimensionality {ndim}D (expected {expected_shape})." ) + requested_learning_rate = learning_rate + registration = sitk.ImageRegistrationMethod() # --- Metric --- @@ -687,10 +773,12 @@ def register_volume( if affine_initialization is not None: pre_tx = affine_to_sitk_linear_transform(affine_initialization) - - sitk_initial_transform: sitk.Transform = sitk.CompositeTransform(ndim) - sitk_initial_transform.AddTransform(pre_tx) - sitk_initial_transform.AddTransform(sitk_centering_transform) + if transform_type == "bspline": + sitk_initial_transform = sitk.CompositeTransform(ndim) + sitk_initial_transform.AddTransform(pre_tx) + sitk_initial_transform.AddTransform(sitk_centering_transform) + else: + sitk_initial_transform = pre_tx else: sitk_initial_transform = sitk_centering_transform @@ -699,10 +787,10 @@ def register_volume( # Always collect per-iteration metric values so callers get convergence # diagnostics regardless of whether the live progress plot is enabled. metric_values: list[float] = [] - registration.AddCommand( - sitk.sitkIterationEvent, - lambda: metric_values.append(registration.GetMetricValue()), - ) + + def _record_iteration() -> None: + metric_value = float(registration.GetMetricValue()) + metric_values.append(metric_value) needs_fill_value = resample or (show_progress and plot_composite) _fill_value = ( @@ -711,29 +799,69 @@ def register_volume( else (float(moving.min()) if needs_fill_value else None) ) - if show_progress: - from confusius.registration._progress import RegistrationProgressPlotter - - resample_kwargs: dict[str, object] = { - "interpolation": resample_interpolation, - "sitk_threads": sitk_threads, - } - if _fill_value is not None: - resample_kwargs["default_value"] = _fill_value - - plotter = RegistrationProgressPlotter( - registration, - fixed_sitk, - moving_sitk, - plot_metric=plot_metric, - plot_composite=plot_composite, - resample_kwargs=resample_kwargs, + with abort_on_sigint(abort_event) as effective_abort_event: + registration.AddCommand(sitk.sitkIterationEvent, _record_iteration) + registration.AddCommand( + sitk.sitkIterationEvent, + lambda: ( + registration.StopRegistration() + if effective_abort_event.is_set() + else None + ), ) - registration.AddCommand(sitk.sitkIterationEvent, plotter.update) - registration.AddCommand(sitk.sitkEndEvent, plotter.close) - with set_sitk_thread_count(sitk_threads): - sitk_optimized_transform = registration.Execute(fixed_reg, moving_reg) + if show_progress: + from confusius.registration.progress import ( + MatplotlibRegistrationProgressPlotter, + ) + + resample_kwargs: dict[str, object] = { + "interpolation": resample_interpolation, + "fill_value": _fill_value, + "sitk_threads": sitk_threads, + } + + plotter_factory = progress_plotter or MatplotlibRegistrationProgressPlotter + plotter = plotter_factory( + registration, + fixed_sitk, + moving_sitk, + plot_metric=plot_metric, + plot_composite=plot_composite, + resample_kwargs=resample_kwargs, + ) + registration.AddCommand(sitk.sitkIterationEvent, plotter.update) + registration.AddCommand(sitk.sitkEndEvent, plotter.close) + + executed = False + if effective_abort_event.is_set(): + if transform_type == "bspline": + sitk_optimized_transform = sitk_initial_transform + elif affine_initialization is not None: + sitk_optimized_transform = affine_to_sitk_linear_transform( + affine_initialization + ) + else: + sitk_optimized_transform = sitk.TranslationTransform(ndim) + aborted = True + stop_condition = "Registration aborted before optimisation started." + else: + try: + with set_sitk_thread_count(sitk_threads): + sitk_optimized_transform = registration.Execute( + fixed_reg, moving_reg + ) + except RuntimeError as exc: + raise _translate_registration_runtime_error( + exc, + transform_type=transform_type, + learning_rate=requested_learning_rate, + ) from exc + executed = True + aborted = effective_abort_event.is_set() + stop_condition = registration.GetOptimizerStopConditionDescription() + if aborted and not stop_condition.strip(): + stop_condition = "Registration aborted." # When resampling, the output lives on the fixed grid; otherwise the moving volume # is returned unchanged and its own coordinates are preserved. @@ -778,17 +906,20 @@ def register_volume( else: optimized_transform = sitk_linear_transform_to_affine(sitk_optimized_transform) - final_metric_value = ( - float(metric_values[-1]) - if metric_values - else float(registration.GetMetricValue()) - ) + final_metric_value: float + if metric_values: + final_metric_value = float(metric_values[-1]) + elif executed: + final_metric_value = float(registration.GetMetricValue()) + else: + final_metric_value = float("nan") diagnostics = RegistrationDiagnostics( metric=metric, metric_values=np.asarray(metric_values, dtype=float), final_metric_value=final_metric_value, n_iterations=len(metric_values), - stop_condition=registration.GetOptimizerStopConditionDescription(), + stop_condition=stop_condition, + status="aborted" if aborted else "completed", ) return result, optimized_transform, diagnostics diff --git a/src/confusius/registration/volumewise.py b/src/confusius/registration/volumewise.py index 677108aa..d3e0b2cf 100644 --- a/src/confusius/registration/volumewise.py +++ b/src/confusius/registration/volumewise.py @@ -2,7 +2,7 @@ from collections.abc import Sequence from contextlib import nullcontext -from typing import Literal, cast +from typing import TYPE_CHECKING, Literal, cast import numpy as np import numpy.typing as npt @@ -14,6 +14,11 @@ from confusius.registration.volume import register_volume from confusius.validation import validate_fusi_dataarray +if TYPE_CHECKING: + from threading import Event + + from confusius.registration.volumewise_progress import VolumewiseProgressReporter + def register_volumewise( data: xr.DataArray, @@ -23,7 +28,7 @@ def register_volumewise( transform: Literal["translation", "rigid", "affine"] = "rigid", metric: Literal["correlation", "mattes_mi"] = "correlation", number_of_histogram_bins: int = 50, - learning_rate: float | Literal["auto"] = "auto", + learning_rate: float | Literal["auto"] = 0.01, number_of_iterations: int = 100, convergence_minimum_value: float = 1e-6, convergence_window_size: int = 10, @@ -35,6 +40,8 @@ def register_volumewise( smoothing_sigmas: Sequence[int] = (6, 2, 1), resample_interpolation: Literal["linear", "bspline"] = "linear", show_progress: bool = True, + progress_reporter: "VolumewiseProgressReporter | None" = None, + abort_event: "Event | None" = None, keep_diagnostics: bool = False, ) -> xr.DataArray: """Register all volumes in a fUSI recording to a reference volume. @@ -61,7 +68,7 @@ def register_volumewise( number_of_histogram_bins : int, default: 50 Number of histogram bins used by Mattes mutual information. Only relevant when `metric="mattes_mi"`. - learning_rate : float or "auto", default: "auto" + learning_rate : float or "auto", default: 0.01 Optimizer step size in normalised units (after `SetOptimizerScalesFromPhysicalShift`). `"auto"` re-estimates the rate at every iteration. A float uses that value directly; if registration diverges or fails to converge, reduce @@ -112,6 +119,15 @@ def register_volumewise( cost of speed. show_progress : bool, default: True Whether to display a progress bar while registering volumes. + progress_reporter : VolumewiseProgressReporter, optional + Thread-safe reporter notified whenever one frame completes. Useful for GUI + progress bars or progressively filling an output layer while frames finish. + abort_event : threading.Event, optional + Cooperative cancellation flag shared across frames. If set before or during + execution, in-flight frame registrations stop at the next optimiser iteration + boundary and this function returns the partial dataset collected so far. Frames + that were not started are left blank (filled with the data minimum), and + per-frame `motion_params` rows are marked via the diagnostics status. keep_diagnostics : bool, default: False Whether to keep the full per-frame [`RegistrationDiagnostics`][confusius.registration.RegistrationDiagnostics] @@ -128,7 +144,7 @@ def register_volumewise( Registered data with the same coordinates as input, input attributes, and added motion metadata in `attrs["reference_time"]` and `attrs["motion_params"]`. `motion_params` always carries per-frame - `final_metric_value` and `n_iterations` columns. When + `final_metric_value`, `n_iterations`, and `status` columns. When `keep_diagnostics=True`, `attrs["registration_diagnostics"]` also carries a list of [`RegistrationDiagnostics`][confusius.registration.RegistrationDiagnostics] @@ -187,49 +203,95 @@ def register_volumewise( progress_context = joblib_progress("Registering volumes...", total=n_frames) - with progress_context: - results = cast( - "list[tuple[xr.DataArray, npt.NDArray[np.floating] | None, RegistrationDiagnostics]]", # noqa: E501 - Parallel(n_jobs=n_jobs)( - delayed(register_volume)( - volume, - ref_da, - transform_type=transform, + parallel_kwargs: dict[str, object] = {"n_jobs": n_jobs} + if abort_event is not None or progress_reporter is not None: + # Use threads when cancellation or progress reporting is enabled so + # every worker sees the shared reporter / event instance. + parallel_kwargs["prefer"] = "threads" + + def _register_one( + frame_index: int, + volume: xr.DataArray, + ) -> tuple[ + int, xr.DataArray, npt.NDArray[np.floating] | None, RegistrationDiagnostics + ]: + # Once aborted, skip cheaply: building SimpleITK images and resampling is + # pure-Python/GIL-bound work that, multiplied across joblib threads, starves the + # GUI thread. Return the original frame with a zero-iteration "aborted" + # diagnostic instead. + if abort_event is not None and abort_event.is_set(): + return ( + frame_index, + volume, + None, + RegistrationDiagnostics( metric=metric, - number_of_histogram_bins=number_of_histogram_bins, - learning_rate=learning_rate, - number_of_iterations=number_of_iterations, - convergence_minimum_value=convergence_minimum_value, - convergence_window_size=convergence_window_size, - initialization=initialization, - optimizer_weights=optimizer_weights, - use_multi_resolution=use_multi_resolution, - shrink_factors=shrink_factors, - smoothing_sigmas=smoothing_sigmas, - resample=True, - resample_interpolation=resample_interpolation, - # Restrict SimpleITK to 1 thread per worker to avoid - # over-subscribing the CPU when joblib spawns many workers. - sitk_threads=1, - show_progress=False, - ) - for volume in data_moved - ), + metric_values=np.empty(0, dtype=float), + final_metric_value=float("nan"), + n_iterations=0, + stop_condition="Registration aborted before frame started.", + status="aborted", + ), + ) + registered_da, frame_affine, frame_diag = register_volume( + volume, + ref_da, + transform_type=transform, + metric=metric, + number_of_histogram_bins=number_of_histogram_bins, + learning_rate=learning_rate, + number_of_iterations=number_of_iterations, + convergence_minimum_value=convergence_minimum_value, + convergence_window_size=convergence_window_size, + initialization=initialization, + optimizer_weights=optimizer_weights, + use_multi_resolution=use_multi_resolution, + shrink_factors=shrink_factors, + smoothing_sigmas=smoothing_sigmas, + resample=True, + resample_interpolation=resample_interpolation, + # Restrict SimpleITK to 1 thread per worker to avoid + # over-subscribing the CPU when joblib spawns many workers. + sitk_threads=1, + show_progress=False, + abort_event=abort_event, ) + return frame_index, registered_da, frame_affine, frame_diag arr = data_moved.values - output = np.zeros_like(arr) - affines: list[npt.NDArray[np.floating] | None] = [] - final_metric_values: list[float] = [] - n_iterations_per_frame: list[int] = [] - diagnostics: list[RegistrationDiagnostics] = [] - for t, (registered_da, frame_affine, frame_diag) in enumerate(results): - output[t] = registered_da.values - affines.append(frame_affine) - final_metric_values.append(frame_diag.final_metric_value) - n_iterations_per_frame.append(frame_diag.n_iterations) - if keep_diagnostics: - diagnostics.append(frame_diag) + # Aborted/un-started frames are left blank (filled with the data minimum, + # i.e. background) rather than copying the unregistered input, so the partial + # result visibly shows which frames were skipped. + output = np.full_like(arr, arr.min()) + affines: list[npt.NDArray[np.floating] | None] = [None] * n_frames + final_metric_values = [float("nan")] * n_frames + n_iterations_per_frame = [0] * n_frames + statuses = ["aborted"] * n_frames + diagnostics: list[RegistrationDiagnostics | None] = [None] * n_frames + + try: + with progress_context: + results = Parallel(return_as="generator_unordered", **parallel_kwargs)( + delayed(_register_one)(t, volume) + for t, volume in enumerate(data_moved) + if abort_event is None or not abort_event.is_set() + ) + for t, registered_da, frame_affine, frame_diag in results: + skipped = ( + frame_diag.status == "aborted" and frame_diag.n_iterations == 0 + ) + if not skipped: + output[t] = registered_da.values + affines[t] = None if skipped else frame_affine + final_metric_values[t] = frame_diag.final_metric_value + n_iterations_per_frame[t] = frame_diag.n_iterations + statuses[t] = frame_diag.status + diagnostics[t] = frame_diag + if progress_reporter is not None: + progress_reporter.frame_completed(t, registered_da, frame_diag) + finally: + if progress_reporter is not None: + progress_reporter.close() time_coords = ( data_moved.coords["time"].values if "time" in data_moved.coords else None @@ -242,6 +304,7 @@ def register_volumewise( # for spotting frames that failed to converge, so we always keep them. motion_df["final_metric_value"] = final_metric_values motion_df["n_iterations"] = n_iterations_per_frame + motion_df["status"] = statuses result = xr.DataArray( output, @@ -255,6 +318,8 @@ def register_volumewise( if keep_diagnostics: # The full diagnostics list carries every frame's optimizer metric # trace, which adds up over long recordings — gated behind the flag. - result.attrs["registration_diagnostics"] = diagnostics + result.attrs["registration_diagnostics"] = [ + cast("RegistrationDiagnostics", d) for d in diagnostics + ] return result.transpose(*data.dims) diff --git a/src/confusius/registration/volumewise_progress.py b/src/confusius/registration/volumewise_progress.py new file mode 100644 index 00000000..6666e099 --- /dev/null +++ b/src/confusius/registration/volumewise_progress.py @@ -0,0 +1,42 @@ +"""Progress reporting protocol for `register_volumewise`.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Protocol + +if TYPE_CHECKING: + import xarray as xr + + from confusius.registration import RegistrationDiagnostics + + +class VolumewiseProgressReporter(Protocol): + """Duck-typed contract for `register_volumewise` progress reporting. + + Implementations may be called from worker threads when volumewise + registration runs in parallel. Any GUI updates must therefore be marshalled + via thread-safe mechanisms such as Qt signals. + """ + + def frame_completed( + self, + frame_index: int, + registered_frame: "xr.DataArray", + diagnostics: "RegistrationDiagnostics", + ) -> None: + """Report that one frame finished and provide its registered output. + + Parameters + ---------- + frame_index : int + Index of the completed frame. + registered_frame : xarray.DataArray + Registered frame output. + diagnostics : confusius.registration.RegistrationDiagnostics + Diagnostics collected for the completed frame. + """ + ... + + def close(self) -> None: + """Report that the full volumewise run has ended.""" + ... diff --git a/src/confusius/xarray/registration.py b/src/confusius/xarray/registration.py index 17417131..d0cfe865 100644 --- a/src/confusius/xarray/registration.py +++ b/src/confusius/xarray/registration.py @@ -1,6 +1,7 @@ """Xarray accessor for registration.""" -from collections.abc import Sequence +from collections.abc import Callable, Sequence +from threading import Event from typing import Literal import numpy as np @@ -8,8 +9,10 @@ import xarray as xr from confusius.registration.diagnostics import RegistrationDiagnostics +from confusius.registration.progress import RegistrationProgress from confusius.registration.volume import register_volume from confusius.registration.volumewise import register_volumewise +from confusius.registration.volumewise_progress import VolumewiseProgressReporter class FUSIRegistrationAccessor: @@ -34,6 +37,8 @@ def to_volume( self, fixed: xr.DataArray, *, + fixed_mask: xr.DataArray | None = None, + moving_mask: xr.DataArray | None = None, transform: Literal["translation", "rigid", "affine", "bspline"] = "rigid", metric: Literal["correlation", "mattes_mi"] = "correlation", number_of_histogram_bins: int = 50, @@ -51,10 +56,13 @@ def to_volume( smoothing_sigmas: Sequence[int] = (6, 2, 1), resample: bool = False, resample_interpolation: Literal["linear", "bspline"] = "linear", + fill_value: float | None = None, + sitk_threads: int = -1, show_progress: bool = False, plot_metric: bool = True, plot_composite: bool = True, - fill_value: float | None = None, + progress_plotter: Callable[..., RegistrationProgress] | None = None, + abort_event: Event | None = None, ) -> "tuple[xr.DataArray, npt.NDArray[np.floating] | xr.DataArray | None, RegistrationDiagnostics]": # noqa: E501 """Register this volume to a fixed reference volume. @@ -62,6 +70,10 @@ def to_volume( ---------- fixed : xarray.DataArray Reference volume to register to. + fixed_mask : xarray.DataArray, optional + Boolean mask for the fixed volume. + moving_mask : xarray.DataArray, optional + Boolean mask for this moving volume. transform : {"translation", "rigid", "affine", "bspline"}, default: "rigid" Type of transform to use for registration. metric : {"correlation", "mattes_mi"}, default: "correlation" @@ -69,8 +81,10 @@ def to_volume( number_of_histogram_bins : int, default: 50 Number of histogram bins (only used when `metric="mattes_mi"`). learning_rate : float or "auto", default: "auto" - Optimizer step size in normalised units (after `SetOptimizerScalesFromPhysicalShift`). - `"auto"` re-estimates the rate at every iteration. + Optimizer step size in normalised units (after + `SetOptimizerScalesFromPhysicalShift`). `"auto"` re-estimates the rate at + every iteration. A float uses that value directly; if registration diverges + or fails to converge, reduce it. number_of_iterations : int, default: 100 Maximum number of optimizer iterations. convergence_minimum_value : float, default: 1e-6 @@ -113,6 +127,12 @@ def to_volume( estimated and the moving volume is returned unchanged. resample_interpolation : {"linear", "bspline"}, default: "linear" Interpolation method used for the final resample step. + fill_value : float, optional + Fill value for voxels outside the moving image's field of view after + resampling. If not provided, defaults to the minimum of the moving + image. See [`register_volume`][confusius.registration.register_volume]. + sitk_threads : int, default: -1 + Number of threads SimpleITK may use internally. show_progress : bool, default: False Whether to display a live progress plot during registration. plot_metric : bool, default: True @@ -121,10 +141,12 @@ def to_volume( plot_composite : bool, default: True Whether to include a fixed/moving composite overlay in the progress plot. Ignored when `show_progress=False`. - fill_value : float, optional - Fill value for voxels outside the moving image's field of view after - resampling. If not provided, defaults to the minimum of the moving - image. See [`register_volume`][confusius.registration.register_volume]. + progress_plotter : callable, optional + Custom progress reporter factory. If not provided, the default + [`MatplotlibRegistrationProgressPlotter`][confusius.registration.MatplotlibRegistrationProgressPlotter] + is used. See [`register_volume`][confusius.registration.register_volume]. + abort_event : threading.Event, optional + Cooperative cancellation flag. Returns ------- @@ -149,6 +171,8 @@ def to_volume( return register_volume( self._obj, fixed, + fixed_mask=fixed_mask, + moving_mask=moving_mask, transform_type=transform, metric=metric, number_of_histogram_bins=number_of_histogram_bins, @@ -164,10 +188,13 @@ def to_volume( smoothing_sigmas=smoothing_sigmas, resample=resample, resample_interpolation=resample_interpolation, + fill_value=fill_value, + sitk_threads=sitk_threads, show_progress=show_progress, plot_metric=plot_metric, plot_composite=plot_composite, - fill_value=fill_value, + progress_plotter=progress_plotter, + abort_event=abort_event, ) def volumewise( @@ -178,7 +205,7 @@ def volumewise( transform: Literal["translation", "rigid", "affine"] = "rigid", metric: Literal["correlation", "mattes_mi"] = "correlation", number_of_histogram_bins: int = 50, - learning_rate: float | Literal["auto"] = "auto", + learning_rate: float | Literal["auto"] = 0.01, number_of_iterations: int = 100, convergence_minimum_value: float = 1e-6, convergence_window_size: int = 10, @@ -190,6 +217,8 @@ def volumewise( smoothing_sigmas: Sequence[int] = (6, 2, 1), resample_interpolation: Literal["linear", "bspline"] = "linear", show_progress: bool = True, + progress_reporter: VolumewiseProgressReporter | None = None, + abort_event: Event | None = None, keep_diagnostics: bool = False, ) -> xr.DataArray: """Register all volumes to a reference time point. @@ -207,9 +236,11 @@ def volumewise( Similarity metric for registration. number_of_histogram_bins : int, default: 50 Number of histogram bins (only used when `metric="mattes_mi"`). - learning_rate : float or "auto", default: "auto" - Optimizer step size in normalised units (after `SetOptimizerScalesFromPhysicalShift`). - `"auto"` re-estimates the rate at every iteration. + learning_rate : float or "auto", default: 0.01 + Optimizer step size in normalised units (after + `SetOptimizerScalesFromPhysicalShift`). `"auto"` re-estimates the rate at + every iteration. A float uses that value directly; if registration diverges + or fails to converge, reduce it. number_of_iterations : int, default: 100 Maximum number of optimizer iterations. convergence_minimum_value : float, default: 1e-6 @@ -244,6 +275,11 @@ def volumewise( Interpolation method used for the final resample step. show_progress : bool, default: True Whether to display a progress bar while registering volumes. + progress_reporter : VolumewiseProgressReporter, optional + Thread-safe reporter notified whenever one frame completes. If not + provided, no per-frame callback is used. + abort_event : threading.Event, optional + Cooperative cancellation flag shared across frames. keep_diagnostics : bool, default: False Whether to keep per-frame registration diagnostics on the result. See @@ -278,5 +314,7 @@ def volumewise( smoothing_sigmas=smoothing_sigmas, resample_interpolation=resample_interpolation, show_progress=show_progress, + progress_reporter=progress_reporter, + abort_event=abort_event, keep_diagnostics=keep_diagnostics, ) diff --git a/tests/unit/test_datasets/test_registry.py b/tests/unit/test_datasets/test_registry.py index 40b186c6..30faae14 100644 --- a/tests/unit/test_datasets/test_registry.py +++ b/tests/unit/test_datasets/test_registry.py @@ -2,7 +2,6 @@ from __future__ import annotations -import pytest from confusius.datasets import list_datasets from confusius.datasets._registry import _REGISTRY diff --git a/tests/unit/test_io/test_loadsave.py b/tests/unit/test_io/test_loadsave.py index 8b8f289c..1a9dc074 100644 --- a/tests/unit/test_io/test_loadsave.py +++ b/tests/unit/test_io/test_loadsave.py @@ -1,6 +1,6 @@ """Unit tests for confusius.io.loadsave module.""" -from unittest.mock import MagicMock, call, patch +from unittest.mock import MagicMock, patch import numpy as np import numpy.testing as npt diff --git a/tests/unit/test_napari/test_registration_metric_plotter.py b/tests/unit/test_napari/test_registration_metric_plotter.py new file mode 100644 index 00000000..b3e6e72e --- /dev/null +++ b/tests/unit/test_napari/test_registration_metric_plotter.py @@ -0,0 +1,51 @@ +"""Unit tests for the bottom-dock registration metric plotter.""" + +from __future__ import annotations + +import pytest + + +@pytest.fixture +def registration_metric_plotter(make_napari_viewer_proxy): + from confusius._napari._registration._metric_plotter import ( + RegistrationMetricPlotter, + ) + + viewer = make_napari_viewer_proxy() + return RegistrationMetricPlotter(viewer) + + +class TestRegistrationMetricPlotterBuffer: + """Pure-logic: add_metric / reset / metric_values.""" + + def test_empty_after_construction(self, registration_metric_plotter) -> None: + assert registration_metric_plotter.metric_values == [] + + def test_add_metric_appends_value(self, registration_metric_plotter) -> None: + registration_metric_plotter.add_metric(0.5) + registration_metric_plotter.add_metric(0.25) + registration_metric_plotter.add_metric(0.1) + assert registration_metric_plotter.metric_values == [0.5, 0.25, 0.1] + + def test_metric_values_returns_a_copy( + self, registration_metric_plotter + ) -> None: + registration_metric_plotter.add_metric(1.0) + snapshot = registration_metric_plotter.metric_values + snapshot.append(99.0) + assert registration_metric_plotter.metric_values == [1.0] + + def test_reset_clears_buffer(self, registration_metric_plotter) -> None: + registration_metric_plotter.add_metric(0.5) + registration_metric_plotter.add_metric(0.25) + registration_metric_plotter.reset() + assert registration_metric_plotter.metric_values == [] + + def test_add_metric_after_reset_starts_new_run( + self, registration_metric_plotter + ) -> None: + registration_metric_plotter.add_metric(0.5) + registration_metric_plotter.add_metric(0.25) + registration_metric_plotter.reset() + registration_metric_plotter.add_metric(0.1) + assert registration_metric_plotter.metric_values == [0.1] diff --git a/tests/unit/test_napari/test_registration_panel.py b/tests/unit/test_napari/test_registration_panel.py new file mode 100644 index 00000000..33a673af --- /dev/null +++ b/tests/unit/test_napari/test_registration_panel.py @@ -0,0 +1,1487 @@ +"""Unit tests for the napari registration panel.""" + +from __future__ import annotations + +from threading import Event +from typing import Any, cast + +import numpy as np +import pytest +import xarray as xr +from qtpy.QtWidgets import QApplication + +from confusius._napari._registration._panel_progress import ( + create_volume_progress_plotter, + setup_volumewise_progress, + teardown_volume_progress, + update_progress_layer, +) +from confusius._napari._registration._panel_results import on_registration_finished +from confusius._napari._registration._panel_transforms import ( + apply_selected_inverse_transform, + apply_selected_transform, + refresh_transform_controls, +) +from confusius._napari._registration._transform_payloads import ( + get_affine_transform_from_payload, + get_bspline_transform_from_payload, + get_input_grid_from_payload, + get_output_grid_from_payload, + load_transform_payload, + make_affine_transform_payload, + make_bspline_transform_payload, + save_transform_payload, +) +from confusius.registration import ( + RegistrationDiagnostics, + resample_like, + resample_volume, +) + + +@pytest.fixture +def viewer(make_napari_viewer_proxy): + return make_napari_viewer_proxy() + + +@pytest.fixture +def real_viewer(make_napari_viewer): + return make_napari_viewer() + + +@pytest.fixture +def registration_panel(viewer): + from confusius._napari._registration._panel import RegistrationPanel + + return RegistrationPanel(viewer) + + +@pytest.fixture +def real_registration_panel(real_viewer): + from confusius._napari._registration._panel import RegistrationPanel + + return RegistrationPanel(real_viewer) + + +def _FakeDiagnostics( + *, + metric: str = "correlation", + metric_values: np.ndarray | None = None, + final_metric_value: float = -1.0, + n_iterations: int = 1, + stop_condition: str = "done", + status: str = "completed", +) -> RegistrationDiagnostics: + return RegistrationDiagnostics( + metric=cast("Any", metric), + metric_values=np.array([-1.0]) if metric_values is None else metric_values, + final_metric_value=final_metric_value, + n_iterations=n_iterations, + stop_condition=stop_condition, + status=cast("Any", status), + ) + + +def _make_bspline_transform() -> xr.DataArray: + return xr.DataArray( + np.arange(2 * 3 * 4, dtype=float).reshape(2, 3, 4), + dims=["component", "y", "x"], + coords={ + "component": xr.DataArray([0, 1], dims=["component"]), + "y": xr.DataArray(np.arange(3) * 0.2, dims=["y"]), + "x": xr.DataArray(np.arange(4) * 0.1, dims=["x"]), + }, + attrs={ + "transform_type": "bspline_transform", + "order": 3, + "direction": [[1.0, 0.0], [0.0, 1.0]], + "affines": {"bspline_initialization": np.eye(3).tolist()}, + }, + ) + + +def _install_immediate_thread_worker(monkeypatch: pytest.MonkeyPatch) -> None: + class _Signal: + def __init__(self) -> None: + self._slots: list[Any] = [] + + def connect(self, slot: Any) -> None: + self._slots.append(slot) + + def emit(self, *args: Any) -> None: + for slot in list(self._slots): + slot(*args) + + class _Worker: + def __init__(self, func: Any, args: tuple[Any, ...], kwargs: dict[str, Any]): + self._func = func + self._args = args + self._kwargs = kwargs + self.returned = _Signal() + self.errored = _Signal() + self.finished = _Signal() + + def start(self) -> None: + try: + self.returned.emit(self._func(*self._args, **self._kwargs)) + except Exception as exc: # noqa: BLE001 + self.errored.emit(exc) + finally: + self.finished.emit() + + def _thread_worker(func: Any) -> Any: + def _runner(*args: Any, **kwargs: Any) -> _Worker: + return _Worker(func, args, kwargs) + + return _runner + + monkeypatch.setattr( + "confusius._napari._registration._panel_transforms.thread_worker", + _thread_worker, + ) + + +class TestRefreshLayers: + def test_combo_populated_on_layer_add(self, viewer, registration_panel): + assert registration_panel._moving_combo.count() == 0 + viewer.add_image(np.zeros((4, 6, 8)), name="vol") + assert registration_panel._moving_combo.count() == 1 + assert registration_panel._moving_combo.itemText(0) == "vol" + + def test_mask_combos_only_list_labels_layers(self, viewer, registration_panel): + viewer.add_image(np.zeros((4, 6, 8)), name="vol") + viewer.add_labels(np.zeros((4, 6, 8), dtype=np.int32), name="mask") + + registration_panel._refresh_layers() + + assert registration_panel._fixed_mask_combo.count() == 2 + assert registration_panel._fixed_mask_combo.itemText(0) == "" + assert registration_panel._fixed_mask_combo.itemText(1) == "mask" + assert registration_panel._moving_mask_combo.count() == 2 + assert registration_panel._moving_mask_combo.itemText(1) == "mask" + + def test_ignores_lazy_non_numpy_layers(self, viewer, registration_panel): + import dask.array as da + + viewer.add_image(np.zeros((4, 6, 8)), name="vol") + viewer.add_image(da.zeros((5, 4, 6), chunks=(1, 4, 6)), name="video") + + registration_panel._refresh_layers() + + assert registration_panel._moving_combo.count() == 1 + assert registration_panel._moving_combo.itemText(0) == "vol" + + +class TestOperationMode: + def test_panel_switch_shows_one_subpanel(self, registration_panel): + assert not registration_panel._register_panel.isHidden() + assert registration_panel._transforms_panel.isHidden() + + registration_panel._transforms_panel_radio.setChecked(True) + registration_panel._on_panel_changed() + + assert registration_panel._register_panel.isHidden() + assert not registration_panel._transforms_panel.isHidden() + + def test_volumewise_hides_fixed_selector(self, registration_panel): + registration_panel._time_series_radio.setChecked(True) + assert registration_panel._fixed_combo.isHidden() + assert not registration_panel._reference_time_spin.isHidden() + assert not registration_panel._n_jobs_row.isHidden() + + def test_between_scan_shows_masks_and_sitk_threads(self, registration_panel): + registration_panel._advanced_toggle.setChecked(True) + + assert not registration_panel._fixed_mask_row.isHidden() + assert not registration_panel._moving_mask_row.isHidden() + assert not registration_panel._sitk_threads_row.isHidden() + + registration_panel._time_series_radio.setChecked(True) + + assert registration_panel._fixed_mask_row.isHidden() + assert registration_panel._moving_mask_row.isHidden() + assert registration_panel._sitk_threads_row.isHidden() + + def test_volume_shows_fixed_selector(self, registration_panel): + registration_panel._time_series_radio.setChecked(True) + registration_panel._single_volume_radio.setChecked(True) + assert not registration_panel._fixed_combo.isHidden() + assert registration_panel._reference_time_spin.isHidden() + assert registration_panel._n_jobs_row.isHidden() + + def test_learning_rate_auto_disables_edit(self, registration_panel): + assert registration_panel._learning_rate_auto_check.isChecked() + assert not registration_panel._learning_rate_edit.isEnabled() + registration_panel._learning_rate_auto_check.setChecked(False) + assert registration_panel._learning_rate_edit.isEnabled() + + def test_volumewise_learning_rate_defaults_to_fixed_0_01(self, registration_panel): + registration_panel._time_series_radio.setChecked(True) + + assert registration_panel._learning_rate_auto_check.isHidden() + assert not registration_panel._learning_rate_auto_check.isChecked() + assert registration_panel._learning_rate_edit.isEnabled() + assert registration_panel._learning_rate_edit.value() == pytest.approx(0.01) + + def test_mode_switch_preserves_session_parameters(self, registration_panel): + registration_panel._time_series_radio.setChecked(True) + registration_panel._learning_rate_edit.setValue(0.23) + registration_panel._n_jobs_spin.setValue(3) + registration_panel._scale_combo.setCurrentText("square root") + + registration_panel._single_volume_radio.setChecked(True) + registration_panel._learning_rate_auto_check.setChecked(True) + registration_panel._learning_rate_edit.setValue(0.42) + registration_panel._scale_combo.setCurrentText("none") + registration_panel._time_series_radio.setChecked(True) + + assert registration_panel._learning_rate_auto_check.isHidden() + assert not registration_panel._learning_rate_auto_check.isChecked() + assert registration_panel._learning_rate_edit.value() == pytest.approx(0.23) + assert registration_panel._n_jobs_spin.value() == 3 + assert registration_panel._scale_combo.currentText() == "square root" + + def test_advanced_group_is_collapsed_by_default(self, registration_panel): + assert not registration_panel._advanced_toggle.isChecked() + assert registration_panel._advanced_content.isHidden() + + registration_panel._advanced_toggle.click() + + assert not registration_panel._advanced_content.isHidden() + + def test_opening_advanced_group_does_not_widen_panel_minimum( + self, registration_panel + ): + # Regression test for the issue #183 overflow pattern: advanced rows + # must wrap on narrow docks instead of raising the panel's minimum + # width, which forced horizontal overflow in the sidebar scroll area. + registration_panel.show() + QApplication.processEvents() + closed_min_width = registration_panel.minimumSizeHint().width() + registration_panel._advanced_toggle.setChecked(True) + QApplication.processEvents() + assert registration_panel.minimumSizeHint().width() <= closed_min_width + + def test_scientific_notation_spinboxes_parse_values(self, registration_panel): + registration_panel._learning_rate_auto_check.setChecked(False) + registration_panel._learning_rate_edit.lineEdit().setText("1e-5") + registration_panel._learning_rate_edit.interpretText() + assert registration_panel._learning_rate_edit.value() == pytest.approx(1e-5) + + registration_panel._convergence_min_edit.lineEdit().setText("2.5e-7") + registration_panel._convergence_min_edit.interpretText() + assert registration_panel._convergence_min_edit.value() == pytest.approx(2.5e-7) + + def test_default_parameter_values(self, registration_panel): + assert registration_panel._transform_combo.currentText() == "rigid" + assert registration_panel._scale_combo.currentText() == "decibel" + assert registration_panel._learning_rate_edit.minimum() == pytest.approx(1e-10) + assert registration_panel._learning_rate_edit.value() == pytest.approx(0.1) + assert registration_panel._convergence_min_edit.minimum() == pytest.approx( + 1e-10 + ) + assert registration_panel._convergence_min_edit.value() == pytest.approx(1e-6) + assert registration_panel._iterations_spin.singleStep() == 100 + + registration_panel._time_series_radio.setChecked(True) + + assert registration_panel._transform_combo.currentText() == "rigid" + + def test_scale_preprocessing_resets_gamma_for_previews( + self, real_viewer, real_registration_panel + ): + moving_data = xr.DataArray( + np.ones((4, 6), dtype=np.float32), + dims=["y", "x"], + coords={ + "y": xr.DataArray(np.arange(4) * 0.2, dims=["y"]), + "x": xr.DataArray(np.arange(6) * 0.1, dims=["x"]), + }, + ) + fixed = xr.DataArray( + 2 * np.ones((4, 6), dtype=np.float32), + dims=["y", "x"], + coords=moving_data.coords, + ) + moving = real_viewer.add_image(moving_data.values, name="moving") + fixed_layer = real_viewer.add_image(fixed.values, name="fixed") + moving.gamma = 0.4 + fixed_layer.gamma = 0.6 + + create_volume_progress_plotter( + real_registration_panel, + moving_layer=moving, + fixed_layer=fixed_layer, + moving=moving_data, + fixed=fixed, + layer_name="Registered (rigid)", + scale_mode="sqrt", + ) + assert real_viewer.layers["Fixed"].gamma == pytest.approx(1.0) + assert real_viewer.layers["Moving"].gamma == pytest.approx(1.0) + assert real_viewer.layers["Registered (rigid)"].gamma == pytest.approx(1.0) + + teardown_volume_progress(real_registration_panel) + create_volume_progress_plotter( + real_registration_panel, + moving_layer=moving, + fixed_layer=fixed_layer, + moving=moving_data, + fixed=fixed, + layer_name="Registered (rigid)", + scale_mode="off", + ) + assert real_viewer.layers["Fixed"].gamma == pytest.approx(0.6) + assert real_viewer.layers["Moving"].gamma == pytest.approx(0.4) + assert real_viewer.layers["Registered (rigid)"].gamma == pytest.approx(0.4) + + def test_create_volume_progress_plotter_preserves_camera_view( + self, real_viewer, real_registration_panel + ): + moving_data = xr.DataArray( + np.ones((5, 4, 6), dtype=np.float32), + dims=["z", "y", "x"], + coords={ + "z": xr.DataArray(np.arange(5) * 0.3, dims=["z"]), + "y": xr.DataArray(np.arange(4) * 0.2, dims=["y"]), + "x": xr.DataArray(np.arange(6) * 0.1, dims=["x"]), + }, + ) + fixed = xr.DataArray( + 2 * np.ones((5, 4, 6), dtype=np.float32), + dims=["z", "y", "x"], + coords=moving_data.coords, + ) + moving = real_viewer.add_image(moving_data.values, name="moving") + fixed_layer = real_viewer.add_image(fixed.values, name="fixed") + + # User navigates to a custom 3D view before launching the run. + real_viewer.dims.ndisplay = 3 + real_viewer.camera.center = (1.0, 2.0, 3.0) + real_viewer.camera.zoom = 7.0 + before = ( + tuple(real_viewer.camera.center), + real_viewer.camera.zoom, + real_viewer.dims.ndisplay, + ) + + create_volume_progress_plotter( + real_registration_panel, + moving_layer=moving, + fixed_layer=fixed_layer, + moving=moving_data, + fixed=fixed, + layer_name="Registered (rigid)", + scale_mode="off", + ) + + after = ( + tuple(real_viewer.camera.center), + real_viewer.camera.zoom, + real_viewer.dims.ndisplay, + ) + assert after == before + + def test_metric_specific_rows_follow_metric(self, registration_panel): + registration_panel._advanced_toggle.setChecked(True) + assert registration_panel._metric_combo.currentText() == "correlation" + assert registration_panel._histogram_bins_row.isHidden() + + registration_panel._metric_combo.setCurrentText("mattes_mi") + assert not registration_panel._histogram_bins_row.isHidden() + + def test_multi_resolution_toggle_hides_dependent_inputs(self, registration_panel): + registration_panel._advanced_toggle.setChecked(True) + assert not registration_panel._multi_resolution_check.isChecked() + assert registration_panel._shrink_factors_row.isHidden() + assert registration_panel._smoothing_sigmas_row.isHidden() + + registration_panel._multi_resolution_check.setChecked(True) + + assert not registration_panel._shrink_factors_row.isHidden() + assert not registration_panel._smoothing_sigmas_row.isHidden() + + def test_mesh_size_is_basic_and_only_visible_for_bspline(self, registration_panel): + assert registration_panel._mesh_size_row.isHidden() + assert registration_panel._optimizer_weights_check.isEnabled() + + registration_panel._transform_combo.setCurrentText("bspline") + + assert not registration_panel._mesh_size_row.isHidden() + assert not registration_panel._optimizer_weights_check.isEnabled() + + registration_panel._transform_combo.setCurrentText("rigid") + + assert registration_panel._mesh_size_row.isHidden() + assert registration_panel._optimizer_weights_check.isEnabled() + + registration_panel._time_series_radio.setChecked(True) + + assert registration_panel._mesh_size_row.isHidden() + + +class TestRunRegistration: + def test_mask_buttons_create_named_layers(self, viewer, registration_panel): + moving = xr.DataArray( + np.zeros((2, 4, 6, 8), dtype=np.float32), + dims=["time", "z", "y", "x"], + coords={ + "time": xr.DataArray(np.arange(2), dims=["time"]), + "z": xr.DataArray(np.arange(4) * 0.3, dims=["z"]), + "y": xr.DataArray(np.arange(6) * 0.2, dims=["y"]), + "x": xr.DataArray(np.arange(8) * 0.1, dims=["x"]), + }, + ) + layer = viewer.add_image( + moving.values, name="moving", metadata={"xarray": moving} + ) + layer.scale = (1.0, 0.3, 0.2, 0.1) + layer.translate = (0.0, 1.0, 2.0, 3.0) + + registration_panel._new_moving_mask_btn.click() + registration_panel._new_fixed_mask_btn.click() + + moving_mask = viewer.layers["Moving mask"] + fixed_mask = viewer.layers["Fixed mask"] + assert np.asarray(moving_mask.data).shape == (4, 6, 8) + assert tuple(moving_mask.scale) == (0.3, 0.2, 0.1) + assert tuple(moving_mask.translate) == (1.0, 2.0, 3.0) + assert np.asarray(fixed_mask.data).shape == (4, 6, 8) + + def test_between_scan_run_uses_selected_initial_transform( + self, viewer, registration_panel, monkeypatch + ): + moving = xr.DataArray( + np.zeros((4, 6), dtype=np.float32), + dims=["y", "x"], + coords={ + "y": xr.DataArray(np.arange(4) * 0.2, dims=["y"]), + "x": xr.DataArray(np.arange(6) * 0.1, dims=["x"]), + }, + ) + fixed = xr.DataArray( + np.ones((4, 6), dtype=np.float32), + dims=["y", "x"], + coords={ + "y": xr.DataArray(np.arange(4) * 0.2, dims=["y"]), + "x": xr.DataArray(np.arange(6) * 0.1, dims=["x"]), + }, + ) + affine = np.array( + [[1.0, 0.0, 0.5], [0.0, 1.0, -0.25], [0.0, 0.0, 1.0]], + dtype=float, + ) + transform_payload = make_affine_transform_payload( + affine, + reference=fixed, + source_layer_name="moving", + target_layer_name="fixed", + operation="register_volume", + transform_model="rigid", + metric="correlation", + diagnostics=_FakeDiagnostics(), + ) + + viewer.add_image(moving.values, name="moving", metadata={"xarray": moving}) + viewer.add_image(fixed.values, name="fixed", metadata={"xarray": fixed}) + viewer.add_image( + fixed.values, + name="Previous registered", + metadata={"confusius_transform": transform_payload}, + ) + registration_panel._refresh_layers() + refresh_transform_controls(registration_panel) + registration_panel._moving_combo.setCurrentText("moving") + registration_panel._fixed_combo.setCurrentText("fixed") + registration_panel._scale_combo.setCurrentText("square root") + for i in range(registration_panel._initialization_combo.count()): + if registration_panel._initialization_combo.itemData(i) == ( + "layer", + "Previous registered", + ): + registration_panel._initialization_combo.setCurrentIndex(i) + break + + captured: dict[str, object] = {} + + class _FakeSignal: + def connect(self, _slot): + return None + + class _FakeWorker: + def __init__(self) -> None: + self.returned = _FakeSignal() + self.errored = _FakeSignal() + self.finished = _FakeSignal() + + def start(self) -> None: + return None + + def _fake_thread_worker(func): + def _runner(*args, **kwargs): + captured["func"] = func + captured["args"] = args + captured["kwargs"] = kwargs + return _FakeWorker() + + return _runner + + monkeypatch.setattr( + "confusius._napari._registration._panel.thread_worker", + _fake_thread_worker, + ) + monkeypatch.setattr( + "confusius._napari._registration._panel.create_volume_progress_plotter", + lambda *_args, **_kwargs: None, + ) + + registration_panel._run_registration() + + kwargs = cast("dict[str, Any]", captured["kwargs"]) + args = cast("tuple[Any, ...]", captured["args"]) + np.testing.assert_array_equal(kwargs["initialization"], affine) + np.testing.assert_allclose(args[0].values, np.sqrt(moving.values)) + np.testing.assert_allclose(args[1].values, np.sqrt(fixed.values)) + assert registration_panel._worker is not None + + def test_between_scan_run_uses_selected_manual_napari_transform( + self, viewer, registration_panel, monkeypatch + ): + moving = xr.DataArray( + np.zeros((2, 4, 6, 8), dtype=np.float32), + dims=["time", "z", "y", "x"], + coords={ + "time": xr.DataArray(np.arange(2), dims=["time"]), + "z": xr.DataArray(np.arange(4) * 0.3, dims=["z"]), + "y": xr.DataArray(np.arange(6) * 0.2, dims=["y"]), + "x": xr.DataArray(np.arange(8) * 0.1, dims=["x"]), + }, + ) + fixed = xr.DataArray( + np.ones((2, 4, 6, 8), dtype=np.float32), + dims=["time", "z", "y", "x"], + coords=moving.coords, + ) + + moving_layer = viewer.add_image( + moving.values, + name="moving", + metadata={"xarray": moving}, + ) + viewer.add_image(fixed.values, name="fixed", metadata={"xarray": fixed}) + + registration_panel._refresh_layers() + registration_panel._moving_combo.setCurrentText("moving") + registration_panel._fixed_combo.setCurrentText("fixed") + + manual_affine = np.eye(5) + manual_affine[0, 4] = 9.0 + manual_affine[1, 4] = 0.5 + manual_affine[2, 4] = -0.25 + manual_affine[3, 3] = 1.1 + manual_affine[3, 4] = 0.75 + moving_layer.affine = manual_affine + QApplication.processEvents() + + for i in range(registration_panel._initialization_combo.count()): + if registration_panel._initialization_combo.itemData(i) == ( + "manual", + "moving", + ): + registration_panel._initialization_combo.setCurrentIndex(i) + break + + captured: dict[str, object] = {} + + class _FakeSignal: + def connect(self, _slot): + return None + + class _FakeWorker: + def __init__(self) -> None: + self.returned = _FakeSignal() + self.errored = _FakeSignal() + self.finished = _FakeSignal() + + def start(self) -> None: + return None + + def _fake_thread_worker(func): + def _runner(*args, **kwargs): + captured["func"] = func + captured["args"] = args + captured["kwargs"] = kwargs + return _FakeWorker() + + return _runner + + monkeypatch.setattr( + "confusius._napari._registration._panel.thread_worker", + _fake_thread_worker, + ) + monkeypatch.setattr( + "confusius._napari._registration._panel.create_volume_progress_plotter", + lambda *_args, **_kwargs: None, + ) + + registration_panel._run_registration() + + kwargs = cast("dict[str, Any]", captured["kwargs"]) + args = cast("tuple[Any, ...]", captured["args"]) + expected = np.array( + [ + [1.0, 0.0, 0.0, -0.5], + [0.0, 1.0, 0.0, 0.25], + [0.0, 0.0, 1.0 / 1.1, -0.75 / 1.1], + [0.0, 0.0, 0.0, 1.0], + ] + ) + np.testing.assert_allclose(kwargs["initialization"], expected) + assert args[0].dims == ("z", "y", "x") + assert registration_panel._worker is not None + + @pytest.mark.parametrize( + ("within_scan", "transform", "weights"), + [ + (False, "rigid", [0.1, 0.2, 0.3, 0.4, 0.5, 0.6]), + (True, "translation", [0.7, 0.8, 0.9]), + ], + ) + def test_run_passes_mode_specific_worker_kwargs( + self, viewer, registration_panel, monkeypatch, within_scan, transform, weights + ): + captured: dict[str, object] = {} + + class _FakeSignal: + def connect(self, _slot): + return None + + class _FakeWorker: + def __init__(self) -> None: + self.returned = _FakeSignal() + self.errored = _FakeSignal() + self.finished = _FakeSignal() + + def start(self) -> None: + return None + + def _fake_thread_worker(func): + def _runner(*args, **kwargs): + captured["func"] = func + captured["args"] = args + captured["kwargs"] = kwargs + return _FakeWorker() + + return _runner + + monkeypatch.setattr( + "confusius._napari._registration._panel.thread_worker", + _fake_thread_worker, + ) + + if within_scan: + moving = xr.DataArray( + np.zeros((2, 4, 6, 8), dtype=np.float32), + dims=["time", "z", "y", "x"], + coords={ + "time": xr.DataArray(np.arange(2), dims=["time"]), + "z": xr.DataArray(np.arange(4) * 0.3, dims=["z"]), + "y": xr.DataArray(np.arange(6) * 0.2, dims=["y"]), + "x": xr.DataArray(np.arange(8) * 0.1, dims=["x"]), + }, + ) + viewer.add_image( + moving.values, name="moving", metadata={"xarray": moving} + ) + registration_panel._time_series_radio.setChecked(True) + registration_panel._refresh_layers() + registration_panel._moving_combo.setCurrentText("moving") + monkeypatch.setattr( + "confusius._napari._registration._panel.setup_volumewise_progress", + lambda *_args, **_kwargs: None, + ) + else: + moving = xr.DataArray( + np.zeros((4, 6, 8), dtype=np.float32), + dims=["z", "y", "x"], + coords={ + "z": xr.DataArray(np.arange(4) * 0.3, dims=["z"]), + "y": xr.DataArray(np.arange(6) * 0.2, dims=["y"]), + "x": xr.DataArray(np.arange(8) * 0.1, dims=["x"]), + }, + ) + fixed = xr.DataArray( + np.ones((4, 6, 8), dtype=np.float32), + dims=["z", "y", "x"], + coords=moving.coords, + ) + viewer.add_image( + moving.values, name="moving", metadata={"xarray": moving} + ) + viewer.add_image(fixed.values, name="fixed", metadata={"xarray": fixed}) + viewer.add_labels( + np.ones((4, 6, 8), dtype=np.int32), name="fixed mask" + ) + viewer.add_labels( + np.ones((4, 6, 8), dtype=np.int32), name="moving mask" + ) + registration_panel._refresh_layers() + registration_panel._moving_combo.setCurrentText("moving") + registration_panel._fixed_combo.setCurrentText("fixed") + registration_panel._fixed_mask_combo.setCurrentText("fixed mask") + registration_panel._moving_mask_combo.setCurrentText("moving mask") + registration_panel._sitk_threads_spin.setValue(3) + monkeypatch.setattr( + "confusius._napari._registration._panel.create_volume_progress_plotter", + lambda *_args, **_kwargs: None, + ) + + registration_panel._transform_combo.setCurrentText(transform) + registration_panel._optimizer_weights_check.setChecked(True) + for spin, value in zip( + registration_panel._optimizer_weight_spins, + weights, + strict=False, + ): + spin.setValue(value) + + registration_panel._run_registration() + + kwargs = cast("dict[str, Any]", captured["kwargs"]) + assert kwargs["optimizer_weights"] == weights + if within_scan: + assert kwargs["reference_time"] == 0 + else: + assert kwargs["sitk_threads"] == 3 + assert kwargs["fixed_mask"].dtype == bool + assert kwargs["moving_mask"].dtype == bool + assert registration_panel._worker is not None + + +class TestAbort: + def test_abort_sets_cancellation_event(self, registration_panel): + registration_panel._worker = object() + registration_panel._abort_event = Event() + registration_panel._begin_work() + + assert registration_panel._run_btn.isHidden() + assert not registration_panel._abort_btn.isHidden() + + registration_panel._abort_registration() + + assert registration_panel._abort_event.is_set() + assert not registration_panel._abort_btn.isEnabled() + assert registration_panel._abort_btn.text() == "Aborting…" + + +class TestValidation: + def test_same_moving_and_fixed_is_flagged( + self, real_viewer, real_registration_panel + ): + real_viewer.add_image(np.zeros((4, 6), dtype=np.float32), name="same") + real_registration_panel._refresh_layers() + real_registration_panel._moving_combo.setCurrentText("same") + real_registration_panel._fixed_combo.setCurrentText("same") + + assert not real_registration_panel._validate_registration_selection() + assert not real_registration_panel._layer_validation.isHidden() + assert "must be different" in real_registration_panel._layer_validation.text() + + def test_within_scan_requires_time_dimension(self, viewer, registration_panel): + viewer.add_image(np.zeros((4, 6), dtype=np.float32), name="vol") + registration_panel._refresh_layers() + registration_panel._time_series_radio.setChecked(True) + registration_panel._moving_combo.setCurrentText("vol") + + assert not registration_panel._validate_registration_selection() + assert ( + "Within-scan registration requires" + in registration_panel._layer_validation.text() + ) + + def test_between_scans_accepts_time_series_by_averaging( + self, viewer, registration_panel + ): + moving = xr.DataArray( + np.zeros((3, 4, 6), dtype=np.float32), + dims=["time", "y", "x"], + coords={ + "time": xr.DataArray(np.arange(3), dims=["time"]), + "y": xr.DataArray(np.arange(4), dims=["y"]), + "x": xr.DataArray(np.arange(6), dims=["x"]), + }, + ) + fixed = xr.DataArray( + np.ones((3, 4, 6), dtype=np.float32), + dims=["time", "y", "x"], + coords={ + "time": xr.DataArray(np.arange(3), dims=["time"]), + "y": xr.DataArray(np.arange(4), dims=["y"]), + "x": xr.DataArray(np.arange(6), dims=["x"]), + }, + ) + viewer.add_image(moving.values, name="moving", metadata={"xarray": moving}) + viewer.add_image(fixed.values, name="fixed", metadata={"xarray": fixed}) + registration_panel._refresh_layers() + registration_panel._moving_combo.setCurrentText("moving") + registration_panel._fixed_combo.setCurrentText("fixed") + + assert registration_panel._validate_registration_selection() + + +class TestTransforms: + def test_affine_payload_roundtrip(self, tmp_path): + reference = xr.DataArray( + np.ones((4, 6), dtype=np.float32), + dims=["y", "x"], + coords={ + "y": xr.DataArray(np.arange(4) * 0.2, dims=["y"]), + "x": xr.DataArray(np.arange(6) * 0.1, dims=["x"]), + }, + ) + payload = make_affine_transform_payload( + np.eye(3), + reference=reference, + source=reference, + source_layer_name="moving", + target_layer_name="fixed", + operation="register_volume", + transform_model="rigid", + metric="correlation", + diagnostics=_FakeDiagnostics(), + ) + + path = tmp_path / "transform.json" + save_transform_payload(path, payload) + loaded = load_transform_payload(path) + + assert loaded["source_layer_name"] == "moving" + assert loaded["name"] == "moving → fixed (rigid)" + assert get_output_grid_from_payload(loaded)["shape"] == [4, 6] + assert get_input_grid_from_payload(loaded)["shape"] == [4, 6] + np.testing.assert_array_equal( + get_affine_transform_from_payload(loaded), np.eye(3) + ) + + def test_bspline_payload_roundtrip(self, tmp_path): + reference = xr.DataArray( + np.ones((3, 4), dtype=np.float32), + dims=["y", "x"], + coords={ + "y": xr.DataArray(np.arange(3) * 0.2, dims=["y"]), + "x": xr.DataArray(np.arange(4) * 0.1, dims=["x"]), + }, + ) + transform = _make_bspline_transform() + payload = make_bspline_transform_payload( + transform, + reference=reference, + source=reference, + source_layer_name="moving", + target_layer_name="fixed", + operation="register_volume", + transform_model="bspline", + metric="correlation", + diagnostics=_FakeDiagnostics(), + ) + + path = tmp_path / "bspline.zarr" + save_transform_payload(path, payload) + loaded = load_transform_payload(path) + + assert loaded["name"] == "moving → fixed (bspline)" + assert loaded["kind"] == "bspline" + assert get_output_grid_from_payload(loaded)["shape"] == [3, 4] + assert get_input_grid_from_payload(loaded)["shape"] == [3, 4] + xr.testing.assert_identical( + get_bspline_transform_from_payload(loaded), + transform.astype(float), + ) + + def test_bspline_transform_is_not_offered_for_initialization( + self, viewer, registration_panel + ): + moving = xr.DataArray( + np.zeros((3, 4), dtype=np.float32), + dims=["y", "x"], + coords={ + "y": xr.DataArray(np.arange(3) * 0.2, dims=["y"]), + "x": xr.DataArray(np.arange(4) * 0.1, dims=["x"]), + }, + ) + payload = make_bspline_transform_payload( + _make_bspline_transform(), + reference=moving, + source_layer_name="moving", + target_layer_name="fixed", + operation="register_volume", + transform_model="bspline", + metric="correlation", + diagnostics=_FakeDiagnostics(), + ) + viewer.add_image( + moving.values, + name="Registered (bspline)", + metadata={"xarray": moving, "confusius_transform": payload}, + ) + + refresh_transform_controls(registration_panel) + + transform_items = [ + registration_panel._transform_source_combo.itemText(i) + for i in range(registration_panel._transform_source_combo.count()) + ] + initialization_items = [ + registration_panel._initialization_combo.itemText(i) + for i in range(registration_panel._initialization_combo.count()) + ] + + assert "moving → fixed (bspline)" in transform_items + assert "moving → fixed (bspline)" not in initialization_items + + def test_apply_transform_uses_bspline_payload( + self, viewer, registration_panel, monkeypatch + ): + moving = xr.DataArray( + np.arange(12, dtype=np.float32).reshape(3, 4), + dims=["y", "x"], + coords={ + "y": xr.DataArray(np.arange(3) * 0.2, dims=["y"]), + "x": xr.DataArray(np.arange(4) * 0.1, dims=["x"]), + }, + ) + transform = _make_bspline_transform().astype(float) + payload = make_bspline_transform_payload( + transform, + reference=moving, + source_layer_name="moving", + target_layer_name="fixed", + operation="register_volume", + transform_model="bspline", + metric="correlation", + diagnostics=_FakeDiagnostics(), + ) + viewer.add_image(moving.values, name="moving", metadata={"xarray": moving}) + viewer.add_image( + moving.values, + name="Registered (bspline)", + metadata={"xarray": moving, "confusius_transform": payload}, + ) + refresh_transform_controls(registration_panel) + registration_panel._transform_source_combo.setCurrentText( + "moving → fixed (bspline)" + ) + registration_panel._transform_target_combo.setCurrentText("moving") + _install_immediate_thread_worker(monkeypatch) + + output_grid = get_output_grid_from_payload(payload) + expected = xr.DataArray( + np.full(output_grid["shape"], 7.0, dtype=np.float32), + dims=output_grid["dims"], + coords={ + dim: xr.DataArray( + output_grid["origin"][i] + + np.arange(output_grid["shape"][i]) * output_grid["spacing"][i], + dims=[dim], + ) + for i, dim in enumerate(output_grid["dims"]) + }, + ) + + def _fake_resample_volume(*args: Any, **kwargs: Any) -> xr.DataArray: + return expected + + monkeypatch.setattr( + "confusius._napari._registration._panel_transforms.resample_volume", + _fake_resample_volume, + ) + + apply_selected_transform(registration_panel) + + layer = viewer.layers["moving → fixed"] + result = layer.metadata["xarray"] + np.testing.assert_array_equal(result.values, expected.values) + np.testing.assert_allclose(result.coords["y"], expected.coords["y"]) + np.testing.assert_allclose(result.coords["x"], expected.coords["x"]) + assert layer.metadata["transform_source"] == "moving → fixed (bspline)" + assert layer.metadata["registration_operation"] == "apply_transform" + + def test_apply_inverse_transform_uses_inverse_affine_and_input_grid( + self, viewer, registration_panel, monkeypatch + ): + source = xr.DataArray( + np.arange(12, dtype=np.float32).reshape(3, 4), + dims=["y", "x"], + coords={ + "y": xr.DataArray(np.arange(3) * 0.2, dims=["y"]), + "x": xr.DataArray(np.arange(4) * 0.1, dims=["x"]), + }, + ) + target = xr.DataArray( + np.arange(30, dtype=np.float32).reshape(5, 6), + dims=["y", "x"], + coords={ + "y": xr.DataArray(np.arange(5) * 0.3, dims=["y"]), + "x": xr.DataArray(np.arange(6) * 0.15, dims=["x"]), + }, + ) + affine = np.array( + [[1.0, 0.0, 0.5], [0.0, 1.0, -0.25], [0.0, 0.0, 1.0]], + dtype=float, + ) + payload = make_affine_transform_payload( + affine, + reference=target, + source=source, + source_layer_name="source", + target_layer_name="target", + operation="register_volume", + transform_model="affine", + metric="correlation", + diagnostics=_FakeDiagnostics(), + ) + viewer.add_image(source.values, name="source", metadata={"xarray": source}) + viewer.add_image(target.values, name="target", metadata={"xarray": target}) + viewer.add_image( + target.values, + name="Registered", + metadata={"xarray": target, "confusius_transform": payload}, + ) + refresh_transform_controls(registration_panel) + registration_panel._transform_source_combo.setCurrentText( + "source → target (affine)" + ) + registration_panel._transform_target_combo.setCurrentText("target") + _install_immediate_thread_worker(monkeypatch) + + apply_selected_inverse_transform(registration_panel) + + input_grid = get_input_grid_from_payload(payload) + expected = resample_volume( + target, + np.linalg.inv(affine), + shape=input_grid["shape"], + spacing=input_grid["spacing"], + origin=input_grid["origin"], + dims=input_grid["dims"], + interpolation="linear", + ) + layer = viewer.layers["target → source"] + result = layer.metadata["xarray"] + np.testing.assert_allclose(result.values, expected.values) + assert tuple(result.dims) == tuple(source.dims) + np.testing.assert_allclose(result.coords["y"], source.coords["y"]) + np.testing.assert_allclose(result.coords["x"], source.coords["x"]) + assert layer.metadata["transform_source"] == "source → target (affine)" + assert layer.metadata["registration_operation"] == "apply_inverse_transform" + + +class TestVolumewiseProgress: + def test_setup_updates_progress_bar_and_output_layer( + self, viewer, registration_panel + ): + moving = xr.DataArray( + np.linspace(-2.0, 3.0, 3 * 4 * 6, dtype=np.float32).reshape(3, 4, 6), + dims=["time", "y", "x"], + coords={ + "time": xr.DataArray(np.arange(3), dims=["time"]), + "y": xr.DataArray(np.arange(4) * 0.2, dims=["y"]), + "x": xr.DataArray(np.arange(6) * 0.1, dims=["x"]), + }, + ) + moving_layer = viewer.add_image( + moving.values, + name="series", + metadata={"xarray": moving}, + ) + progress = setup_volumewise_progress( + registration_panel, + moving_layer=moving_layer, + moving=moving, + layer_name="Motion corrected", + scale_mode="off", + ) + + assert registration_panel._volumewise_progress_layer is not None + assert registration_panel._volumewise_moving_preview_layer is not None + assert registration_panel._progress.maximum() == 3 + assert registration_panel._progress.value() == 0 + np.testing.assert_array_equal( + np.asarray(viewer.layers["Motion corrected"].data), + np.full(moving.shape, float(moving.min()), dtype=np.float32), + ) + + frame = xr.DataArray( + np.ones((4, 6), dtype=np.float32), + dims=["y", "x"], + coords={ + "y": xr.DataArray(np.arange(4) * 0.2, dims=["y"]), + "x": xr.DataArray(np.arange(6) * 0.1, dims=["x"]), + }, + ) + progress.frame_completed(1, frame, _FakeDiagnostics(n_iterations=2)) + + assert registration_panel._progress.value() == 1 + np.testing.assert_array_equal( + np.asarray(viewer.layers["Motion corrected"].data)[1], + np.asarray(frame.values), + ) + + +class TestFinishedCallbacks: + def test_volume_result_adds_new_layer_with_transform_metadata( + self, viewer, registration_panel + ): + fixed = xr.DataArray( + np.ones((4, 6), dtype=np.float32), + dims=["y", "x"], + coords={ + "y": xr.DataArray(np.arange(4) * 0.2, dims=["y"]), + "x": xr.DataArray(np.arange(6) * 0.1, dims=["x"]), + }, + ) + registered = fixed.copy() + transform = np.eye(3) + diagnostics = _FakeDiagnostics() + + payload = { + "operation": "register_volume", + "moving_layer_name": "moving", + "fixed_layer_name": "fixed", + "transform": "rigid", + "metric": "correlation", + "learning_rate": "auto", + "number_of_iterations": 100, + "use_multi_resolution": False, + "resample_interpolation": "linear", + } + + on_registration_finished( + registration_panel, + payload, + (registered, transform, diagnostics), + ) + + layer = viewer.layers["Registered (rigid)"] + assert layer.metadata["registration_transform"] is transform + assert layer.metadata["registration_diagnostics"] is diagnostics + assert layer.metadata["registration_status"] == "completed" + np.testing.assert_array_equal( + get_affine_transform_from_payload(layer.metadata["confusius_transform"]), + transform, + ) + assert ( + layer.metadata["xarray"].attrs["registration_operation"] + == "register_volume" + ) + + def test_volume_result_adds_bspline_transform_metadata( + self, viewer, registration_panel + ): + fixed = xr.DataArray( + np.ones((3, 4), dtype=np.float32), + dims=["y", "x"], + coords={ + "y": xr.DataArray(np.arange(3) * 0.2, dims=["y"]), + "x": xr.DataArray(np.arange(4) * 0.1, dims=["x"]), + }, + ) + registered = fixed.copy() + transform = _make_bspline_transform() + diagnostics = _FakeDiagnostics() + + payload = { + "operation": "register_volume", + "moving_layer_name": "moving", + "fixed_layer_name": "fixed", + "transform": "bspline", + "metric": "correlation", + "learning_rate": "auto", + "number_of_iterations": 100, + "use_multi_resolution": False, + "resample_interpolation": "linear", + } + + on_registration_finished( + registration_panel, + payload, + (registered, transform, diagnostics), + ) + + layer = viewer.layers["Registered (bspline)"] + assert layer.metadata["registration_status"] == "completed" + assert layer.metadata["confusius_transform"]["kind"] == "bspline" + xr.testing.assert_identical( + get_bspline_transform_from_payload(layer.metadata["confusius_transform"]), + transform.astype(float), + ) + + def test_volume_result_replaces_preview_layer( + self, real_viewer, real_registration_panel + ): + moving_data = xr.DataArray( + np.zeros((4, 6), dtype=np.float32), + dims=["y", "x"], + coords={ + "y": xr.DataArray(np.arange(4) * 0.2, dims=["y"]), + "x": xr.DataArray(np.arange(6) * 0.1, dims=["x"]), + }, + ) + moving = real_viewer.add_image( + np.zeros((4, 6), dtype=np.float32), name="moving" + ) + fixed = xr.DataArray( + np.ones((4, 6), dtype=np.float32), + dims=["y", "x"], + coords={ + "y": xr.DataArray(np.arange(4) * 0.2, dims=["y"]), + "x": xr.DataArray(np.arange(6) * 0.1, dims=["x"]), + }, + ) + fixed_layer = real_viewer.add_image( + np.ones((4, 6), dtype=np.float32), name="fixed" + ) + + create_volume_progress_plotter( + real_registration_panel, + moving_layer=moving, + fixed_layer=fixed_layer, + moving=moving_data, + fixed=fixed, + layer_name="Registered (rigid)", + scale_mode="off", + ) + assert {"Fixed", "Moving", "Registered (rigid)"}.issubset( + {layer.name for layer in real_viewer.layers} + ) + + registered = fixed.copy() + transform = np.eye(3) + diagnostics = _FakeDiagnostics() + payload = { + "operation": "register_volume", + "moving_layer_name": "moving", + "fixed_layer_name": "fixed", + "transform": "rigid", + "metric": "correlation", + "learning_rate": "auto", + "number_of_iterations": 100, + "use_multi_resolution": False, + "resample_interpolation": "linear", + } + on_registration_finished( + real_registration_panel, + payload, + (registered, transform, diagnostics), + ) + + assert real_registration_panel._progress_layer is None + assert real_registration_panel._progress_bridge is None + result_layer = real_viewer.layers["Registered (rigid)"] + np.testing.assert_array_equal( + np.asarray(result_layer.data), + np.asarray(registered.values), + ) + assert {"Fixed", "Moving", "Registered (rigid)"}.issubset( + {layer.name for layer in real_viewer.layers} + ) + + def test_create_volume_progress_plotter_applies_initial_transform_to_preview_layers( + self, real_viewer, real_registration_panel + ): + moving_data = xr.DataArray( + np.arange(24, dtype=np.float32).reshape(4, 6), + dims=["y", "x"], + coords={ + "y": xr.DataArray(np.arange(4) * 0.2, dims=["y"]), + "x": xr.DataArray(np.arange(6) * 0.1, dims=["x"]), + }, + ) + moving = real_viewer.add_image(moving_data.values, name="moving") + fixed = xr.DataArray( + np.zeros((4, 6), dtype=np.float32), + dims=["y", "x"], + coords=moving_data.coords, + ) + fixed_layer = real_viewer.add_image( + np.zeros((4, 6), dtype=np.float32), name="fixed" + ) + initial_transform = np.array( + [[1.0, 0.0, 0.2], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]], + dtype=float, + ) + + create_volume_progress_plotter( + real_registration_panel, + moving_layer=moving, + fixed_layer=fixed_layer, + moving=moving_data, + fixed=fixed, + layer_name="Registered (rigid)", + initial_transform=initial_transform, + scale_mode="off", + ) + + expected = resample_like(moving_data, fixed, initial_transform) + np.testing.assert_allclose( + np.asarray(real_viewer.layers["Moving"].data), + np.asarray(expected.data), + ) + np.testing.assert_allclose( + np.asarray(real_viewer.layers["Registered (rigid)"].data), + np.asarray(expected.data), + ) + + def test_progress_layer_data_updates_on_iteration( + self, real_viewer, real_registration_panel + ): + moving_data = xr.DataArray( + np.zeros((4, 6), dtype=np.float32), + dims=["y", "x"], + coords={ + "y": xr.DataArray(np.arange(4) * 0.2, dims=["y"]), + "x": xr.DataArray(np.arange(6) * 0.1, dims=["x"]), + }, + ) + moving = real_viewer.add_image( + np.zeros((4, 6), dtype=np.float32), name="moving" + ) + fixed = xr.DataArray( + np.zeros((4, 6), dtype=np.float32), + dims=["y", "x"], + coords={ + "y": xr.DataArray(np.arange(4) * 0.2, dims=["y"]), + "x": xr.DataArray(np.arange(6) * 0.1, dims=["x"]), + }, + ) + fixed_layer = real_viewer.add_image( + np.zeros((4, 6), dtype=np.float32), name="fixed" + ) + + create_volume_progress_plotter( + real_registration_panel, + moving_layer=moving, + fixed_layer=fixed_layer, + moving=moving_data, + fixed=fixed, + layer_name="Registered (rigid)", + scale_mode="off", + ) + + next_arr = np.full((4, 6), 0.5, dtype=np.float32) + update_progress_layer(real_registration_panel, next_arr) + np.testing.assert_array_equal( + np.asarray(real_viewer.layers["Registered (rigid)"].data), next_arr + ) + + update_progress_layer( + real_registration_panel, np.zeros((3, 6), dtype=np.float32) + ) + np.testing.assert_array_equal( + np.asarray(real_viewer.layers["Registered (rigid)"].data), next_arr + ) + + teardown_volume_progress(real_registration_panel) + assert real_registration_panel._progress_layer is None + assert "Registered (rigid)" not in {layer.name for layer in real_viewer.layers} + + def test_volumewise_result_adds_registered_layer(self, viewer, registration_panel): + moving = xr.DataArray( + np.zeros((3, 4, 6), dtype=np.float32), + dims=["time", "y", "x"], + coords={ + "time": xr.DataArray(np.arange(3), dims=["time"]), + "y": xr.DataArray(np.arange(4) * 0.2, dims=["y"]), + "x": xr.DataArray(np.arange(6) * 0.1, dims=["x"]), + }, + ) + moving_layer = viewer.add_image( + moving.values, + name="series", + metadata={"xarray": moving}, + ) + setup_volumewise_progress( + registration_panel, + moving_layer=moving_layer, + moving=moving, + layer_name="Motion corrected", + scale_mode="off", + ) + + registered = xr.DataArray( + np.ones((3, 4, 6), dtype=np.float32), + dims=["time", "y", "x"], + coords=moving.coords, + attrs={"motion_params": object()}, + ) + payload = { + "operation": "register_volumewise", + "moving_layer_name": "series", + "transform": "rigid", + "metric": "correlation", + "learning_rate": "auto", + "number_of_iterations": 100, + "use_multi_resolution": False, + "resample_interpolation": "linear", + "reference_time": 1, + } + + on_registration_finished(registration_panel, payload, registered) + + layer = viewer.layers["Motion corrected"] + np.testing.assert_array_equal(np.asarray(layer.data), registered.values) + assert layer.metadata["reference_time"] == 1 + assert layer.metadata["registration_operation"] == "register_volumewise" + assert {"Moving", "Motion corrected"}.issubset( + {existing.name for existing in viewer.layers} + ) + assert "registration_status" not in layer.metadata + assert ( + layer.metadata["xarray"].attrs["registration_operation"] + == "register_volumewise" + ) + + def test_unique_transform_and_result_names(self, viewer, registration_panel): + fixed = xr.DataArray( + np.ones((4, 6), dtype=np.float32), + dims=["y", "x"], + coords={ + "y": xr.DataArray(np.arange(4) * 0.2, dims=["y"]), + "x": xr.DataArray(np.arange(6) * 0.1, dims=["x"]), + }, + ) + payload = { + "operation": "register_volume", + "moving_layer_name": "moving", + "fixed_layer_name": "fixed", + "transform": "rigid", + "metric": "correlation", + "learning_rate": "auto", + "number_of_iterations": 100, + "use_multi_resolution": False, + "resample_interpolation": "linear", + } + transform = np.eye(3) + diagnostics = _FakeDiagnostics() + + on_registration_finished( + registration_panel, + payload, + (fixed.copy(), transform, diagnostics), + ) + on_registration_finished( + registration_panel, + payload, + (fixed.copy(), transform, diagnostics), + ) + + assert "Registered (rigid)" in {layer.name for layer in viewer.layers} + assert "Registered (rigid) [1]" in {layer.name for layer in viewer.layers} + names = [ + viewer.layers[name].metadata["confusius_transform"]["name"] + for name in ("Registered (rigid)", "Registered (rigid) [1]") + ] + assert names == [ + "moving → fixed (rigid)", + "moving → fixed (rigid) [1]", + ] diff --git a/tests/unit/test_napari/test_registration_progress.py b/tests/unit/test_napari/test_registration_progress.py new file mode 100644 index 00000000..bc0b86cb --- /dev/null +++ b/tests/unit/test_napari/test_registration_progress.py @@ -0,0 +1,282 @@ +"""Unit tests for the napari-backed registration progress reporter.""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import pytest +import SimpleITK as sitk + +from confusius._napari._registration._progress import ( + NapariRegistrationProgressPlotter, + NapariRegistrationProgressPlotterBridge, + NapariRegistrationProgressReporter, + NapariRegistrationProgressReporterBridge, + make_napari_progress_factory, +) + + +@pytest.fixture +def fixed_img_2d(): + """Small 2D SimpleITK image with a bright square.""" + arr = np.zeros((16, 16), dtype=np.float32) + arr[6:10, 6:10] = 1.0 + img = sitk.GetImageFromArray(arr.T) + img.SetSpacing((1.0, 1.0)) + return img + + +@pytest.fixture +def moving_img_2d(fixed_img_2d): + """Same image shifted by one pixel.""" + arr = sitk.GetArrayFromImage(fixed_img_2d).T + shifted = np.roll(arr, 1, axis=0).astype(np.float32) + img = sitk.GetImageFromArray(shifted.T) + img.SetSpacing(fixed_img_2d.GetSpacing()) + return img + + +def _make_registration_method(ndim: int = 2) -> sitk.ImageRegistrationMethod: + """Return a minimally configured ImageRegistrationMethod.""" + reg = sitk.ImageRegistrationMethod() + reg.SetMetricAsCorrelation() + reg.SetInterpolator(sitk.sitkLinear) + reg.SetOptimizerAsGradientDescent( + learningRate=0.5, + numberOfIterations=5, + convergenceMinimumValue=1e-7, + convergenceWindowSize=3, + ) + reg.SetShrinkFactorsPerLevel(shrinkFactors=[1]) + reg.SetSmoothingSigmasPerLevel(smoothingSigmas=[0]) + reg.SmoothingSigmasAreSpecifiedInPhysicalUnitsOff() + reg.SetInitialTransform(sitk.TranslationTransform(ndim), inPlace=True) + return reg + + +class _SignalSpy: + """Collect emitted payloads from a Qt signal.""" + + def __init__(self) -> None: + self.payloads: list[Any] = [] + + def __call__(self, payload: Any) -> None: + self.payloads.append(payload) + + +class TestNapariRegistrationProgressPlotter: + """Per-iteration reporter behaviour.""" + + def test_update_resamples_and_emits_array(self, qtbot, fixed_img_2d, moving_img_2d): + reg = _make_registration_method(ndim=2) + bridge = NapariRegistrationProgressPlotterBridge() + spy = _SignalSpy() + bridge.iterated.connect(spy) + + reporter = NapariRegistrationProgressPlotter( + bridge, + reg, + fixed_img_2d, + moving_img_2d, + resample_kwargs={"interpolation": "linear", "fill_value": 0.0}, + ) + + with qtbot.waitSignal(bridge.iterated, timeout=2000): + reporter.update() + + assert len(spy.payloads) == 1 + arr = spy.payloads[0] + # `.T` restores numpy axis order, matching `register_volume`. + assert arr.shape == (16, 16) + assert arr.dtype == np.float32 + + def test_update_emits_metric_value(self, qtbot, fixed_img_2d, moving_img_2d): + """`update()` also forwards the current optimizer metric value.""" + reg = _make_registration_method(ndim=2) + bridge = NapariRegistrationProgressPlotterBridge() + metric_spy = _SignalSpy() + bridge.metric_updated.connect(metric_spy) + + reporter = NapariRegistrationProgressPlotter( + bridge, + reg, + fixed_img_2d, + moving_img_2d, + resample_kwargs={"fill_value": 0.0}, + ) + + expected_metric = float(reg.GetMetricValue()) + + with qtbot.waitSignal(bridge.metric_updated, timeout=2000): + reporter.update() + + assert len(metric_spy.payloads) == 1 + emitted_metric = metric_spy.payloads[0] + if np.isnan(expected_metric): + assert np.isnan(emitted_metric) + else: + assert emitted_metric == pytest.approx(expected_metric) + + def test_update_skips_metric_when_plot_metric_false( + self, qtbot, fixed_img_2d, moving_img_2d + ): + """`plot_metric=False` suppresses the metric_updated emission.""" + reg = _make_registration_method(ndim=2) + bridge = NapariRegistrationProgressPlotterBridge() + metric_spy = _SignalSpy() + bridge.metric_updated.connect(metric_spy) + + reporter = NapariRegistrationProgressPlotter( + bridge, + reg, + fixed_img_2d, + moving_img_2d, + plot_metric=False, + resample_kwargs={"fill_value": 0.0}, + ) + # Iterate and confirm the metric signal never fires. We trigger the + # iterated signal first to give the metric a chance to emit, then + # check the spy. + with qtbot.waitSignal(bridge.iterated, timeout=2000): + reporter.update() + assert metric_spy.payloads == [] + + def test_close_emits_finished_signal(self, qtbot, fixed_img_2d, moving_img_2d): + reg = _make_registration_method(ndim=2) + bridge = NapariRegistrationProgressPlotterBridge() + reporter = NapariRegistrationProgressPlotter( + bridge, + reg, + fixed_img_2d, + moving_img_2d, + resample_kwargs={"fill_value": 0.0}, + ) + with qtbot.waitSignal(bridge.finished, timeout=1000): + reporter.close() + + +class TestNapariRegistrationProgressReporter: + """Aggregate per-frame progress for volumewise registration.""" + + def test_frame_completed_emits_progress_and_array(self, qtbot): + import xarray as xr + + bridge = NapariRegistrationProgressReporterBridge() + reporter = NapariRegistrationProgressReporter(bridge, n_frames=3) + progress_payloads: list[tuple[int, int]] = [] + frame_payloads: list[tuple[int, np.ndarray]] = [] + bridge.frame_progress.connect( + lambda completed, total: progress_payloads.append((completed, total)) + ) + bridge.frame_completed.connect( + lambda index, array: frame_payloads.append((index, array)) + ) + frame = xr.DataArray(np.ones((2, 2), dtype=np.float32), dims=("y", "x")) + diagnostics = object() + + with qtbot.waitSignals( + [bridge.frame_progress, bridge.frame_completed], timeout=1000 + ): + reporter.frame_completed(1, frame, diagnostics) # type: ignore[arg-type] + + assert progress_payloads == [(1, 3)] + assert len(frame_payloads) == 1 + assert frame_payloads[0][0] == 1 + np.testing.assert_array_equal(frame_payloads[0][1], frame.values) + + def test_frame_completed_accumulates_unique_progress(self, qtbot): + import xarray as xr + + bridge = NapariRegistrationProgressReporterBridge() + reporter = NapariRegistrationProgressReporter(bridge, n_frames=3) + progress_payloads: list[tuple[int, int]] = [] + bridge.frame_progress.connect( + lambda completed, total: progress_payloads.append((completed, total)) + ) + frame = xr.DataArray(np.ones((2, 2), dtype=np.float32), dims=("y", "x")) + diagnostics = object() + + reporter.frame_completed(1, frame, diagnostics) # type: ignore[arg-type] + reporter.frame_completed(2, frame, diagnostics) # type: ignore[arg-type] + + qtbot.waitUntil(lambda: len(progress_payloads) == 2, timeout=1000) + assert progress_payloads == [(1, 3), (2, 3)] + + def test_close_emits_finished_signal(self, qtbot): + bridge = NapariRegistrationProgressReporterBridge() + reporter = NapariRegistrationProgressReporter(bridge, n_frames=3) + + with qtbot.waitSignal(bridge.finished, timeout=1000): + reporter.close() + + +class TestMakeNapariProgressFactory: + """Factory closure behaviour.""" + + def test_factory_returns_napari_volume_progress( + self, qtbot, fixed_img_2d, moving_img_2d + ): + bridge = NapariRegistrationProgressPlotterBridge() + factory = make_napari_progress_factory(bridge) + reg = _make_registration_method(ndim=2) + + plotter = factory( + reg, + fixed_img_2d, + moving_img_2d, + plot_metric=True, + plot_composite=True, + resample_kwargs={"fill_value": 0.0}, + ) + + assert isinstance(plotter, NapariRegistrationProgressPlotter) + + with qtbot.waitSignals( + [bridge.metric_updated, bridge.iterated, bridge.finished], timeout=2000 + ): + plotter.update() + plotter.close() + + +class TestRegisterVolumeWithNapariFactory: + """End-to-end: register_volume calls the injected napari factory.""" + + def test_factory_is_invoked_and_iterated_signal_fires(self, qtbot): + import xarray as xr + + from confusius.registration.volume import register_volume + + arr = np.zeros((16, 16), dtype=np.float32) + arr[6:10, 6:10] = 1.0 + da = xr.DataArray( + arr, + dims=("y", "x"), + coords={ + "y": np.arange(16) * 0.1, + "x": np.arange(16) * 0.1, + }, + ) + + bridge = NapariRegistrationProgressPlotterBridge() + spy = _SignalSpy() + bridge.iterated.connect(spy) + factory = make_napari_progress_factory(bridge) + + with qtbot.waitSignal(bridge.finished, timeout=5000): + result, _transform, _diagnostics = register_volume( + da, + da, + transform_type="translation", + show_progress=True, + progress_plotter=factory, + plot_metric=True, + plot_composite=False, + ) + + # The translator iterates at least once, so we should have received + # at least one intermediate resampled array. + assert len(spy.payloads) >= 1 + for payload in spy.payloads: + assert payload.shape == da.shape + assert result.shape == da.shape diff --git a/tests/unit/test_registration/test_progress.py b/tests/unit/test_registration/test_progress.py index d08ca2c3..333240ed 100644 --- a/tests/unit/test_registration/test_progress.py +++ b/tests/unit/test_registration/test_progress.py @@ -1,4 +1,8 @@ -"""Unit tests for RegistrationProgressPlotter.""" +"""Unit tests for MatplotlibRegistrationProgressPlotter.""" + +import builtins +import sys +import types import matplotlib import numpy as np @@ -7,8 +11,8 @@ matplotlib.use("Agg") -from confusius.registration._progress import ( # noqa: E402 - RegistrationProgressPlotter, +from confusius.registration.progress import ( # noqa: E402 + MatplotlibRegistrationProgressPlotter, ) # --------------------------------------------------------------------------- @@ -75,17 +79,42 @@ def _make_registration_method(): # --------------------------------------------------------------------------- -# RegistrationProgressPlotter +# MatplotlibRegistrationProgressPlotter # --------------------------------------------------------------------------- -class TestRegistrationProgressPlotterInstantiation: +class TestMatplotlibRegistrationProgressPlotterInstantiation: """Smoke tests for plotter construction.""" + def test_importerror_from_ipython_detection_falls_back_to_script_mode( + self, fixed_img_2d, moving_img_2d, monkeypatch + ): + """Missing IPython support falls back cleanly to non-notebook mode.""" + reg = _make_registration_method() + original_import = builtins.__import__ + + def _guarded_import(name, *args, **kwargs): + if name == "IPython.core.getipython": + raise ImportError("no ipython") + return original_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", _guarded_import) + + plotter = MatplotlibRegistrationProgressPlotter( + reg, + fixed_img_2d, + moving_img_2d, + plot_metric=True, + plot_composite=False, + ) + + assert plotter._notebook is False + plotter.figure.clf() + def test_metric_only(self, fixed_img_2d, moving_img_2d): """Plotter with only metric panel is created without error.""" reg = _make_registration_method() - plotter = RegistrationProgressPlotter( + plotter = MatplotlibRegistrationProgressPlotter( reg, fixed_img_2d, moving_img_2d, @@ -97,7 +126,7 @@ def test_metric_only(self, fixed_img_2d, moving_img_2d): def test_composite_only(self, fixed_img_2d, moving_img_2d): """Plotter with only composite panel is created without error.""" reg = _make_registration_method() - plotter = RegistrationProgressPlotter( + plotter = MatplotlibRegistrationProgressPlotter( reg, fixed_img_2d, moving_img_2d, @@ -109,7 +138,7 @@ def test_composite_only(self, fixed_img_2d, moving_img_2d): def test_both_panels(self, fixed_img_2d, moving_img_2d): """Plotter with both panels is created without error.""" reg = _make_registration_method() - plotter = RegistrationProgressPlotter( + plotter = MatplotlibRegistrationProgressPlotter( reg, fixed_img_2d, moving_img_2d, @@ -119,15 +148,55 @@ def test_both_panels(self, fixed_img_2d, moving_img_2d): plotter.figure.clf() -class TestRegistrationProgressPlotterUpdate: +class TestMatplotlibRegistrationProgressPlotterUpdate: """Tests for metric_values population and composite rendering.""" + def test_notebook_mode_uses_display_and_closes_figure( + self, fixed_img_2d, moving_img_2d, monkeypatch + ): + """Notebook mode renders via IPython display and closes on finish.""" + import matplotlib.pyplot as plt + + reg = _make_registration_method() + display_calls: list[tuple[object, bool]] = [] + closed_figures: list[object] = [] + + fake_getipython = types.ModuleType("IPython.core.getipython") + + class ZMQInteractiveShell: + pass + + fake_getipython.get_ipython = lambda: ZMQInteractiveShell() + fake_display = types.ModuleType("IPython.display") + fake_display.display = ( + lambda fig, clear=False: display_calls.append((fig, clear)) + ) + monkeypatch.setitem(sys.modules, "IPython.core.getipython", fake_getipython) + monkeypatch.setitem(sys.modules, "IPython.display", fake_display) + monkeypatch.setattr(plt, "close", lambda fig: closed_figures.append(fig)) + + plotter = MatplotlibRegistrationProgressPlotter( + reg, + fixed_img_2d, + moving_img_2d, + plot_metric=False, + plot_composite=True, + ) + + plotter.update() + plotter.close() + + assert display_calls + assert display_calls[-1][0] is plotter.figure + assert display_calls[-1][1] is True + assert closed_figures == [plotter.figure] + def test_metric_values_populated_after_registration( self, fixed_img_2d, moving_img_2d ): """metric_values contains one entry per iteration after registration.""" reg = _make_registration_method() - plotter = RegistrationProgressPlotter( + plotter = MatplotlibRegistrationProgressPlotter( reg, fixed_img_2d, moving_img_2d, @@ -146,7 +215,7 @@ def test_metric_values_populated_after_registration( def test_metric_values_are_floats(self, fixed_img_2d, moving_img_2d): """All recorded metric values are finite floats.""" reg = _make_registration_method() - plotter = RegistrationProgressPlotter( + plotter = MatplotlibRegistrationProgressPlotter( reg, fixed_img_2d, moving_img_2d, @@ -167,7 +236,7 @@ def test_composite_panel_rendered_after_registration( ): """Composite panel renders without error after at least one iteration.""" reg = _make_registration_method() - plotter = RegistrationProgressPlotter( + plotter = MatplotlibRegistrationProgressPlotter( reg, fixed_img_2d, moving_img_2d, @@ -198,7 +267,7 @@ def test_3d_composite_panel_rendered(self, fixed_img_3d, moving_img_3d): reg.SmoothingSigmasAreSpecifiedInPhysicalUnitsOff() reg.SetInitialTransform(sitk.TranslationTransform(3), inPlace=True) - plotter = RegistrationProgressPlotter( + plotter = MatplotlibRegistrationProgressPlotter( reg, fixed_img_3d, moving_img_3d, @@ -214,37 +283,75 @@ def test_3d_composite_panel_rendered(self, fixed_img_3d, moving_img_3d): plotter.figure.clf() -class TestRegistrationProgressPlotterResampleKwargs: - """Tests for resample_kwargs fill-value behaviour.""" +class TestMatplotlibRegistrationProgressPlotterResampleKwargs: + """Tests for intermediate-resample settings.""" + + def test_none_interpolation_falls_back_to_linear( + self, fixed_img_2d, moving_img_2d + ): + """A `None` interpolation override falls back to linear at render time.""" + reg = _make_registration_method() + plotter = MatplotlibRegistrationProgressPlotter( + reg, + fixed_img_2d, + moving_img_2d, + plot_metric=False, + plot_composite=True, + resample_kwargs={"interpolation": None}, + ) + + plotter.update() + + assert plotter._composite_im is not None + plotter.figure.clf() + + def test_invalid_interpolation_raises_on_update( + self, fixed_img_2d, moving_img_2d + ): + """Unknown interpolation names raise a clear ValueError during rendering.""" + reg = _make_registration_method() + plotter = MatplotlibRegistrationProgressPlotter( + reg, + fixed_img_2d, + moving_img_2d, + plot_metric=False, + plot_composite=True, + resample_kwargs={"interpolation": "bogus"}, + ) + + with pytest.raises(ValueError, match="Invalid `interpolation`"): + plotter.update() + + plotter.figure.clf() def test_default_fill_value_is_moving_min(self, fixed_img_2d, moving_img_2d): - """When resample_kwargs omits default_value, it is set to moving_img.min().""" + """When resample_kwargs omits fill_value, it defaults to moving_img.min().""" reg = _make_registration_method() - plotter = RegistrationProgressPlotter( + plotter = MatplotlibRegistrationProgressPlotter( reg, fixed_img_2d, moving_img_2d, plot_metric=False, plot_composite=True ) expected = float(sitk.GetArrayFromImage(moving_img_2d).min()) - assert plotter._resample_kwargs["default_value"] == pytest.approx(expected) + assert plotter._fill_value == pytest.approx(expected) plotter.figure.clf() def test_explicit_fill_value_is_respected(self, fixed_img_2d, moving_img_2d): - """Explicit default_value in resample_kwargs overrides the auto-default.""" + """Explicit fill_value in resample_kwargs overrides the auto-default.""" reg = _make_registration_method() - plotter = RegistrationProgressPlotter( + plotter = MatplotlibRegistrationProgressPlotter( reg, fixed_img_2d, moving_img_2d, plot_metric=False, plot_composite=True, - resample_kwargs={"default_value": -60.0}, + resample_kwargs={"fill_value": -60.0}, ) - assert plotter._resample_kwargs["default_value"] == pytest.approx(-60.0) + assert plotter._fill_value == pytest.approx(-60.0) plotter.figure.clf() def test_explicit_interpolation_is_stored(self, fixed_img_2d, moving_img_2d): """interpolation key in resample_kwargs is stored and later used.""" reg = _make_registration_method() - plotter = RegistrationProgressPlotter( + plotter = MatplotlibRegistrationProgressPlotter( reg, fixed_img_2d, moving_img_2d, @@ -252,7 +359,7 @@ def test_explicit_interpolation_is_stored(self, fixed_img_2d, moving_img_2d): plot_composite=True, resample_kwargs={"interpolation": "nearest"}, ) - assert plotter._resample_kwargs["interpolation"] == "nearest" + assert plotter._interpolation == "nearest" plotter.figure.clf() diff --git a/tests/unit/test_registration/test_volume.py b/tests/unit/test_registration/test_volume.py index e551cb41..730cae9a 100644 --- a/tests/unit/test_registration/test_volume.py +++ b/tests/unit/test_registration/test_volume.py @@ -1,5 +1,8 @@ """Unit tests for single-volume registration.""" +import signal +from threading import Event + import numpy as np import pytest import xarray as xr @@ -10,6 +13,130 @@ from confusius.registration.volume import register_volume +class TestRegisterVolumeSigint: + """Ctrl+C handling exposed through the public `register_volume` API.""" + + def test_first_ctrl_c_returns_aborted_result_and_restores_handler( + self, sample_2d_dataarray_spatial, monkeypatch + ): + """First Ctrl+C sets the cooperative abort event and restores SIGINT afterwards.""" + import SimpleITK as sitk + + previous_handler = signal.getsignal(signal.SIGINT) + + def fake_execute(self, fixed, moving): + del self, fixed, moving + handler = signal.getsignal(signal.SIGINT) + assert callable(handler) + handler(signal.SIGINT, None) + return sitk.TranslationTransform(2) + + monkeypatch.setattr(sitk.ImageRegistrationMethod, "Execute", fake_execute) + + _result, _transform, diagnostics = register_volume( + sample_2d_dataarray_spatial, + sample_2d_dataarray_spatial, + transform_type="translation", + ) + + assert diagnostics.status == "aborted" + assert signal.getsignal(signal.SIGINT) is previous_handler + + def test_second_ctrl_c_raises_keyboardinterrupt( + self, sample_2d_dataarray_spatial, monkeypatch + ): + """Second Ctrl+C falls back to the previous default SIGINT handler.""" + import SimpleITK as sitk + + previous_handler = signal.getsignal(signal.SIGINT) + + def fake_execute(self, fixed, moving): + del self, fixed, moving + handler = signal.getsignal(signal.SIGINT) + assert callable(handler) + handler(signal.SIGINT, None) + handler(signal.SIGINT, None) + return sitk.TranslationTransform(2) + + monkeypatch.setattr(sitk.ImageRegistrationMethod, "Execute", fake_execute) + + with pytest.raises(KeyboardInterrupt): + register_volume( + sample_2d_dataarray_spatial, + sample_2d_dataarray_spatial, + transform_type="translation", + ) + + assert signal.getsignal(signal.SIGINT) is previous_handler + + def test_second_ctrl_c_ignores_when_previous_handler_ignores( + self, sample_2d_dataarray_spatial, monkeypatch + ): + """Second Ctrl+C is ignored when the previous SIGINT handler ignored it.""" + import SimpleITK as sitk + + previous_handler = signal.getsignal(signal.SIGINT) + signal.signal(signal.SIGINT, signal.SIG_IGN) + + def fake_execute(self, fixed, moving): + del self, fixed, moving + handler = signal.getsignal(signal.SIGINT) + assert callable(handler) + handler(signal.SIGINT, None) + handler(signal.SIGINT, None) + return sitk.TranslationTransform(2) + + monkeypatch.setattr(sitk.ImageRegistrationMethod, "Execute", fake_execute) + + try: + _result, _transform, diagnostics = register_volume( + sample_2d_dataarray_spatial, + sample_2d_dataarray_spatial, + transform_type="translation", + ) + finally: + signal.signal(signal.SIGINT, previous_handler) + + assert diagnostics.status == "aborted" + + def test_second_ctrl_c_calls_previous_custom_handler( + self, sample_2d_dataarray_spatial, monkeypatch + ): + """Second Ctrl+C delegates to a previous custom handler when one is installed.""" + import SimpleITK as sitk + + previous_handler = signal.getsignal(signal.SIGINT) + calls: list[tuple[int, object]] = [] + + def custom_handler(signum: int, frame: object) -> None: + calls.append((signum, frame)) + + signal.signal(signal.SIGINT, custom_handler) + + def fake_execute(self, fixed, moving): + del self, fixed, moving + handler = signal.getsignal(signal.SIGINT) + assert callable(handler) + handler(signal.SIGINT, None) + handler(signal.SIGINT, None) + return sitk.TranslationTransform(2) + + monkeypatch.setattr(sitk.ImageRegistrationMethod, "Execute", fake_execute) + + try: + _result, _transform, diagnostics = register_volume( + sample_2d_dataarray_spatial, + sample_2d_dataarray_spatial, + transform_type="translation", + ) + finally: + signal.signal(signal.SIGINT, previous_handler) + + assert diagnostics.status == "aborted" + assert len(calls) == 1 + assert calls[0][0] == signal.SIGINT + + class TestRegisterVolumeValidation: """Input validation for register_volume.""" @@ -81,6 +208,96 @@ def test_shape_mismatch_no_error( ) assert result.shape == moving.shape + def test_abort_event_returns_partial_result(self, sample_2d_dataarray_spatial): + """A pre-set abort event returns an aborted diagnostics record.""" + abort_event = Event() + abort_event.set() + + result, _transform, diagnostics = register_volume( + sample_2d_dataarray_spatial, + sample_2d_dataarray_spatial, + transform_type="translation", + abort_event=abort_event, + ) + + assert result.shape == sample_2d_dataarray_spatial.shape + assert diagnostics.status == "aborted" + assert diagnostics.n_iterations == 0 + + def test_unknown_runtime_error_is_passed_through( + self, sample_2d_dataarray_spatial, monkeypatch + ): + """Unknown SimpleITK runtime errors are re-raised unchanged.""" + import SimpleITK as sitk + + error = RuntimeError("boom") + + def fake_execute(self, fixed, moving): + del self, fixed, moving + raise error + + monkeypatch.setattr(sitk.ImageRegistrationMethod, "Execute", fake_execute) + + with pytest.raises(RuntimeError) as excinfo: + register_volume( + sample_2d_dataarray_spatial, + sample_2d_dataarray_spatial, + transform_type="translation", + ) + + assert excinfo.value is error + + def test_bspline_scale_error_raises_clearer_message( + self, sample_2d_dataarray_spatial, monkeypatch + ): + """Known SimpleITK scale failures are rewritten to actionable errors.""" + import SimpleITK as sitk + + def fake_execute(self, fixed, moving): + del self, fixed, moving + raise RuntimeError( + "Exception thrown in SimpleITK ImageRegistrationMethod_Execute: " + "ITK ERROR: GradientDescentOptimizerv4Template: " + "m_Scales values must be > epsilon.[1e-20, 1e-12]" + ) + + monkeypatch.setattr(sitk.ImageRegistrationMethod, "Execute", fake_execute) + + with pytest.raises(RuntimeError, match="could not compute valid optimizer scales"): + register_volume( + sample_2d_dataarray_spatial, + sample_2d_dataarray_spatial, + transform_type="bspline", + learning_rate=1.0, + ) + + def test_bspline_scale_error_with_auto_learning_rate_suggests_fixed_rate( + self, sample_2d_dataarray_spatial, monkeypatch + ): + """Auto-learning-rate scale failures suggest retrying with a fixed rate.""" + import SimpleITK as sitk + + def fake_execute(self, fixed, moving): + del self, fixed, moving + raise RuntimeError( + "Exception thrown in SimpleITK ImageRegistrationMethod_Execute: " + "ITK ERROR: GradientDescentOptimizerv4Template: " + "m_Scales values must be > epsilon.[1e-20, 1e-12]" + ) + + monkeypatch.setattr(sitk.ImageRegistrationMethod, "Execute", fake_execute) + + with pytest.raises( + RuntimeError, + match='Retry with a fixed `learning_rate` such as `0.1` or `0.01`', + ): + register_volume( + sample_2d_dataarray_spatial, + sample_2d_dataarray_spatial, + transform_type="bspline", + learning_rate="auto", + ) + class TestRegisterVolumeOutput: """Output properties for register_volume.""" @@ -109,7 +326,7 @@ def test_bspline_returns_dataarray_transform(self, sample_2d_dataarray_spatial): transform_type="bspline", ) assert isinstance(bspline_tx, xr.DataArray) - assert bspline_tx.attrs.get("type") == "bspline_transform" + assert bspline_tx.attrs.get("transform_type") == "bspline_transform" assert bspline_tx.dims[0] == "component" def test_resample_true_coords_match_fixed( @@ -931,3 +1148,126 @@ def test_default_fill_value_is_moving_min(self): assert float(result.values[0, 0]) == pytest.approx( float(moving.min()), abs=1e-5 ) + + +class TestRegisterVolumePreSetAbort: + """Pre-set abort_event short-circuits before SimpleITK Execute is called.""" + + def test_bspline_abort_returns_initial_bspline_transform( + self, sample_2d_dataarray_spatial + ): + """Pre-aborted bspline returns a DataArray without forcing a bspline fit.""" + abort_event = Event() + abort_event.set() + + _, transform, diagnostics = register_volume( + sample_2d_dataarray_spatial, + sample_2d_dataarray_spatial, + transform_type="bspline", + abort_event=abort_event, + ) + + assert diagnostics.status == "aborted" + assert diagnostics.n_iterations == 0 + assert ( + diagnostics.stop_condition + == "Registration aborted before optimisation started." + ) + # The returned DataArray wraps the initial (unoptimised) bspline — its + # coefficients differ from a real registration only in that no iterations ran. + assert isinstance(transform, xr.DataArray) + assert transform.attrs.get("transform_type") == "bspline_transform" + + def test_affine_initialization_abort_returns_initialization_affine( + self, sample_2d_dataarray_spatial + ): + """Pre-aborted linear registration returns the provided affine initialization. + + The transform must match the initialization matrix — not the default + identity/TranslationTransform fallback used when no initialization is set — + so downstream consumers can rely on a coherent aborted transform. + """ + pre_affine = np.array([[1.0, 0.0, 0.5], [0.0, 1.0, -0.25], [0.0, 0.0, 1.0]]) + + abort_event = Event() + abort_event.set() + + _, transform, diagnostics = register_volume( + sample_2d_dataarray_spatial, + sample_2d_dataarray_spatial, + transform_type="rigid", + initialization=pre_affine, + abort_event=abort_event, + ) + + assert diagnostics.status == "aborted" + assert diagnostics.n_iterations == 0 + assert_allclose(transform, pre_affine) + + +class TestRegisterVolumeConvergesBeforeFirstIteration: + """`final_metric_value` falls back to the optimizer's metric when no iteration event fires.""" + + def test_final_metric_value_pulled_from_optimizer_when_no_iterations( + self, sample_2d_dataarray_spatial + ): + """When SimpleITK converges before any iteration event, final_metric_value is + the optimizer's current metric, not NaN. + + Achieved by raising `convergence_minimum_value` above the metric for identical + images and shrinking the window to 1, so the convergence checker passes at + iteration 0 before any iteration event fires. + """ + _, _, diagnostics = register_volume( + sample_2d_dataarray_spatial, + sample_2d_dataarray_spatial, + transform_type="translation", + number_of_iterations=100, + convergence_minimum_value=1.0, + convergence_window_size=1, + ) + + assert diagnostics.n_iterations == 0 + assert diagnostics.status == "completed" + assert np.isfinite(diagnostics.final_metric_value) + assert "Convergence checker passed at iteration 0" in diagnostics.stop_condition + + +class TestRegisterVolumeFromWorkerThread: + """`register_volume` works when called from a non-main thread.""" + + def test_register_volume_runs_in_non_main_thread(self, sample_2d_dataarray_spatial): + """Calling `register_volume` from a worker thread bypasses SIGINT wiring. + + The non-main-thread branch of `abort_on_sigint` skips installing a SIGINT + handler and simply yields the abort event, so registration runs to + completion without trying to mutate the main thread's signal handlers. + """ + import threading + + from confusius.registration.diagnostics import RegistrationDiagnostics + + result_holder: dict[str, object] = {} + + def worker() -> None: + assert threading.current_thread() is not threading.main_thread() + result, transform, diagnostics = register_volume( + sample_2d_dataarray_spatial, + sample_2d_dataarray_spatial, + transform_type="translation", + number_of_iterations=2, + ) + result_holder["result"] = result + result_holder["transform"] = transform + result_holder["diagnostics"] = diagnostics + + thread = threading.Thread(target=worker) + thread.start() + thread.join() + + result = result_holder["result"] + diagnostics = result_holder["diagnostics"] + assert isinstance(result, xr.DataArray) + assert isinstance(diagnostics, RegistrationDiagnostics) + assert result.shape == sample_2d_dataarray_spatial.shape + assert diagnostics.status == "completed" diff --git a/tests/unit/test_registration/test_volumewise.py b/tests/unit/test_registration/test_volumewise.py index 32ddae37..665eb816 100644 --- a/tests/unit/test_registration/test_volumewise.py +++ b/tests/unit/test_registration/test_volumewise.py @@ -1,5 +1,7 @@ """Unit tests for volumewise registration functions.""" +from threading import Event + import numpy as np import pytest import xarray as xr @@ -9,6 +11,23 @@ from confusius.registration.volumewise import register_volumewise +class _FakeVolumewiseProgressReporter: + def __init__(self) -> None: + self.completed_frames: list[int] = [] + self.closed = False + + def frame_completed( + self, + frame_index: int, + registered_frame: xr.DataArray, + diagnostics: RegistrationDiagnostics, + ) -> None: + self.completed_frames.append(frame_index) + + def close(self) -> None: + self.closed = True + + class TestRegisterVolumewise: """Tests for register_volumewise function.""" @@ -69,6 +88,125 @@ def _guarded_import(name, *args, **kwargs): assert result.shape == sample_2d_dataarray.shape + def test_abort_event_returns_partial_dataset(self, sample_2d_dataarray): + """A pre-set abort event returns an aborted partial dataset.""" + abort_event = Event() + abort_event.set() + + result = register_volumewise( + sample_2d_dataarray, + n_jobs=2, + transform="translation", + abort_event=abort_event, + ) + + assert result.shape == sample_2d_dataarray.shape + assert set(result.attrs["motion_params"]["status"]) == {"aborted"} + assert_allclose( + result.values, + np.full_like(sample_2d_dataarray.values, sample_2d_dataarray.values.min()), + ) + + def test_progress_reporter_receives_frame_updates( + self, sample_2d_dataarray, monkeypatch + ): + reporter = _FakeVolumewiseProgressReporter() + + def _fake_register_volume(_volume, _ref_da, **kwargs): + diagnostics = RegistrationDiagnostics( + metric="correlation", + metric_values=np.asarray([-1.0, -0.5]), + final_metric_value=-0.5, + n_iterations=2, + stop_condition="done", + status="completed", + ) + return _volume.copy(), np.eye(3), diagnostics + + monkeypatch.setattr( + "confusius.registration.volumewise.register_volume", + _fake_register_volume, + ) + + result = register_volumewise( + sample_2d_dataarray, + n_jobs=1, + transform="translation", + show_progress=False, + progress_reporter=reporter, + ) + + assert result.shape == sample_2d_dataarray.shape + assert sorted(reporter.completed_frames) == list( + range(sample_2d_dataarray.sizes["time"]) + ) + assert reporter.closed + + def test_abort_during_run_skips_not_yet_started_frames( + self, sample_2d_dataarray, monkeypatch + ): + """Already-scheduled frames hit the cheap aborted-frame fast path.""" + import joblib + + abort_event = Event() + calls = {"count": 0} + + def _fake_register_volume(volume, _ref_da, **kwargs): + calls["count"] += 1 + if calls["count"] == 1: + abort_event.set() + diagnostics = RegistrationDiagnostics( + metric="correlation", + metric_values=np.asarray([-1.0]), + final_metric_value=-1.0, + n_iterations=1, + stop_condition="done", + status="completed", + ) + return volume.copy(), np.eye(3), diagnostics + + class _FakeParallel: + def __init__(self, *args, **kwargs): + del args, kwargs + + def __call__(self, tasks): + scheduled = list(tasks) + + def _run(): + for task in scheduled: + yield task() + + return _run() + + def _fake_delayed(func): + def _wrap(*args, **kwargs): + return lambda: func(*args, **kwargs) + + return _wrap + + monkeypatch.setattr( + "confusius.registration.volumewise.register_volume", + _fake_register_volume, + ) + monkeypatch.setattr(joblib, "Parallel", _FakeParallel) + monkeypatch.setattr(joblib, "delayed", _fake_delayed) + + result = register_volumewise( + sample_2d_dataarray, + n_jobs=2, + transform="translation", + show_progress=False, + abort_event=abort_event, + ) + + statuses = list(result.attrs["motion_params"]["status"]) + assert statuses[0] == "completed" + assert all(status == "aborted" for status in statuses[1:]) + assert calls["count"] == 1 + + background = sample_2d_dataarray.values.min() + assert np.all(result.values[1:] == background) + def test_wrong_dimensionality_raises(self): """Data that is neither 2D+t nor 3D+t raises ValueError.""" # 1D+time = 2D total. diff --git a/tests/unit/test_xarray/test_wrapper_calls.py b/tests/unit/test_xarray/test_wrapper_calls.py index 87c1dcf0..cf04b099 100644 --- a/tests/unit/test_xarray/test_wrapper_calls.py +++ b/tests/unit/test_xarray/test_wrapper_calls.py @@ -226,9 +226,15 @@ def _volumewise(data, **kwargs): ) fixed = sample_3d_volume.copy() + fixed_mask = fixed > 0 + moving_mask = sample_3d_volume > 0 + progress_plotter = object() + abort_event = object() assert ( sample_3d_volume.fusi.register.to_volume( fixed, + fixed_mask=fixed_mask, + moving_mask=moving_mask, transform="affine", metric="mattes_mi", number_of_histogram_bins=40, @@ -248,6 +254,9 @@ def _volumewise(data, **kwargs): plot_metric=False, plot_composite=False, fill_value=-1.0, + sitk_threads=2, + progress_plotter=progress_plotter, + abort_event=abort_event, ) is reg_result ) @@ -255,6 +264,8 @@ def _volumewise(data, **kwargs): sample_3d_volume, fixed, { + "fixed_mask": fixed_mask, + "moving_mask": moving_mask, "transform_type": "affine", "metric": "mattes_mi", "number_of_histogram_bins": 40, @@ -270,13 +281,18 @@ def _volumewise(data, **kwargs): "smoothing_sigmas": (3, 1, 0), "resample": True, "resample_interpolation": "bspline", + "fill_value": -1.0, + "sitk_threads": 2, "show_progress": True, "plot_metric": False, "plot_composite": False, - "fill_value": -1.0, + "progress_plotter": progress_plotter, + "abort_event": abort_event, }, ) + progress_reporter = object() + abort_event = object() assert ( sample_3d_volume.fusi.register.volumewise( reference_time=2, @@ -295,6 +311,8 @@ def _volumewise(data, **kwargs): smoothing_sigmas=(2, 0), resample_interpolation="bspline", show_progress=False, + progress_reporter=progress_reporter, + abort_event=abort_event, keep_diagnostics=True, ) is volumewise_result @@ -318,6 +336,8 @@ def _volumewise(data, **kwargs): "smoothing_sigmas": (2, 0), "resample_interpolation": "bspline", "show_progress": False, + "progress_reporter": progress_reporter, + "abort_event": abort_event, "keep_diagnostics": True, }, ) diff --git a/tools/prefetch_doc_datasets.py b/tools/prefetch_doc_datasets.py index f4fcc4eb..327166f5 100644 --- a/tools/prefetch_doc_datasets.py +++ b/tools/prefetch_doc_datasets.py @@ -47,7 +47,7 @@ def _prefetch_nunez_elizalde() -> None: # docs/images/gui/generate.py fetch_nunez_elizalde_2022( subjects="CR022", - sessions="20201011", + sessions=["20201007", "20201011"], tasks="spontaneous", acqs="slice04", ) @@ -99,6 +99,14 @@ def _prefetch_cybis_pereira() -> None: acqs="slice37", ) + # docs/images/gui/generate.py (within-scan registration GIF) + fetch_cybis_pereira_2026( + datasets="rawdata", + subjects="rat75", + sessions="20220523", + acqs="slice32", + ) + # docs/examples/registration/register_volume_same_subject.py fetch_cybis_pereira_2026( datasets="rawdata",