Thomas Ales
756 lines
34 KiB
Python
756 lines
34 KiB
Python
"""FFT Options and Manual Alignment dialogs."""
|
||
|
||
from typing import TYPE_CHECKING
|
||
|
||
import matplotlib as mpl
|
||
import numpy as np
|
||
from matplotlib.backends.backend_qtagg import NavigationToolbar2QT
|
||
from PyQt6.QtCore import QSignalBlocker, pyqtSignal
|
||
from PyQt6.QtWidgets import (
|
||
QButtonGroup, QComboBox, QDialog, QDialogButtonBox, QGroupBox,
|
||
QHBoxLayout, QLabel, QMessageBox, QPushButton, QRadioButton, QSpinBox,
|
||
QVBoxLayout, QWidget,
|
||
)
|
||
|
||
import sras_compute as compute
|
||
from sras_compute import (
|
||
PYFFTW_AVAILABLE, ManualAngleParams, build_manual_alignment,
|
||
delete_manual_alignment, save_manual_alignment,
|
||
)
|
||
from sras_format import SrasFile
|
||
from sras_workers import Ch4MaskWorker, CrossCorrelateWorker
|
||
|
||
from .canvases import ManualAlignOverlayCanvas
|
||
from .common import (
|
||
_CSS_HINT, _CSS_MUTED, _CSS_WARN, Jobs, _axes_extent, _form, _group, _make_dspin,
|
||
_scroll_panel, _wrap_label,
|
||
)
|
||
|
||
if TYPE_CHECKING:
|
||
from .main_window import SrasViewerWindow
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# FFT Options dialog
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class FftOptionsDialog(QDialog):
|
||
"""Configure FFT backend and zero-padding.
|
||
|
||
Changes take effect only when the user clicks Apply. Cancel discards
|
||
all pending edits. The live 'frequency resolution' label updates as
|
||
the user adjusts the pad factor so they can see the trade-off before
|
||
committing.
|
||
"""
|
||
|
||
def __init__(self, parent=None, *,
|
||
current_backend: str,
|
||
current_pad_factor: int,
|
||
samples_per_frame: int | None,
|
||
sample_rate_hz: float | None,
|
||
grating_um: float,
|
||
cached_pad_factor: int | None = None):
|
||
super().__init__(parent)
|
||
self.setWindowTitle("FFT Options")
|
||
self.setModal(True)
|
||
self.setMinimumWidth(380)
|
||
|
||
self._samples_per_frame = samples_per_frame
|
||
self._sample_rate_hz = sample_rate_hz
|
||
self._grating_um = grating_um
|
||
# The pad the open file's stored FFT cache was computed at, if it has
|
||
# one — picking anything else here makes that cache unreadable, which
|
||
# is worth saying before Apply rather than after.
|
||
self._cached_pad_factor = cached_pad_factor
|
||
|
||
layout = QVBoxLayout(self)
|
||
|
||
# ---- Backend ---------------------------------------------------
|
||
grp_backend = QGroupBox("FFT Backend")
|
||
bl = QVBoxLayout(grp_backend)
|
||
|
||
self._btn_scipy = QRadioButton("SciPy FFT (pocketfft) (always available)")
|
||
self._btn_pyfftw = QRadioButton(
|
||
"pyFFTW (faster for large arrays)" if PYFFTW_AVAILABLE
|
||
else "pyFFTW (not installed — run: pip install pyfftw)")
|
||
self._btn_pyfftw.setEnabled(PYFFTW_AVAILABLE)
|
||
|
||
self._backend_group = QButtonGroup(self)
|
||
self._backend_group.addButton(self._btn_scipy, id=0)
|
||
self._backend_group.addButton(self._btn_pyfftw, id=1)
|
||
|
||
if current_backend == "pyfftw" and PYFFTW_AVAILABLE:
|
||
self._btn_pyfftw.setChecked(True)
|
||
else:
|
||
self._btn_scipy.setChecked(True)
|
||
|
||
bl.addWidget(self._btn_scipy)
|
||
bl.addWidget(self._btn_pyfftw)
|
||
layout.addWidget(grp_backend)
|
||
|
||
# ---- Zero-padding ----------------------------------------------
|
||
grp_zp = QGroupBox("Zero-Padding")
|
||
zl = QVBoxLayout(grp_zp)
|
||
|
||
pad_row = QHBoxLayout()
|
||
pad_row.addWidget(QLabel("Pad factor:"))
|
||
self._spin_pad = QSpinBox()
|
||
self._spin_pad.setRange(1, 256)
|
||
self._spin_pad.setValue(max(1, current_pad_factor))
|
||
self._spin_pad.setToolTip(
|
||
"Multiply the waveform length by this factor via zero-padding\n"
|
||
"before computing the FFT.\n"
|
||
"1 = no padding (natural length).\n"
|
||
"Powers of 2 (2, 4, 8 …) give the best performance."
|
||
)
|
||
self._spin_pad.valueChanged.connect(self._update_info)
|
||
pad_row.addWidget(self._spin_pad)
|
||
zl.addLayout(pad_row)
|
||
|
||
self._lbl_nfft = QLabel()
|
||
self._lbl_freq_res = QLabel()
|
||
self._lbl_vel_res = QLabel()
|
||
for lbl in (self._lbl_nfft, self._lbl_freq_res, self._lbl_vel_res):
|
||
lbl.setStyleSheet(_CSS_HINT)
|
||
zl.addWidget(lbl)
|
||
|
||
self._lbl_cache = QLabel()
|
||
self._lbl_cache.setWordWrap(True)
|
||
self._lbl_cache.setStyleSheet(_CSS_WARN)
|
||
zl.addWidget(self._lbl_cache)
|
||
|
||
layout.addWidget(grp_zp)
|
||
|
||
# ---- Buttons ---------------------------------------------------
|
||
buttons = QDialogButtonBox()
|
||
buttons.addButton("Apply", QDialogButtonBox.ButtonRole.AcceptRole
|
||
).clicked.connect(self.accept)
|
||
buttons.addButton("Cancel", QDialogButtonBox.ButtonRole.RejectRole
|
||
).clicked.connect(self.reject)
|
||
layout.addWidget(buttons)
|
||
|
||
self._update_info()
|
||
|
||
def _update_info(self):
|
||
spf = self._samples_per_frame
|
||
sr = self._sample_rate_hz
|
||
pad = self._spin_pad.value()
|
||
|
||
self._lbl_cache.setText(
|
||
"" if self._cached_pad_factor in (None, pad) else
|
||
f"! This file's stored FFT cache was computed at pad "
|
||
f"{self._cached_pad_factor}x, so at {pad}x it can't be used and "
|
||
f"CH1/Velocity will recompute. Re-run Convert -> Batch Compute FFT "
|
||
f"to cache at {pad}x.")
|
||
|
||
if spf is None or sr is None:
|
||
self._lbl_nfft.setText("Load a file to preview FFT parameters.")
|
||
self._lbl_freq_res.setText("")
|
||
self._lbl_vel_res.setText("")
|
||
return
|
||
|
||
n_fft = spf * pad
|
||
freq_res_hz = sr / n_fft
|
||
freq_res_mhz = freq_res_hz / 1e6
|
||
# v (m/s) = freq (MHz) × grating (µm)
|
||
vel_res_ms = freq_res_mhz * self._grating_um
|
||
|
||
self._lbl_nfft.setText(f"FFT points: {spf} × {pad} = {n_fft:,}")
|
||
self._lbl_freq_res.setText(
|
||
f"Frequency bin: {freq_res_mhz:.4f} MHz ({freq_res_hz / 1e3:.2f} kHz)")
|
||
self._lbl_vel_res.setText(
|
||
f"Velocity bin: {vel_res_ms:.3f} m/s "
|
||
f"(at grating = {self._grating_um:.2f} µm)")
|
||
|
||
def get_backend(self) -> str:
|
||
return "pyfftw" if self._btn_pyfftw.isChecked() and PYFFTW_AVAILABLE else "scipy"
|
||
|
||
def get_pad_factor(self) -> int:
|
||
return max(1, self._spin_pad.value())
|
||
|
||
|
||
|
||
class ManualAlignmentDialog(QDialog):
|
||
"""Non-modal manual angle-alignment editor (Fusion -> Manual Alignment...).
|
||
|
||
Shows every angle's binarized CH4 (Bias B) mask overlaid in a distinct
|
||
color at partial opacity on one shared canvas, so translation/rotation
|
||
misalignment is visible by eye. Reference angle (always index 0) is
|
||
ground truth and never moves; every other angle is aligned to it. The
|
||
user picks an "active" angle and nudges its rotation+translation with
|
||
the keyboard; Auto Cross-Correlate finds every non-reference angle's
|
||
rotation *and* translation by registering its image against the
|
||
reference's (see compute.register_angle_to_reference) — meant to get every
|
||
angle stacked on top of each other so keyboard nudging only has to make
|
||
small corrections, not find an alignment from scratch; Auto De-rotate is
|
||
the weaker fallback that just seeds rotation from the stage's reported
|
||
angle, leaving translation alone. Save writes a JSON sidecar next to the
|
||
.sras file and hands a freshly-built, full-resolution AlignmentResult back
|
||
to the main window — the exact same object shape compute_angle_alignment
|
||
produces, so every existing Aligned-View code path (apply_alignment,
|
||
_aligned_canvas_axes, the pixel-inspector inverse-transform) works
|
||
completely unmodified.
|
||
|
||
Non-modal by design (shown via .show(), never .exec() or setModal(True))
|
||
so the user can still interact with the main window. Talks back to
|
||
SrasViewerWindow two ways: it reuses parent._run_worker/_jobs directly
|
||
for its background mask-fetch and cross-correlate steps, so the main
|
||
window's existing shutdown/lifecycle plumbing covers both for free, and
|
||
it emits alignment_saved / alignment_cleared signals for the two moments
|
||
that should actually mutate the main window's persistent state —
|
||
everything else (nudging, Auto De-rotate, Auto Cross-Correlate, threshold
|
||
edits) stays purely local to this dialog until Save.
|
||
"""
|
||
|
||
alignment_saved = pyqtSignal(object, str) # AlignmentResult, sidecar path (str)
|
||
alignment_cleared = pyqtSignal()
|
||
|
||
_PREVIEW_MARGIN_FRAC = 0.15
|
||
_BASE_ALPHA = 0.42
|
||
_ACTIVE_ALPHA = 0.75
|
||
_MAX_PREVIEW_DIM = 1024
|
||
|
||
# (label, sources passed to compute.register_angle_to_reference). "Both"
|
||
# registers on each and keeps whichever scores higher per angle, which
|
||
# costs roughly double but removes the failure mode where the single
|
||
# chosen source is the one that happens to be uninformative for one angle.
|
||
_CORRELATE_SOURCES = (
|
||
("Both, keep best (recommended)", ("signal", "mask")),
|
||
("Raw signal", ("signal",)),
|
||
("Thresholded mask", ("mask",)),
|
||
)
|
||
|
||
def __init__(self, parent: "SrasViewerWindow", sras: SrasFile, *,
|
||
ref_angle_idx: int, dc_threshold_mv: float,
|
||
seed_per_angle: dict[int, ManualAngleParams] | None,
|
||
cached_dc4_mv: dict[int, np.ndarray]):
|
||
super().__init__(parent)
|
||
self._parent = parent
|
||
self._sras = sras
|
||
self._ref_angle_idx = ref_angle_idx
|
||
self._downsample = (1, 1) # (rows, cols) block-mean factors
|
||
self._dc4_mv: dict[int, np.ndarray] = {}
|
||
self._masks_small: dict[int, np.ndarray] = {}
|
||
self._preview_layers: dict[int, np.ndarray] = {}
|
||
self._preview_origin_mm = (0.0, 0.0)
|
||
self._preview_shape = (1, 1)
|
||
self._preview_pitch_mm = (1.0, 1.0)
|
||
self._masks_ready = False
|
||
self._fit_notes: dict[int, tuple[float, str]] = {}
|
||
self._derotate_sign_flipped = False
|
||
|
||
self.setWindowTitle(f"Manual Alignment — {sras.path.name}")
|
||
self.resize(1150, 760)
|
||
|
||
self._seed_initial_params(seed_per_angle)
|
||
n = sras.n_angles
|
||
cmap = mpl.colormaps["tab10"] if n <= 10 else mpl.colormaps["tab20"]
|
||
self._angle_colors = {a: cmap(a % cmap.N)[:3] for a in range(n)}
|
||
self._active_angle = 1 if ref_angle_idx == 0 and n > 1 else 0
|
||
|
||
self._build_ui(dc_threshold_mv)
|
||
self._set_controls_enabled(False) # re-enabled once masks are ready
|
||
self._start_mask_prep(cached_dc4_mv)
|
||
|
||
def showEvent(self, event):
|
||
super().showEvent(event)
|
||
self.canvas.setFocus()
|
||
|
||
# ------------------------------------------------------------------
|
||
# Construction
|
||
# ------------------------------------------------------------------
|
||
|
||
def _seed_initial_params(self, seed_per_angle: dict[int, ManualAngleParams] | None):
|
||
seed = seed_per_angle or {}
|
||
self._angle_params: dict[int, ManualAngleParams] = {
|
||
a: (ManualAngleParams(seed[a].rotation_deg, seed[a].shift_mm)
|
||
if a in seed else ManualAngleParams())
|
||
for a in range(self._sras.n_angles)
|
||
}
|
||
self._angle_params[self._ref_angle_idx] = ManualAngleParams()
|
||
|
||
def _build_ui(self, dc_threshold_mv: float):
|
||
root = QHBoxLayout(self)
|
||
|
||
self.canvas = ManualAlignOverlayCanvas()
|
||
left = QWidget()
|
||
left_l = QVBoxLayout(left)
|
||
left_l.setContentsMargins(0, 0, 0, 0)
|
||
left_l.setSpacing(4)
|
||
left_l.addWidget(NavigationToolbar2QT(self.canvas, left))
|
||
left_l.addWidget(self.canvas)
|
||
root.addWidget(left, stretch=1)
|
||
|
||
panel = QWidget()
|
||
panel_l = QVBoxLayout(panel)
|
||
panel_l.setContentsMargins(0, 0, 0, 0)
|
||
panel_l.setSpacing(8)
|
||
panel_l.addWidget(self._build_angle_group())
|
||
panel_l.addWidget(self._build_adjust_group())
|
||
panel_l.addWidget(self._build_step_group())
|
||
panel_l.addWidget(self._build_threshold_group(dc_threshold_mv))
|
||
panel_l.addWidget(self._build_correlate_group())
|
||
panel_l.addWidget(self._build_actions_group())
|
||
self.lbl_status = _wrap_label("", _CSS_MUTED)
|
||
panel_l.addWidget(self.lbl_status)
|
||
panel_l.addStretch()
|
||
|
||
root.addWidget(_scroll_panel(panel, 320))
|
||
self._connect_controls()
|
||
|
||
def _build_angle_group(self) -> QWidget:
|
||
grp_angle, al = _group("Active Angle")
|
||
self.combo_active_angle = QComboBox()
|
||
for a in range(self._sras.n_angles):
|
||
label = f"Angle {a} ({self._sras.angles_deg[a]:.1f}°)"
|
||
if a == self._ref_angle_idx:
|
||
label += " [reference]"
|
||
self.combo_active_angle.addItem(label)
|
||
al.addWidget(self.combo_active_angle)
|
||
self.lbl_active_note = _wrap_label("", _CSS_WARN)
|
||
al.addWidget(self.lbl_active_note)
|
||
return grp_angle
|
||
|
||
def _build_adjust_group(self) -> QWidget:
|
||
self.grp_manual_adjust, mform_box = _group("Manual Adjustment")
|
||
mform = _form()
|
||
self.spin_active_rotation_deg = _make_dspin(-3600.0, 3600.0, 3, suffix=" °")
|
||
mform.addRow("Rotation:", self.spin_active_rotation_deg)
|
||
|
||
self.spin_active_shift_x_mm = _make_dspin(-1e5, 1e5, 4, suffix=" mm")
|
||
mform.addRow("Shift X:", self.spin_active_shift_x_mm)
|
||
|
||
self.spin_active_shift_y_mm = _make_dspin(-1e5, 1e5, 4, suffix=" mm")
|
||
mform.addRow("Shift Y:", self.spin_active_shift_y_mm)
|
||
mform_box.addLayout(mform)
|
||
return self.grp_manual_adjust
|
||
|
||
def _build_step_group(self) -> QWidget:
|
||
self.grp_step_sizes, sl = _group("Nudge Step Sizes")
|
||
sform = _form()
|
||
self.spin_step_translate_mm = _make_dspin(0.0001, 1000.0, 4,
|
||
suffix=" mm", value=0.01)
|
||
sform.addRow("Translate step:", self.spin_step_translate_mm)
|
||
|
||
self.spin_step_rotate_deg = _make_dspin(0.001, 90.0, 3,
|
||
suffix=" °", value=0.1)
|
||
sform.addRow("Rotate step:", self.spin_step_rotate_deg)
|
||
|
||
self.spin_step_multiplier = _make_dspin(1.0, 1000.0, 1, value=10.0)
|
||
sform.addRow("Coarse × (Shift):", self.spin_step_multiplier)
|
||
sl.addLayout(sform)
|
||
sl.addWidget(_wrap_label(
|
||
"Arrow keys nudge X/Y translation; Q/E nudge rotation (CCW/CW). "
|
||
"Hold Shift for the coarse step. Click the image once so it has "
|
||
"keyboard focus.", _CSS_HINT))
|
||
return self.grp_step_sizes
|
||
|
||
def _build_threshold_group(self, dc_threshold_mv: float) -> QWidget:
|
||
self.grp_mask_threshold, tl = _group("Mask Threshold")
|
||
tform = _form()
|
||
self.spin_mask_threshold_mv = _make_dspin(-500.0, 500.0, 3,
|
||
suffix=" mV", value=dc_threshold_mv)
|
||
tform.addRow("DC threshold:", self.spin_mask_threshold_mv)
|
||
tl.addLayout(tform)
|
||
return self.grp_mask_threshold
|
||
|
||
def _build_correlate_group(self) -> QWidget:
|
||
self.grp_correlate, cl = _group("Cross-Correlate (FFT)")
|
||
cform = _form()
|
||
self.combo_correlate_source = QComboBox()
|
||
for label, sources in self._CORRELATE_SOURCES:
|
||
self.combo_correlate_source.addItem(label, sources)
|
||
cform.addRow("Correlate on:", self.combo_correlate_source)
|
||
|
||
self.spin_correlate_search_deg = _make_dspin(0.0, 180.0, 1, suffix=" °",
|
||
value=6.0, step=1.0)
|
||
cform.addRow("Rotation search (±):", self.spin_correlate_search_deg)
|
||
cl.addLayout(cform)
|
||
self.btn_auto_correlate = QPushButton("Auto Cross-Correlate (vs Reference)")
|
||
cl.addWidget(self.btn_auto_correlate)
|
||
cl.addWidget(_wrap_label(
|
||
"Finds each non-reference angle's rotation *and* translation by "
|
||
"cross-correlating its image against the reference's — the stage's "
|
||
"reported angle is only the starting point of the search, and both "
|
||
"of its signs are tried. Run this first, then nudge only for small "
|
||
"corrections.", _CSS_HINT))
|
||
return self.grp_correlate
|
||
|
||
def _build_actions_group(self) -> QWidget:
|
||
grp_actions, acl = _group("Actions")
|
||
self.btn_auto_derotate = QPushButton("Auto De-rotate (use known angles)")
|
||
self.btn_save = QPushButton("Save Alignment")
|
||
self.btn_clear = QPushButton("Clear Alignment…")
|
||
self.btn_close = QPushButton("Close")
|
||
for btn in (self.btn_auto_derotate, self.btn_save, self.btn_clear, self.btn_close):
|
||
acl.addWidget(btn)
|
||
return grp_actions
|
||
|
||
def _connect_controls(self):
|
||
self.combo_active_angle.currentIndexChanged.connect(self._on_active_angle_changed)
|
||
self.spin_active_rotation_deg.editingFinished.connect(self._on_rotation_spin_edited)
|
||
self.spin_active_shift_x_mm.editingFinished.connect(self._on_shift_spin_edited)
|
||
self.spin_active_shift_y_mm.editingFinished.connect(self._on_shift_spin_edited)
|
||
self.spin_mask_threshold_mv.editingFinished.connect(self._on_mask_threshold_edited)
|
||
self.btn_auto_derotate.clicked.connect(self._on_auto_derotate)
|
||
self.btn_auto_correlate.clicked.connect(self._on_auto_correlate)
|
||
self.btn_save.clicked.connect(self._on_save)
|
||
self.btn_clear.clicked.connect(self._on_clear)
|
||
self.btn_close.clicked.connect(self.close)
|
||
self.canvas.nudge_translate.connect(self._on_nudge_translate)
|
||
self.canvas.nudge_rotate.connect(self._on_nudge_rotate)
|
||
|
||
with QSignalBlocker(self.combo_active_angle):
|
||
self.combo_active_angle.setCurrentIndex(self._active_angle)
|
||
self._on_active_angle_changed(self._active_angle)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Mask preparation (initial CH4 fetch + threshold + downsample)
|
||
# ------------------------------------------------------------------
|
||
|
||
def _start_mask_prep(self, cached_dc4_mv: dict[int, np.ndarray]):
|
||
self._dc4_mv = dict(cached_dc4_mv)
|
||
missing = [a for a in range(self._sras.n_angles) if a not in self._dc4_mv]
|
||
if not missing:
|
||
self._finish_mask_prep()
|
||
return
|
||
self.lbl_status.setText(f"Preparing masks: 0/{len(missing)} angle(s) needed…")
|
||
started = self._parent._run_worker(
|
||
Jobs.MANUAL_ALIGN_MASKS, Ch4MaskWorker(self._sras, missing),
|
||
connect=(
|
||
("angle_done", self._on_mask_angle_done),
|
||
("error", lambda msg: self.lbl_status.setText(f"Mask prep error: {msg}")),
|
||
),
|
||
on_done=self._finish_mask_prep)
|
||
if not started:
|
||
self.lbl_status.setText(
|
||
"Could not start mask preparation (busy) — close and reopen.")
|
||
|
||
def _on_mask_angle_done(self, angle_idx: int, dc4_mv: np.ndarray):
|
||
self._dc4_mv[angle_idx] = dc4_mv
|
||
self.lbl_status.setText(
|
||
f"Preparing masks: {len(self._dc4_mv)}/{self._sras.n_angles} ready…")
|
||
|
||
def _finish_mask_prep(self):
|
||
if len(self._dc4_mv) < self._sras.n_angles:
|
||
return # a mask-worker error left some angles unfetched
|
||
# Rows and columns get their own factor. A real scan is ~7500 frames
|
||
# wide but only ~750 rows tall, so one shared factor sized for the
|
||
# frames would throw away 8x more row detail than the preview needs and
|
||
# leave the overlay too coarse in y to judge alignment by eye.
|
||
max_rows = max(img.shape[0] for img in self._dc4_mv.values())
|
||
max_cols = max(img.shape[1] for img in self._dc4_mv.values())
|
||
self._downsample = (
|
||
max(1, int(np.ceil(max_rows / self._MAX_PREVIEW_DIM))),
|
||
max(1, int(np.ceil(max_cols / self._MAX_PREVIEW_DIM))))
|
||
self._recompute_masks_small()
|
||
self._rebuild_preview_canvas()
|
||
self._set_controls_enabled(True)
|
||
self.lbl_status.setText("Ready.")
|
||
|
||
def _recompute_masks_small(self):
|
||
"""Threshold + downsample every angle's already-in-memory full-res
|
||
CH4 mV image. Cheap (a compare + block-mean), so this re-runs in
|
||
full whenever the mask-threshold spin box changes — no re-fetch.
|
||
Purely for the overlay's visuals: no alignment geometry depends on this
|
||
threshold, only which pixels the overlay paints."""
|
||
threshold = self.spin_mask_threshold_mv.value()
|
||
fy, fx = self._downsample
|
||
self._masks_small = {
|
||
a: compute.block_mean_2d((img >= threshold).astype(np.float32), fy, fx)
|
||
for a, img in self._dc4_mv.items()
|
||
}
|
||
|
||
# ------------------------------------------------------------------
|
||
# Preview canvas: full rebuild vs. incremental single-layer refresh
|
||
# ------------------------------------------------------------------
|
||
|
||
def _rebuild_preview_canvas(self):
|
||
"""Full geometry rebuild: recomputes the shared preview canvas's
|
||
origin/shape (rotation can grow the union bbox — translation alone
|
||
cannot, per the padding baked in via _PREVIEW_MARGIN_FRAC) and every
|
||
angle's reprojected mask layer. Triggered by: dialog open,
|
||
mask-threshold change, Auto De-rotate, a rotation nudge/edit of the
|
||
active angle. NOT triggered by a translation-only nudge — see
|
||
_refresh_active_preview_layer."""
|
||
dx_ref, dy_ref = compute.pixel_pitch_mm(self._sras, self._ref_angle_idx)
|
||
fy, fx = self._downsample
|
||
pitch = (dx_ref * fx, dy_ref * fy)
|
||
origin, shape = compute.canvas_for_params(
|
||
self._sras, self._ref_angle_idx, pitch, self._angle_params,
|
||
margin_frac=self._PREVIEW_MARGIN_FRAC, snap=False)
|
||
self._preview_origin_mm, self._preview_shape = origin, shape
|
||
self._preview_pitch_mm = pitch
|
||
self._preview_layers = {
|
||
a: self._reproject(a) for a in range(self._sras.n_angles)
|
||
}
|
||
self._redraw_overlay()
|
||
|
||
def _reproject(self, angle_idx: int) -> np.ndarray:
|
||
"""One angle's downsampled mask on the current preview canvas.
|
||
src_downsample must match _masks_small's block-mean factors, or the
|
||
layer lands magnified and offset instead of where the alignment
|
||
actually puts it."""
|
||
p = self._angle_params[angle_idx]
|
||
return compute.reproject_mask(
|
||
self._sras, angle_idx, self._ref_angle_idx,
|
||
self._masks_small[angle_idx], p.rotation_deg, p.shift_mm,
|
||
self._preview_pitch_mm, self._preview_origin_mm, self._preview_shape,
|
||
src_downsample=self._downsample)
|
||
|
||
def _refresh_active_preview_layer(self):
|
||
"""Cheap path for a translation-only nudge/edit of the active angle:
|
||
reproject just that one angle's downsampled mask onto the *existing*
|
||
preview canvas — every other angle's cached layer is untouched."""
|
||
self._preview_layers[self._active_angle] = self._reproject(self._active_angle)
|
||
self._redraw_overlay()
|
||
|
||
def _redraw_overlay(self):
|
||
"""Alpha-composite every angle's colored mask layer into one RGBA
|
||
image ("all thresholds overlaid with varying opacity"). Each angle
|
||
keeps a fixed, distinct color regardless of which is active; the
|
||
active angle is drawn last (on top) at a visibly higher alpha so
|
||
it's easy to track while nudging."""
|
||
if not self._preview_layers:
|
||
return # mask prep hasn't finished yet — nothing to draw
|
||
n_rows, n_cols = self._preview_shape
|
||
rgba = np.zeros((n_rows, n_cols, 4), dtype=np.float32)
|
||
order = sorted(range(self._sras.n_angles), key=lambda a: a == self._active_angle)
|
||
for a in order:
|
||
layer = self._preview_layers.get(a)
|
||
if layer is None:
|
||
continue
|
||
alpha = self._ACTIVE_ALPHA if a == self._active_angle else self._BASE_ALPHA
|
||
color = self._angle_colors[a]
|
||
fg_a = layer * alpha
|
||
for c in range(3):
|
||
rgba[..., c] = color[c] * fg_a + rgba[..., c] * rgba[..., 3] * (1 - fg_a)
|
||
rgba[..., 3] = fg_a + rgba[..., 3] * (1 - fg_a)
|
||
|
||
x0, y0 = self._preview_origin_mm
|
||
dx, dy = self._preview_pitch_mm
|
||
x_axis = x0 + np.arange(n_cols) * dx
|
||
y_axis = y0 + np.arange(n_rows) * dy
|
||
extent = _axes_extent(x_axis, y_axis, dx, dy)
|
||
title = (f"Angle {self._active_angle} active "
|
||
f"({self._sras.angles_deg[self._active_angle]:.1f}°)")
|
||
self.canvas.show_overlay(rgba, extent, title)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Angle selection / nudge / edit handlers
|
||
# ------------------------------------------------------------------
|
||
|
||
def _on_active_angle_changed(self, angle_idx: int):
|
||
self._active_angle = angle_idx
|
||
is_ref = angle_idx == self._ref_angle_idx
|
||
self.grp_manual_adjust.setEnabled(self._masks_ready and not is_ref)
|
||
self.lbl_active_note.setText(
|
||
"Reference angle — defines the shared origin, not adjustable." if is_ref else "")
|
||
self._sync_active_spinboxes()
|
||
self._redraw_overlay()
|
||
|
||
def _sync_active_spinboxes(self):
|
||
p = self._angle_params[self._active_angle]
|
||
for spin, val in ((self.spin_active_rotation_deg, p.rotation_deg),
|
||
(self.spin_active_shift_x_mm, p.shift_mm[0]),
|
||
(self.spin_active_shift_y_mm, p.shift_mm[1])):
|
||
with QSignalBlocker(spin):
|
||
spin.setValue(val)
|
||
|
||
def _on_nudge_translate(self, dir_x: int, dir_y: int, coarse: bool):
|
||
if not self._masks_ready or self._active_angle == self._ref_angle_idx:
|
||
return
|
||
step = self.spin_step_translate_mm.value()
|
||
if coarse:
|
||
step *= self.spin_step_multiplier.value()
|
||
p = self._angle_params[self._active_angle]
|
||
p.shift_mm = (p.shift_mm[0] + dir_x * step, p.shift_mm[1] + dir_y * step)
|
||
self._sync_active_spinboxes()
|
||
self._refresh_active_preview_layer()
|
||
|
||
def _on_nudge_rotate(self, direction: int, coarse: bool):
|
||
if not self._masks_ready or self._active_angle == self._ref_angle_idx:
|
||
return
|
||
step = self.spin_step_rotate_deg.value()
|
||
if coarse:
|
||
step *= self.spin_step_multiplier.value()
|
||
self._angle_params[self._active_angle].rotation_deg += direction * step
|
||
self._sync_active_spinboxes()
|
||
self._rebuild_preview_canvas()
|
||
|
||
def _on_rotation_spin_edited(self):
|
||
if self._active_angle == self._ref_angle_idx:
|
||
return
|
||
self._angle_params[self._active_angle].rotation_deg = self.spin_active_rotation_deg.value()
|
||
self._rebuild_preview_canvas()
|
||
|
||
def _on_shift_spin_edited(self):
|
||
if self._active_angle == self._ref_angle_idx:
|
||
return
|
||
p = self._angle_params[self._active_angle]
|
||
p.shift_mm = (self.spin_active_shift_x_mm.value(), self.spin_active_shift_y_mm.value())
|
||
self._refresh_active_preview_layer()
|
||
|
||
def _on_mask_threshold_edited(self):
|
||
if not self._masks_ready:
|
||
return
|
||
self._recompute_masks_small()
|
||
self._rebuild_preview_canvas()
|
||
|
||
# ------------------------------------------------------------------
|
||
# Actions
|
||
# ------------------------------------------------------------------
|
||
|
||
def _on_auto_derotate(self):
|
||
"""Seed every angle's rotation from the stage's reported angle.
|
||
|
||
A starting point for nudging by eye, not an alignment: the stage's
|
||
sign convention relative to this module's is not knowable from the
|
||
file, so the sign that lines the scans up is whichever of the two looks
|
||
right in the overlay. Auto Cross-Correlate decides that from the images
|
||
instead, and is the button to reach for first.
|
||
"""
|
||
sign = -1.0 if self._derotate_sign_flipped else 1.0
|
||
self._derotate_sign_flipped = not self._derotate_sign_flipped
|
||
n_changed = 0
|
||
for a in range(self._sras.n_angles):
|
||
if a == self._ref_angle_idx:
|
||
continue
|
||
self._angle_params[a].rotation_deg = sign * compute.nominal_delta_deg(
|
||
self._sras, a, self._ref_angle_idx)
|
||
n_changed += 1
|
||
self._sync_active_spinboxes()
|
||
self._rebuild_preview_canvas()
|
||
self.lbl_status.setText(
|
||
f"Rotation set to the stage angle ({'−' if sign < 0 else '+'}delta) "
|
||
f"for {n_changed} angle(s); translation untouched. Click again to "
|
||
"try the opposite sign.")
|
||
|
||
def _on_auto_correlate(self):
|
||
if not self._masks_ready:
|
||
return
|
||
angles = [a for a in range(self._sras.n_angles) if a != self._ref_angle_idx]
|
||
if not angles:
|
||
return
|
||
worker = CrossCorrelateWorker(
|
||
self._sras, self._ref_angle_idx, angles, self._dc4_mv,
|
||
sources=self.combo_correlate_source.currentData(),
|
||
dc_threshold_mv=self.spin_mask_threshold_mv.value(),
|
||
search_deg=self.spin_correlate_search_deg.value())
|
||
self._correlate_done_count = 0
|
||
self._correlate_total = len(angles)
|
||
self._fit_notes = {}
|
||
self._set_controls_enabled(False)
|
||
self.lbl_status.setText(f"Cross-correlating: 0/{self._correlate_total} angle(s)…")
|
||
started = self._parent._run_worker(
|
||
Jobs.MANUAL_ALIGN_CORRELATE, worker,
|
||
connect=(
|
||
("angle_done", self._on_correlate_angle_done),
|
||
("error", self._on_correlate_error),
|
||
),
|
||
on_done=self._finish_auto_correlate)
|
||
if not started:
|
||
self._set_controls_enabled(True)
|
||
self.lbl_status.setText("Could not start cross-correlation (busy) — try again.")
|
||
|
||
def _on_correlate_angle_done(self, angle_idx: int, rotation_deg: float,
|
||
shift_x_mm: float, shift_y_mm: float,
|
||
score: float, source: str):
|
||
self._angle_params[angle_idx] = ManualAngleParams(rotation_deg, (shift_x_mm, shift_y_mm))
|
||
self._fit_notes[angle_idx] = (score, source)
|
||
self._correlate_done_count += 1
|
||
self.lbl_status.setText(
|
||
f"Cross-correlating: {self._correlate_done_count}/{self._correlate_total} angle(s)…")
|
||
|
||
def _on_correlate_error(self, msg: str):
|
||
self.lbl_status.setText(f"Cross-correlation error: {msg}")
|
||
|
||
def _finish_auto_correlate(self):
|
||
self._sync_active_spinboxes()
|
||
self._rebuild_preview_canvas()
|
||
self._set_controls_enabled(True)
|
||
self.lbl_status.setText(
|
||
f"Cross-correlated {self._correlate_done_count} angle(s) against "
|
||
f"Angle {self._ref_angle_idx}.\n" + self._fit_report())
|
||
|
||
def _fit_report(self) -> str:
|
||
"""Per-angle registration quality, worst first.
|
||
|
||
Surfaced rather than buried because a single bad acquisition (stage
|
||
glitch, laser dropout) registers poorly and would otherwise be fused in
|
||
silently — seeing which angle it is, is what makes dropping it with
|
||
sras_edit_scans.py actionable. The deviation from the stage's own
|
||
reported angle is shown alongside: a large one means the search and the
|
||
stage disagree, which is either a genuine mechanical error or a sign
|
||
that this angle's fit is not to be trusted.
|
||
"""
|
||
if not self._fit_notes:
|
||
return ""
|
||
rows = sorted(self._fit_notes.items(), key=lambda kv: kv[1][0])
|
||
worst = rows[0]
|
||
lines = [f"Worst fit: angle {worst[0]} (score {worst[1][0]:.3f}, "
|
||
f"{worst[1][1]})."]
|
||
drifted = []
|
||
for a, _note in rows:
|
||
nominal = compute.nominal_delta_deg(self._sras, a, self._ref_angle_idx)
|
||
got = self._angle_params[a].rotation_deg
|
||
dev = min(abs(got - nominal), abs(got + nominal))
|
||
if dev > 1.0:
|
||
drifted.append(f"{a} ({dev:.2f}°)")
|
||
if drifted:
|
||
lines.append("Rotation differs from the stage angle by >1° for "
|
||
"angle(s) " + ", ".join(drifted) + ".")
|
||
lines.append("Nudge from here for any remaining fine correction.")
|
||
return " ".join(lines)
|
||
|
||
def _on_save(self):
|
||
threshold = self.spin_mask_threshold_mv.value()
|
||
resolved = dict(self._angle_params) # already concrete floats
|
||
try:
|
||
path = save_manual_alignment(self._sras, self._ref_angle_idx, threshold, resolved)
|
||
result = build_manual_alignment(self._sras, self._ref_angle_idx,
|
||
threshold, resolved)
|
||
except OSError as exc:
|
||
QMessageBox.warning(self, "Save Alignment Failed", str(exc))
|
||
return
|
||
self.lbl_status.setText(f"Saved to {path.name}.")
|
||
self.alignment_saved.emit(result, str(path))
|
||
|
||
def _on_clear(self):
|
||
reply = QMessageBox.question(
|
||
self, "Clear Alignment",
|
||
"This resets every angle back to raw/unaligned (0° rotation, no "
|
||
"shift) and deletes the saved alignment file for this scan, if "
|
||
"any. This cannot be undone. Continue?",
|
||
QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No,
|
||
QMessageBox.StandardButton.No)
|
||
if reply != QMessageBox.StandardButton.Yes:
|
||
return
|
||
try:
|
||
existed = delete_manual_alignment(self._sras)
|
||
except OSError as exc:
|
||
QMessageBox.warning(self, "Clear Alignment Failed",
|
||
f"Could not delete the saved alignment file: {exc}")
|
||
return
|
||
self._angle_params = {a: ManualAngleParams() for a in range(self._sras.n_angles)}
|
||
self._fit_notes = {}
|
||
self._sync_active_spinboxes()
|
||
self._rebuild_preview_canvas()
|
||
self.lbl_status.setText(
|
||
"Alignment cleared; saved file removed." if existed
|
||
else "Alignment cleared (there was no saved file).")
|
||
self.alignment_cleared.emit()
|
||
|
||
def _set_controls_enabled(self, enabled: bool):
|
||
self._masks_ready = enabled
|
||
self.combo_active_angle.setEnabled(enabled)
|
||
self.grp_manual_adjust.setEnabled(enabled and self._active_angle != self._ref_angle_idx)
|
||
self.grp_step_sizes.setEnabled(enabled)
|
||
self.grp_mask_threshold.setEnabled(enabled)
|
||
self.grp_correlate.setEnabled(enabled)
|
||
self.btn_auto_derotate.setEnabled(enabled)
|
||
self.btn_save.setEnabled(enabled)
|
||
self.btn_clear.setEnabled(enabled)
|
||
|
||
|