Split sras_viewer.py into a package (pure move)
The 2790-line module becomes sras_viewer/: common.py (constants + layout helpers), canvases.py (RoiQuad, ImageCanvas, WaveformCanvas, ManualAlignOverlayCanvas), dialogs.py (FftOptionsDialog, ManualAlignmentDialog), main_window.py (SrasViewerWindow + main), with __init__ re-exporting the public names and __main__ keeping `python -m sras_viewer` working. pyproject gains a `sras-viewer` console script. Code moved verbatim; only import headers are new (pyflakes-clean). tests/test_gui.py patch targets follow the classes to their new modules. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Thomas Ales
@@ -0,0 +1,767 @@
|
||||
"""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 pyqtSignal
|
||||
from PyQt6.QtWidgets import (
|
||||
QButtonGroup, QComboBox, QDialog, QDialogButtonBox, QDoubleSpinBox,
|
||||
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, _SPIN_MIN_W, _form, _group,
|
||||
_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):
|
||||
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
|
||||
|
||||
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)
|
||||
|
||||
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()
|
||||
|
||||
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)
|
||||
|
||||
# ---- Active Angle -------------------------------------------------
|
||||
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)
|
||||
panel_l.addWidget(grp_angle)
|
||||
|
||||
# ---- Manual Adjustment ---------------------------------------------
|
||||
self.grp_manual_adjust, mform_box = _group("Manual Adjustment")
|
||||
mform = _form()
|
||||
self.spin_active_rotation_deg = QDoubleSpinBox()
|
||||
self.spin_active_rotation_deg.setRange(-3600.0, 3600.0)
|
||||
self.spin_active_rotation_deg.setDecimals(3)
|
||||
self.spin_active_rotation_deg.setSuffix(" °")
|
||||
self.spin_active_rotation_deg.setMinimumWidth(_SPIN_MIN_W)
|
||||
mform.addRow("Rotation:", self.spin_active_rotation_deg)
|
||||
|
||||
self.spin_active_shift_x_mm = QDoubleSpinBox()
|
||||
self.spin_active_shift_x_mm.setRange(-1e5, 1e5)
|
||||
self.spin_active_shift_x_mm.setDecimals(4)
|
||||
self.spin_active_shift_x_mm.setSuffix(" mm")
|
||||
self.spin_active_shift_x_mm.setMinimumWidth(_SPIN_MIN_W)
|
||||
mform.addRow("Shift X:", self.spin_active_shift_x_mm)
|
||||
|
||||
self.spin_active_shift_y_mm = QDoubleSpinBox()
|
||||
self.spin_active_shift_y_mm.setRange(-1e5, 1e5)
|
||||
self.spin_active_shift_y_mm.setDecimals(4)
|
||||
self.spin_active_shift_y_mm.setSuffix(" mm")
|
||||
self.spin_active_shift_y_mm.setMinimumWidth(_SPIN_MIN_W)
|
||||
mform.addRow("Shift Y:", self.spin_active_shift_y_mm)
|
||||
mform_box.addLayout(mform)
|
||||
panel_l.addWidget(self.grp_manual_adjust)
|
||||
|
||||
# ---- Nudge Step Sizes ------------------------------------------------
|
||||
self.grp_step_sizes, sl = _group("Nudge Step Sizes")
|
||||
sform = _form()
|
||||
self.spin_step_translate_mm = QDoubleSpinBox()
|
||||
self.spin_step_translate_mm.setRange(0.0001, 1000.0)
|
||||
self.spin_step_translate_mm.setDecimals(4)
|
||||
self.spin_step_translate_mm.setSuffix(" mm")
|
||||
self.spin_step_translate_mm.setValue(0.01)
|
||||
self.spin_step_translate_mm.setMinimumWidth(_SPIN_MIN_W)
|
||||
sform.addRow("Translate step:", self.spin_step_translate_mm)
|
||||
|
||||
self.spin_step_rotate_deg = QDoubleSpinBox()
|
||||
self.spin_step_rotate_deg.setRange(0.001, 90.0)
|
||||
self.spin_step_rotate_deg.setDecimals(3)
|
||||
self.spin_step_rotate_deg.setSuffix(" °")
|
||||
self.spin_step_rotate_deg.setValue(0.1)
|
||||
self.spin_step_rotate_deg.setMinimumWidth(_SPIN_MIN_W)
|
||||
sform.addRow("Rotate step:", self.spin_step_rotate_deg)
|
||||
|
||||
self.spin_step_multiplier = QDoubleSpinBox()
|
||||
self.spin_step_multiplier.setRange(1.0, 1000.0)
|
||||
self.spin_step_multiplier.setDecimals(1)
|
||||
self.spin_step_multiplier.setValue(10.0)
|
||||
self.spin_step_multiplier.setMinimumWidth(_SPIN_MIN_W)
|
||||
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))
|
||||
panel_l.addWidget(self.grp_step_sizes)
|
||||
|
||||
# ---- Mask Threshold ---------------------------------------------------
|
||||
self.grp_mask_threshold, tl = _group("Mask Threshold")
|
||||
tform = _form()
|
||||
self.spin_mask_threshold_mv = QDoubleSpinBox()
|
||||
self.spin_mask_threshold_mv.setRange(-500.0, 500.0)
|
||||
self.spin_mask_threshold_mv.setDecimals(3)
|
||||
self.spin_mask_threshold_mv.setSuffix(" mV")
|
||||
self.spin_mask_threshold_mv.setValue(dc_threshold_mv)
|
||||
self.spin_mask_threshold_mv.setMinimumWidth(_SPIN_MIN_W)
|
||||
tform.addRow("DC threshold:", self.spin_mask_threshold_mv)
|
||||
tl.addLayout(tform)
|
||||
panel_l.addWidget(self.grp_mask_threshold)
|
||||
|
||||
# ---- Cross-Correlate (FFT) -----------------------------------------
|
||||
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 = QDoubleSpinBox()
|
||||
self.spin_correlate_search_deg.setRange(0.0, 180.0)
|
||||
self.spin_correlate_search_deg.setSingleStep(1.0)
|
||||
self.spin_correlate_search_deg.setDecimals(1)
|
||||
self.spin_correlate_search_deg.setSuffix(" °")
|
||||
self.spin_correlate_search_deg.setValue(6.0)
|
||||
self.spin_correlate_search_deg.setMinimumWidth(_SPIN_MIN_W)
|
||||
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))
|
||||
panel_l.addWidget(self.grp_correlate)
|
||||
|
||||
# ---- Actions ------------------------------------------------------
|
||||
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)
|
||||
panel_l.addWidget(grp_actions)
|
||||
|
||||
self.lbl_status = _wrap_label("", _CSS_MUTED)
|
||||
panel_l.addWidget(self.lbl_status)
|
||||
panel_l.addStretch()
|
||||
|
||||
root.addWidget(_scroll_panel(panel, 320))
|
||||
|
||||
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)
|
||||
|
||||
self.combo_active_angle.blockSignals(True)
|
||||
self.combo_active_angle.setCurrentIndex(self._active_angle)
|
||||
self.combo_active_angle.blockSignals(False)
|
||||
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(
|
||||
"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 = [x_axis[0] - dx / 2, x_axis[-1] + dx / 2,
|
||||
y_axis[-1] + dy / 2, y_axis[0] - dy / 2]
|
||||
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])):
|
||||
spin.blockSignals(True)
|
||||
spin.setValue(val)
|
||||
spin.blockSignals(False)
|
||||
|
||||
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(
|
||||
"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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user