Add manual angle alignment mode with overlay, keyboard nudge, and save/clear
Automatic alignment's phase-correlation translation search is unreliable, so add a Fusion -> Manual Alignment dialog: every angle's CH4 threshold mask overlaid at once with distinct colors/opacity, a selector for the active angle, arrow keys to nudge translation and Q/E to nudge rotation, an Auto De-rotate button that snaps to the known scan angle, and Save/Clear controls backed by a JSON sidecar that's restored automatically on reload. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
Thomas Ales
+657
-6
@@ -19,29 +19,34 @@ import faulthandler
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib as mpl
|
||||
import numpy as np
|
||||
from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg, NavigationToolbar2QT
|
||||
from matplotlib.figure import Figure
|
||||
from matplotlib.patches import Polygon
|
||||
from matplotlib.path import Path as MplPath
|
||||
from PyQt6.QtCore import QObject, Qt, QThread, pyqtSignal
|
||||
from PyQt6.QtGui import QAction
|
||||
from PyQt6.QtGui import QAction, QKeyEvent
|
||||
from PyQt6.QtWidgets import (
|
||||
QApplication, QButtonGroup, QCheckBox, QComboBox, QDialog, QDialogButtonBox,
|
||||
QDoubleSpinBox, QFileDialog, QFormLayout, QFrame, QGroupBox, QHBoxLayout,
|
||||
QLabel, QMainWindow, QProgressDialog, QPushButton, QRadioButton, QScrollArea,
|
||||
QSizePolicy, QSpinBox, QSplitter, QVBoxLayout, QWidget,
|
||||
QLabel, QMainWindow, QMessageBox, QProgressDialog, QPushButton, QRadioButton,
|
||||
QScrollArea, QSizePolicy, QSpinBox, QSplitter, QVBoxLayout, QWidget,
|
||||
)
|
||||
|
||||
import sras_compute as compute
|
||||
from sras_compute import PYFFTW_AVAILABLE, apply_alignment
|
||||
from sras_compute import (
|
||||
PYFFTW_AVAILABLE, ManualAngleParams, apply_alignment, build_manual_alignment,
|
||||
delete_manual_alignment, load_manual_alignment, save_manual_alignment,
|
||||
sidecar_path,
|
||||
)
|
||||
from sras_format import (
|
||||
CH1_IDX, CH3_IDX, CH4_IDX, CH_NAMES, SrasFile, adc_to_mv, mv_to_adc,
|
||||
_FALLBACK_YMULT_MV, _FALLBACK_YOFF_ADC,
|
||||
)
|
||||
from sras_workers import (
|
||||
AngleAlignmentWorker, BatchCacheWorker, ComputeWorker, DcPrecomputeWorker,
|
||||
LoadWorker,
|
||||
AngleAlignmentWorker, BatchCacheWorker, Ch4MaskWorker, ComputeWorker,
|
||||
DcPrecomputeWorker, LoadWorker,
|
||||
)
|
||||
|
||||
faulthandler.enable() # print a native stack trace on SIGSEGV/SIGABRT/etc.
|
||||
@@ -738,6 +743,542 @@ class FftOptionsDialog(QDialog):
|
||||
return max(1, self._spin_pad.value())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Manual alignment dialog (Fusion -> Manual Alignment...)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class ManualAlignOverlayCanvas(FigureCanvasQTAgg):
|
||||
"""Renders ManualAlignmentDialog's multi-angle mask overlay and turns
|
||||
keyboard input into translate/rotate nudge requests for whichever angle
|
||||
the dialog currently has active.
|
||||
|
||||
A pure input+render widget — it holds no alignment state and never
|
||||
touches SrasFile itself; ManualAlignmentDialog owns all of that and
|
||||
decides, from these signals, whether a cheap single-layer refresh or a
|
||||
full preview-canvas rebuild is needed.
|
||||
|
||||
FigureCanvasQTAgg is a real QWidget, so keyPressEvent works like on any
|
||||
other widget, but Qt only ever delivers key events to whichever widget
|
||||
currently has focus — StrongFocus, plus grabbing focus on click and once
|
||||
right after the dialog is shown, are both required or arrow keys
|
||||
silently do nothing.
|
||||
|
||||
Rotate keys are letters (Q/E), not punctuation (comma/period or
|
||||
brackets): Shift+letter still reports the same Qt.Key on every platform,
|
||||
whereas Shift+comma/bracket can report a different virtual key
|
||||
(Key_Less / Key_BraceLeft) depending on platform and keyboard layout —
|
||||
which would silently break the "Shift = coarse step" modifier for
|
||||
rotation specifically. Arrow keys have no such hazard.
|
||||
"""
|
||||
nudge_translate = pyqtSignal(int, int, bool) # dir_x, dir_y in {-1,0,1}; coarse
|
||||
nudge_rotate = pyqtSignal(int, bool) # dir in {-1,1} (CCW/CW); coarse
|
||||
|
||||
_TRANSLATE_KEYS = {
|
||||
Qt.Key.Key_Left: (-1, 0),
|
||||
Qt.Key.Key_Right: (1, 0),
|
||||
Qt.Key.Key_Up: (0, -1),
|
||||
Qt.Key.Key_Down: (0, 1),
|
||||
}
|
||||
_ROTATE_KEYS = {Qt.Key.Key_Q: 1, Qt.Key.Key_E: -1} # CCW, CW
|
||||
|
||||
def __init__(self, parent=None):
|
||||
fig = Figure(figsize=(6, 6), tight_layout=True)
|
||||
self.ax = fig.add_subplot(111)
|
||||
super().__init__(fig)
|
||||
self.setParent(parent)
|
||||
self.setFocusPolicy(Qt.FocusPolicy.StrongFocus)
|
||||
self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
|
||||
self.mpl_connect("button_press_event", lambda _e: self.setFocus())
|
||||
|
||||
def show_overlay(self, rgba: np.ndarray, extent: list[float], title: str):
|
||||
self.figure.clf()
|
||||
self.ax = self.figure.add_subplot(111)
|
||||
self.ax.imshow(rgba, extent=extent, origin="upper", aspect="auto")
|
||||
self.ax.set_xlabel("X (mm)")
|
||||
self.ax.set_ylabel("Y (mm)")
|
||||
self.ax.set_title(title)
|
||||
self.draw_idle() # coalesces rapid redraws — matters for key-repeat.
|
||||
|
||||
def keyPressEvent(self, event: QKeyEvent):
|
||||
key = event.key()
|
||||
coarse = bool(event.modifiers() & Qt.KeyboardModifier.ShiftModifier)
|
||||
if key in self._TRANSLATE_KEYS:
|
||||
dx, dy = self._TRANSLATE_KEYS[key]
|
||||
self.nudge_translate.emit(dx, dy, coarse)
|
||||
event.accept()
|
||||
elif key in self._ROTATE_KEYS:
|
||||
self.nudge_rotate.emit(self._ROTATE_KEYS[key], coarse)
|
||||
event.accept()
|
||||
else:
|
||||
super().keyPressEvent(event)
|
||||
|
||||
|
||||
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 — the thing the automatic phase-
|
||||
correlation step (compute_angle_alignment) sometimes gets wrong. The user
|
||||
picks an "active" angle and nudges its rotation+translation with the
|
||||
keyboard; Auto De-rotate sets every non-reference angle's rotation to the
|
||||
known, analytic scan-angle delta without touching any translation the
|
||||
user has already dialed in. 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 own (rare) background mask-fetch step, so the main window's
|
||||
existing shutdown/lifecycle plumbing covers it 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, 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
|
||||
|
||||
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_factor = 1
|
||||
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_dx_mm = self._preview_dy_mm = 1.0
|
||||
self._masks_ready = 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)
|
||||
|
||||
# ---- 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, 300))
|
||||
|
||||
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_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
|
||||
max_dim = max(max(img.shape) for img in self._dc4_mv.values())
|
||||
self._downsample_factor = max(1, int(np.ceil(max_dim / 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."""
|
||||
threshold = self.spin_mask_threshold_mv.value()
|
||||
factor = self._downsample_factor
|
||||
self._masks_small = {
|
||||
a: compute._block_mean_downsample(
|
||||
(img >= threshold).astype(np.float32), factor)
|
||||
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)
|
||||
factor = self._downsample_factor
|
||||
dx_c, dy_c = dx_ref * factor, dy_ref * factor
|
||||
origin, shape = compute.union_canvas_mm(
|
||||
self._sras, self._ref_angle_idx, dx_c, dy_c, self._angle_params,
|
||||
margin_frac=self._PREVIEW_MARGIN_FRAC)
|
||||
self._preview_origin_mm, self._preview_shape = origin, shape
|
||||
self._preview_dx_mm, self._preview_dy_mm = dx_c, dy_c
|
||||
self._preview_layers = {
|
||||
a: compute.reproject_mask(
|
||||
self._sras, a, self._ref_angle_idx, self._masks_small[a],
|
||||
self._angle_params[a].rotation_deg, self._angle_params[a].shift_mm,
|
||||
dx_c, dy_c, origin, shape)
|
||||
for a in range(self._sras.n_angles)
|
||||
}
|
||||
self._redraw_overlay()
|
||||
|
||||
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."""
|
||||
a = self._active_angle
|
||||
self._preview_layers[a] = compute.reproject_mask(
|
||||
self._sras, a, self._ref_angle_idx, self._masks_small[a],
|
||||
self._angle_params[a].rotation_deg, self._angle_params[a].shift_mm,
|
||||
self._preview_dx_mm, self._preview_dy_mm,
|
||||
self._preview_origin_mm, self._preview_shape)
|
||||
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_dx_mm, self._preview_dy_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):
|
||||
n_changed = 0
|
||||
for a in range(self._sras.n_angles):
|
||||
if a == self._ref_angle_idx:
|
||||
continue
|
||||
self._angle_params[a].rotation_deg = compute._theta_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 known scan angle for {n_changed} angle(s) "
|
||||
"(translation left untouched).")
|
||||
|
||||
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._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.btn_auto_derotate.setEnabled(enabled)
|
||||
self.btn_save.setEnabled(enabled)
|
||||
self.btn_clear.setEnabled(enabled)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main window
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -788,6 +1329,7 @@ class SrasViewerWindow(QMainWindow):
|
||||
self._alignment_result = None
|
||||
self._alignment_generation: int = 0
|
||||
self._aligned_cache: dict[tuple, np.ndarray] = {}
|
||||
self._manual_align_dialog: ManualAlignmentDialog | None = None
|
||||
|
||||
self._build_ui()
|
||||
|
||||
@@ -1137,6 +1679,16 @@ class SrasViewerWindow(QMainWindow):
|
||||
self._alignment_act.triggered.connect(self._on_angle_alignment)
|
||||
fusion_menu.addAction(self._alignment_act)
|
||||
|
||||
self._manual_align_act = QAction("&Manual Alignment…", self)
|
||||
self._manual_align_act.setStatusTip(
|
||||
"Open an interactive dialog to align angles by eye: overlaid CH4 "
|
||||
"threshold masks, keyboard nudge (translate + rotate), auto "
|
||||
"de-rotate to the known scan angles, and save/clear a persistent "
|
||||
"alignment.")
|
||||
self._manual_align_act.setEnabled(False)
|
||||
self._manual_align_act.triggered.connect(self._on_manual_alignment)
|
||||
fusion_menu.addAction(self._manual_align_act)
|
||||
|
||||
convert_menu = menubar.addMenu("&Convert")
|
||||
self._batch_dc_act = QAction("Batch Compute DC and &Store…", self)
|
||||
self._batch_dc_act.setStatusTip(
|
||||
@@ -1196,6 +1748,13 @@ class SrasViewerWindow(QMainWindow):
|
||||
self._sras = sras
|
||||
self._current_image = None
|
||||
|
||||
# A manual-alignment dialog bound to the previous file must not
|
||||
# survive a reload — its per-angle state (and the sras it was
|
||||
# constructed against) no longer matches the new file's geometry.
|
||||
if self._manual_align_dialog is not None:
|
||||
self._manual_align_dialog.close()
|
||||
self._manual_align_dialog = None
|
||||
|
||||
# Caches (and any in-flight DC precompute) belong to the previous
|
||||
# file's geometry — discard and start fresh. Bumping the generation
|
||||
# counters makes any still-running worker's result get dropped when
|
||||
@@ -1213,6 +1772,26 @@ class SrasViewerWindow(QMainWindow):
|
||||
self.chk_aligned_view.setEnabled(False)
|
||||
self.chk_aligned_view.blockSignals(False)
|
||||
|
||||
# Silently restore a previously-saved manual alignment, if any, so
|
||||
# the work survives closing and reopening the file.
|
||||
sidecar = load_manual_alignment(sras)
|
||||
if sidecar is not None:
|
||||
try:
|
||||
self._alignment_result = build_manual_alignment(
|
||||
sras, sidecar.ref_angle_idx, sidecar.dc_threshold_mv,
|
||||
sidecar.per_angle)
|
||||
self.chk_aligned_view.blockSignals(True)
|
||||
self.chk_aligned_view.setChecked(True)
|
||||
self.chk_aligned_view.blockSignals(False)
|
||||
self.statusBar().showMessage(
|
||||
f"Restored saved manual alignment from "
|
||||
f"{sidecar_path(sras.path).name}")
|
||||
except Exception as exc:
|
||||
# A corrupt/foreign sidecar or a rescan that shrank n_angles
|
||||
# below ref_angle_idx must not block opening the .sras file.
|
||||
self.statusBar().showMessage(
|
||||
f"Could not restore saved alignment: {exc}")
|
||||
|
||||
# A ROI from the previous file no longer matches the new scan's
|
||||
# geometry, so discard it on every load.
|
||||
self.image_canvas.clear_roi()
|
||||
@@ -1318,6 +1897,8 @@ class SrasViewerWindow(QMainWindow):
|
||||
|
||||
self._alignment_act.setEnabled(
|
||||
has_file and s.n_angles > 1 and not self._job_running("align"))
|
||||
self._manual_align_act.setEnabled(
|
||||
has_file and s.n_angles > 1 and not self._job_running("align"))
|
||||
self.chk_aligned_view.setEnabled(enabled and self._alignment_result is not None)
|
||||
self._update_roi_ui()
|
||||
|
||||
@@ -1934,6 +2515,73 @@ class SrasViewerWindow(QMainWindow):
|
||||
f"canvas {nc}×{nr} px).")
|
||||
self._refresh_display()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Fusion: manual alignment
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _on_manual_alignment(self):
|
||||
if self._sras is None or self._sras.n_angles <= 1:
|
||||
return
|
||||
if self._manual_align_dialog is not None:
|
||||
self._manual_align_dialog.raise_()
|
||||
self._manual_align_dialog.activateWindow()
|
||||
return
|
||||
|
||||
ref_idx = 0
|
||||
threshold_mv = self.spin_threshold_mv.value()
|
||||
seed: dict[int, ManualAngleParams] = {}
|
||||
if (self._alignment_result is not None
|
||||
and self._alignment_result.ref_angle_idx == ref_idx):
|
||||
seed = {a: ManualAngleParams(t.rotation_deg, t.shift_mm)
|
||||
for a, t in self._alignment_result.per_angle.items()}
|
||||
threshold_mv = self._alignment_result.dc_threshold_mv
|
||||
else:
|
||||
sidecar = load_manual_alignment(self._sras)
|
||||
if sidecar is not None and sidecar.ref_angle_idx == ref_idx:
|
||||
seed = dict(sidecar.per_angle)
|
||||
threshold_mv = sidecar.dc_threshold_mv
|
||||
|
||||
cached_dc4 = {a: img for (a, ch), img in self._dc_cache.items() if ch == CH4_IDX}
|
||||
dlg = ManualAlignmentDialog(
|
||||
self, self._sras, ref_angle_idx=ref_idx, dc_threshold_mv=threshold_mv,
|
||||
seed_per_angle=seed, cached_dc4_mv=cached_dc4)
|
||||
dlg.alignment_saved.connect(self._on_manual_alignment_saved)
|
||||
dlg.alignment_cleared.connect(self._on_manual_alignment_cleared)
|
||||
dlg.finished.connect(self._on_manual_align_dialog_closed)
|
||||
dlg.setAttribute(Qt.WidgetAttribute.WA_DeleteOnClose)
|
||||
self._manual_align_dialog = dlg
|
||||
dlg.show()
|
||||
|
||||
def _on_manual_align_dialog_closed(self, _result_code: int):
|
||||
self._manual_align_dialog = None
|
||||
|
||||
def _on_manual_alignment_saved(self, result, sidecar_path_str: str):
|
||||
self._alignment_result = result
|
||||
self._aligned_cache = {}
|
||||
self._alignment_generation += 1
|
||||
self.chk_aligned_view.setEnabled(True)
|
||||
self.chk_aligned_view.blockSignals(True)
|
||||
self.chk_aligned_view.setChecked(True)
|
||||
self.chk_aligned_view.blockSignals(False)
|
||||
self._update_controls_enabled(self._sras is not None)
|
||||
self.statusBar().showMessage(
|
||||
f"Manual alignment saved to {Path(sidecar_path_str).name}")
|
||||
if self._current_image is not None:
|
||||
self._refresh_display()
|
||||
|
||||
def _on_manual_alignment_cleared(self):
|
||||
self._alignment_result = None
|
||||
self._aligned_cache = {}
|
||||
self._alignment_generation += 1
|
||||
self.chk_aligned_view.blockSignals(True)
|
||||
self.chk_aligned_view.setChecked(False)
|
||||
self.chk_aligned_view.setEnabled(False)
|
||||
self.chk_aligned_view.blockSignals(False)
|
||||
self._update_controls_enabled(self._sras is not None)
|
||||
self.statusBar().showMessage("Manual alignment cleared.")
|
||||
if self._current_image is not None:
|
||||
self._refresh_display()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# FFT Options
|
||||
# ------------------------------------------------------------------
|
||||
@@ -1960,6 +2608,9 @@ class SrasViewerWindow(QMainWindow):
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def closeEvent(self, event):
|
||||
if self._manual_align_dialog is not None:
|
||||
self._manual_align_dialog.close()
|
||||
|
||||
# Signal every cancellable worker first, then wait. Waiting without
|
||||
# signalling means sitting out whatever is in flight — on a large
|
||||
# scan a single angle is ~40 s.
|
||||
|
||||
Reference in New Issue
Block a user