Thomas Ales
6d0e30b9ce
The alignment machinery never modified a scan: it produced an
AlignmentResult and every consumer resampled on the fly. That is right
for a viewer but means the aligned stack cannot leave the process, so no
other tool can read it and reopening the scan redoes the registration.
sras_align_export.write_aligned_sras bakes an alignment into a new v6
file: each angle is gathered onto the shared canvas by nearest neighbour,
so every output angle shares one grid and the file opens already aligned.
Registering the export against itself returns identity, which is the test
that pins the whole index chain.
Three details that are easy to get wrong and are now covered:
* Rounding must be floor(x + 0.5), not np.rint. scipy's order=0 rounds
halves away from zero, and the canvas is snapped to the reference's
pixel grid, so exact halves are common rather than hypothetical.
* Out-of-bounds must be tested on the fractional coordinate against
[0, n-1], not on the rounded index, or a one-pixel rim gets real data
everywhere the Aligned View shows padding.
* ...but with a tolerance, because the mm-space affine chain lands an
exactly-integer transform a few times 1e-13 off. A bare >= 0 drops
the *reference* angle's entire first row and last column.
Padding is the per-channel ADC code nearest 0 mV, not zero: zero ADC
decodes to ~+100 mV on real calibration and would masquerade as sample.
Source rows are served from sliding in-RAM bands. A rotated angle maps
one output row to a diagonal across the source, so indexing a memmap in
output order refaults nearly the whole angle per row — terabytes of
paging for a gigabyte of data.
Also in compute: crop_alignment_result (a crop is pure index translation,
so it folds into the affine's offset rather than becoming a second
transform), overlap_stats, largest_rect_at_least for a crop that stays
inside the overlap region, and seed/sign/refine knobs on
register_angle_to_reference whose defaults leave existing behaviour and
tests byte-identical.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
838 lines
37 KiB
Python
838 lines
37 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 AlignOverlayCanvas
|
||
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):
|
||
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())
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Row-Averaged FFT Options dialog
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class RowAverageFftOptionsDialog(QDialog):
|
||
"""Configure the same-row, distance-weighted neighbor averaging applied
|
||
to each pixel's CH1 waveform before 'Batch Compute Row-Averaged FFT and
|
||
Store' re-runs the FFT peak search — a same-row SNR cleanup pass, never
|
||
mixing across rows/Y (see sras_compute._row_average_waveforms).
|
||
|
||
Unlike the plain FFT batch action (which stores unmasked and defers
|
||
masking to display time), the DC threshold here is required up front:
|
||
it decides which same-row neighbors are eligible to contribute to a
|
||
pixel's average, so it can't be deferred.
|
||
|
||
Changes take effect only when the user clicks Apply. Cancel discards
|
||
all pending edits.
|
||
"""
|
||
|
||
def __init__(self, parent=None, *,
|
||
current_n: int,
|
||
current_threshold_mv: float,
|
||
pixel_x_mm: float | None):
|
||
super().__init__(parent)
|
||
self.setWindowTitle("Row-Averaged FFT Options")
|
||
self.setModal(True)
|
||
self.setMinimumWidth(380)
|
||
|
||
self._pixel_x_mm = pixel_x_mm
|
||
|
||
layout = QVBoxLayout(self)
|
||
|
||
# ---- Neighbor window ---------------------------------------------
|
||
grp_window = QGroupBox("Same-Row Neighbor Window")
|
||
wl = QVBoxLayout(grp_window)
|
||
|
||
n_row = QHBoxLayout()
|
||
n_row.addWidget(QLabel("Neighbor half-width (n):"))
|
||
self._spin_n = QSpinBox()
|
||
self._spin_n.setRange(1, 50)
|
||
self._spin_n.setValue(max(1, current_n))
|
||
self._spin_n.setToolTip(
|
||
"Each pixel's CH1 waveform is averaged with up to n same-row\n"
|
||
"neighbors on each side, distance-weighted (Gaussian) and\n"
|
||
"counting only neighbors that already pass the DC threshold\n"
|
||
"below. Never mixes across rows/Y.")
|
||
self._spin_n.valueChanged.connect(self._update_info)
|
||
n_row.addWidget(self._spin_n)
|
||
wl.addLayout(n_row)
|
||
|
||
self._lbl_width = QLabel()
|
||
self._lbl_width.setStyleSheet(_CSS_HINT)
|
||
wl.addWidget(self._lbl_width)
|
||
|
||
layout.addWidget(grp_window)
|
||
|
||
# ---- DC threshold ------------------------------------------------
|
||
grp_thr = QGroupBox("Neighbor Validity")
|
||
tl = QVBoxLayout(grp_thr)
|
||
thr_row = QHBoxLayout()
|
||
thr_row.addWidget(QLabel("DC threshold:"))
|
||
self._spin_threshold = _make_dspin(-500.0, 500.0, 3, suffix=" mV",
|
||
value=current_threshold_mv, step=0.025)
|
||
self._spin_threshold.setToolTip(
|
||
"A same-row neighbor only contributes to a pixel's average if\n"
|
||
"its own CH4 signal is at or above this threshold -- the same\n"
|
||
"test used for RF mask display. A pixel below threshold stays\n"
|
||
"masked, exactly as today; it is never rescued by its neighbors.")
|
||
thr_row.addWidget(self._spin_threshold)
|
||
tl.addLayout(thr_row)
|
||
layout.addWidget(grp_thr)
|
||
|
||
# ---- 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):
|
||
n = self._spin_n.value()
|
||
if self._pixel_x_mm is None:
|
||
self._lbl_width.setText("Load a file to preview the window's physical width.")
|
||
return
|
||
width_um = 2 * n * self._pixel_x_mm * 1e3
|
||
self._lbl_width.setText(
|
||
f"Window: ±{n} px = {width_um:.2f} µm full width "
|
||
f"(pixel pitch {self._pixel_x_mm * 1e3:.3g} µm)")
|
||
|
||
def get_half_width(self) -> int:
|
||
return self._spin_n.value()
|
||
|
||
def get_threshold_mv(self) -> float:
|
||
return self._spin_threshold.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 = AlignOverlayCanvas()
|
||
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)
|
||
|
||
|