Thomas Ales
f40c965b74
Alignment was two disconnected menu actions. `Angle Alignment` ran the
registration with ref_angle_idx hard-coded to 0, no exposed parameters and
no way to retry; `Manual Alignment...` was a separate dialog that
deliberately refused to inherit the automatic result, so a bad fit meant
starting over by hand. Neither told you whether the alignment was any
good, and neither produced anything durable beyond an in-memory result.
Fusion -> Alignment Wizard... now covers all of it in three steps:
1. Correlate. Reference angle, DC threshold, source, rotation seed and
sign, search window, coarse step and fine grid are all on the page,
and Run/Re-run is repeatable. Angles start pre-rotated from the
stage angles in the file, so the page is informative before any
correlation runs. The verdict is a picture: every angle's DC mask
reprojected onto the shared canvas and summed, coloured by how many
angles cover each pixel, so a good alignment reads as one saturated
plateau and a bad one as a fringe of low-count halos. A per-angle
fit table flags angles that did not register or that disagree with
their stage angle by more than a degree.
2. Crop. An axis-aligned rectangle on that canvas, with numeric
canvas-pixel boxes synced both ways and a live size estimate.
"Fit to full overlap" uses a largest-rectangle sweep rather than a
bounding box: the overlap region of several rotated scans is roughly
a disc, whose bounding box has corners no angle covers.
3. Save. Writes the aligned, cropped stack to a new .sras.
Because the wizard replaces both actions it absorbs the old dialog's
by-eye nudge editor — without it, a scan the search cannot fit would
have no fallback at all. ManualAlignmentDialog is therefore deleted
rather than left orphaned, and its tests move to the wizard.
Two bugs found while driving it end to end and fixed here: QSpinBox
setRange clamps and emits valueChanged, which committed a 1x1 crop
before the default preset could run; and the mm round trip returns an
exact pixel boundary as 11.000000000000002, so a bare ceil() added a
spurious column on every rectangle edit.
Verified: 103 pass, the 6 test_stored_cache failures are pre-existing on
main, and tools/check_equivalence.py is byte-identical to main.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
1412 lines
58 KiB
Python
1412 lines
58 KiB
Python
"""The alignment wizard: correlate, crop, export.
|
||
|
||
Replaces what used to be two disconnected menu actions. `Angle Alignment` ran
|
||
the registration with every parameter hard-coded and no way to retry, and
|
||
`Manual Alignment…` was a separate dialog that deliberately refused to inherit
|
||
the automatic result — so a bad fit meant starting over by hand, and neither
|
||
path told you whether the alignment was actually any good.
|
||
|
||
The three steps are the three decisions:
|
||
|
||
1. Correlate. Every parameter that affects the fit is on the page, the run is
|
||
repeatable, and the verdict is a picture: each angle's DC mask reprojected
|
||
onto the shared canvas and summed, so the colour of a pixel is *how many
|
||
angles cover it*. A good alignment is one saturated plateau at N; a bad one
|
||
is a fan of low-count halos. Angles start pre-rotated by the stage angles
|
||
stored in the file, so the page is informative before any correlation runs.
|
||
Manual nudging lives here too, since this page is now the only place to
|
||
correct an angle the search cannot fit.
|
||
2. Crop. Axis-aligned only, because that is the only crop .sras geometry can
|
||
express — a free quadrilateral would have to be squared off behind the
|
||
user's back.
|
||
3. Export. Writes the aligned, cropped data to a new .sras.
|
||
|
||
State lives on the wizard rather than in QWizardPage.registerField: the pages
|
||
share numpy arrays, ManualAngleParams and an AlignmentResult, none of which are
|
||
scalar widget properties.
|
||
"""
|
||
|
||
from dataclasses import dataclass, field
|
||
from pathlib import Path
|
||
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, Qt, pyqtSignal
|
||
from PyQt6.QtWidgets import (
|
||
QCheckBox, QComboBox, QFileDialog, QHBoxLayout, QHeaderView, QLabel,
|
||
QLineEdit, QMessageBox, QProgressBar, QPushButton, QSpinBox,
|
||
QTableWidget, QTableWidgetItem, QVBoxLayout, QWidget, QWizard, QWizardPage,
|
||
)
|
||
|
||
import sras_align_export as export
|
||
import sras_compute as compute
|
||
from sras_compute import ManualAngleParams, build_manual_alignment
|
||
from sras_format import SrasFile
|
||
from sras_workers import AlignedExportWorker, Ch4MaskWorker, CrossCorrelateWorker
|
||
|
||
from .canvases import AlignOverlayCanvas, ImageCanvas, RoiQuad, count_colormap
|
||
from .common import (
|
||
_CSS_HINT, _CSS_INFO, _CSS_MUTED, _CSS_WARN, Jobs, _axes_extent, _form,
|
||
_group, _make_dspin, _scroll_panel, _wrap_label,
|
||
)
|
||
|
||
if TYPE_CHECKING:
|
||
from .main_window import SrasViewerWindow
|
||
|
||
# Longest side of the coarsened preview the mask stack is built on. Same
|
||
# reasoning as the old manual dialog: a full-resolution reprojection per angle
|
||
# per edit is far more detail than an alignment can be judged by eye at.
|
||
_MAX_PREVIEW_DIM = 1024
|
||
|
||
# Alpha for the per-angle colour view (the active angle sits on top, brighter).
|
||
_BASE_ALPHA = 0.42
|
||
_ACTIVE_ALPHA = 0.75
|
||
|
||
_PANEL_W = 372
|
||
|
||
# (label, sources for register_angle_to_reference) — "both" costs roughly double
|
||
# but removes the failure mode where the one chosen source happens to be
|
||
# uninformative for a single angle.
|
||
_CORRELATE_SOURCES = (
|
||
("Both, keep best", ("signal", "mask")),
|
||
("Raw signal", ("signal",)),
|
||
("Thresholded mask", ("mask",)),
|
||
)
|
||
|
||
# (label, seed_deg override or None for "the file's stage angle", signs to
|
||
# sweep, sign used for the pre-rotation *preview*). The stage's rotational sense
|
||
# relative to this module's math-positive convention is not knowable from the
|
||
# file, which is why "both" exists and is the default; the explicit signs are
|
||
# for a user who has established which way their own stage turns.
|
||
_PREROTATE_MODES = (
|
||
("Stage angle, both signs", None, (-1, 1), 1),
|
||
("Stage angle, + delta only", None, (1,), 1),
|
||
("Stage angle, − delta only", None, (-1,), -1),
|
||
("No pre-rotation", 0.0, (1,), 0),
|
||
)
|
||
|
||
|
||
@dataclass
|
||
class WizardState:
|
||
"""Everything the three pages share."""
|
||
dc4_mv: dict[int, np.ndarray] = field(default_factory=dict)
|
||
masks_small: dict[int, np.ndarray] = field(default_factory=dict)
|
||
downsample: tuple[int, int] = (1, 1)
|
||
params: dict[int, ManualAngleParams] = field(default_factory=dict)
|
||
fits: dict[int, tuple[float, str]] = field(default_factory=dict)
|
||
ref_angle_idx: int = 0
|
||
threshold_mv: float = 0.0
|
||
result: compute.AlignmentResult | None = None # full, uncropped
|
||
counts: np.ndarray | None = None # coarse overlap counts
|
||
layers: dict[int, np.ndarray] = field(default_factory=dict)
|
||
preview_pitch_mm: tuple[float, float] = (1.0, 1.0)
|
||
preview_shape: tuple[int, int] = (1, 1)
|
||
geometry_generation: int = 0
|
||
crop: tuple[int, int, int, int] | None = None # canvas (row0,col0,nr,nc)
|
||
cropped_result: compute.AlignmentResult | None = None
|
||
out_path: str = ""
|
||
exported_path: str = ""
|
||
|
||
|
||
class AlignmentWizard(QWizard):
|
||
"""Three-page alignment + export flow (Fusion → Alignment Wizard…).
|
||
|
||
Non-modal by design, like the dialog it replaces: an export can take minutes
|
||
on a real scan and the user should still be able to look at the data.
|
||
|
||
Background work goes through parent._run_worker, which inherits the main
|
||
window's job registry, thread joining and shutdown handling. Two rules that
|
||
module's docstring establishes and every launch here follows: disable the
|
||
trigger *before* calling it (a re-entrant click must not be able to start a
|
||
second thread over the first), and never ignore its return value — False
|
||
means another job holds the key.
|
||
"""
|
||
|
||
PAGE_CORRELATE, PAGE_ROI, PAGE_SAVE = 0, 1, 2
|
||
|
||
# cropped AlignmentResult, written path
|
||
alignment_ready = pyqtSignal(object, str)
|
||
|
||
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._cached_dc4 = dict(cached_dc4_mv)
|
||
self._seed = dict(seed_per_angle or {})
|
||
self._closing = False
|
||
|
||
self.state = WizardState(
|
||
ref_angle_idx=ref_angle_idx, threshold_mv=dc_threshold_mv,
|
||
params={a: ManualAngleParams() for a in range(sras.n_angles)})
|
||
|
||
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.setWindowTitle(f"Alignment Wizard — {sras.path.name}")
|
||
self.setWizardStyle(QWizard.WizardStyle.ModernStyle)
|
||
self.setOption(QWizard.WizardOption.HaveHelpButton, False)
|
||
self.setOption(QWizard.WizardOption.NoBackButtonOnStartPage, True)
|
||
# IndependentPages must stay OFF: it suppresses cleanupPage, which is
|
||
# how the ROI page discards a crop whose canvas is about to change.
|
||
self.setOption(QWizard.WizardOption.IndependentPages, False)
|
||
self.resize(1220, 820)
|
||
|
||
self.setPage(self.PAGE_CORRELATE, CorrelatePage(self))
|
||
self.setPage(self.PAGE_ROI, RoiPage(self))
|
||
self.setPage(self.PAGE_SAVE, SavePage(self))
|
||
self.setStartId(self.PAGE_CORRELATE)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Shared helpers the pages use
|
||
# ------------------------------------------------------------------
|
||
|
||
@property
|
||
def sras(self) -> SrasFile:
|
||
return self._sras
|
||
|
||
def seed_params(self) -> dict[int, ManualAngleParams]:
|
||
"""Per-angle starting parameters: a saved sidecar if one matches,
|
||
otherwise identity (the pre-rotation is applied on top by page 1)."""
|
||
return {a: ManualAngleParams(self._seed[a].rotation_deg,
|
||
self._seed[a].shift_mm)
|
||
if a in self._seed else ManualAngleParams()
|
||
for a in range(self._sras.n_angles)}
|
||
|
||
def rebuild_result(self):
|
||
"""Recompute the full AlignmentResult from the current parameters.
|
||
|
||
Closed-form matrix and bbox maths with no per-pixel work, so it is cheap
|
||
enough to call on every edit — which is what keeps the preview honest
|
||
about the canvas a rotation change implies.
|
||
"""
|
||
st = self.state
|
||
st.result = build_manual_alignment(
|
||
self._sras, st.ref_angle_idx, st.threshold_mv, st.params)
|
||
|
||
def rebuild_stack(self):
|
||
"""Reproject every angle's mask onto a coarsened view of the *final*
|
||
canvas and sum them into an overlap-count image.
|
||
|
||
The preview deliberately shares the final canvas's origin and uses a
|
||
pitch that is an integer multiple of it, rather than the padded,
|
||
unsnapped preview canvas the old manual dialog used. That is what lets
|
||
the crop page turn a rectangle drawn in millimetres into an exact
|
||
integer window of the real canvas, with no second coordinate frame to
|
||
reconcile.
|
||
"""
|
||
st = self.state
|
||
if st.result is None or not st.masks_small:
|
||
return
|
||
fy, fx = st.downsample
|
||
n_rows, n_cols = st.result.canvas_shape
|
||
st.preview_pitch_mm = (st.result.canvas_dx_mm * fx,
|
||
st.result.canvas_dy_mm * fy)
|
||
st.preview_shape = (max(1, -(-n_rows // fy)), max(1, -(-n_cols // fx)))
|
||
|
||
st.layers = {}
|
||
counts = np.zeros(st.preview_shape, dtype=np.int16)
|
||
for a in range(self._sras.n_angles):
|
||
p = st.params[a]
|
||
layer = compute.reproject_mask(
|
||
self._sras, a, st.ref_angle_idx, st.masks_small[a],
|
||
p.rotation_deg, p.shift_mm, st.preview_pitch_mm,
|
||
st.result.canvas_origin_mm, st.preview_shape,
|
||
src_downsample=st.downsample)
|
||
st.layers[a] = layer
|
||
counts += (layer > 0.5).astype(np.int16)
|
||
st.counts = counts
|
||
|
||
def preview_extent(self) -> list[float]:
|
||
st = self.state
|
||
x0, y0 = st.result.canvas_origin_mm
|
||
dx, dy = st.preview_pitch_mm
|
||
n_rows, n_cols = st.preview_shape
|
||
return _axes_extent(x0 + np.arange(n_cols) * dx,
|
||
y0 + np.arange(n_rows) * dy, dx, dy)
|
||
|
||
def mm_to_canvas(self, x_mm: float, y_mm: float) -> tuple[float, float]:
|
||
"""(col, row) in *final* canvas pixels, fractional."""
|
||
st = self.state
|
||
x0, y0 = st.result.canvas_origin_mm
|
||
return ((x_mm - x0) / st.result.canvas_dx_mm,
|
||
(y_mm - y0) / st.result.canvas_dy_mm)
|
||
|
||
def canvas_to_mm(self, col: float, row: float) -> tuple[float, float]:
|
||
st = self.state
|
||
x0, y0 = st.result.canvas_origin_mm
|
||
return (x0 + col * st.result.canvas_dx_mm,
|
||
y0 + row * st.result.canvas_dy_mm)
|
||
|
||
def job_running(self) -> bool:
|
||
return any(self._parent._job_running(k) for k in
|
||
(Jobs.ALIGN_MASKS, Jobs.ALIGN_CORRELATE, Jobs.ALIGN_EXPORT))
|
||
|
||
# ------------------------------------------------------------------
|
||
# Lifetime
|
||
# ------------------------------------------------------------------
|
||
|
||
def reject(self):
|
||
"""Refuse to close while a job is in flight.
|
||
|
||
The running worker's signals are connected to bound methods of this
|
||
wizard and its pages; letting Qt delete them under a live thread is a
|
||
crash, not an inconvenience. Ask the worker to stop and let the job's
|
||
own completion close us.
|
||
"""
|
||
if self.job_running():
|
||
for page_id in (self.PAGE_CORRELATE, self.PAGE_SAVE):
|
||
page = self.page(page_id)
|
||
if page is not None:
|
||
page.request_stop()
|
||
self._closing = True
|
||
return
|
||
super().reject()
|
||
|
||
def closeEvent(self, event):
|
||
"""Window-manager close goes through the same deferral as Cancel."""
|
||
if self.job_running():
|
||
self.reject()
|
||
event.ignore()
|
||
return
|
||
super().closeEvent(event)
|
||
|
||
def accept(self):
|
||
st = self.state
|
||
self.alignment_ready.emit(st.cropped_result, st.exported_path)
|
||
super().accept()
|
||
|
||
def maybe_close_after_job(self):
|
||
"""Called by a page when its job finishes; completes a deferred close."""
|
||
if self._closing and not self.job_running():
|
||
self._closing = False
|
||
super().reject()
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Page 1 — cross-correlation
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class CorrelatePage(QWizardPage):
|
||
"""Register every angle against the reference, and show whether it worked.
|
||
|
||
The stack image is the point of the page. Numbers alone ("score 0.42") do
|
||
not tell you whether a five-angle fusion is usable; an overlap-count image
|
||
does, immediately.
|
||
"""
|
||
|
||
def __init__(self, wizard: AlignmentWizard):
|
||
super().__init__(wizard)
|
||
self._wiz = wizard
|
||
self._sras = wizard.sras
|
||
self._busy = False
|
||
self._masks_ready = False
|
||
self._active_angle = 0
|
||
self._worker = None
|
||
self._done_count = 0
|
||
self._total = 0
|
||
|
||
self.setTitle("Step 1 — Cross-correlate the angles")
|
||
self.setSubTitle(
|
||
"Angles start pre-rotated by the stage angles stored in the scan. "
|
||
"Run the correlation, then check the stack: every pixel is coloured "
|
||
"by how many angles cover it, so a good alignment is one solid "
|
||
"plateau.")
|
||
self._build_ui()
|
||
|
||
# ---- construction -------------------------------------------------
|
||
|
||
def _build_ui(self):
|
||
root = QHBoxLayout(self)
|
||
|
||
self.canvas = AlignOverlayCanvas()
|
||
left = QWidget()
|
||
ll = QVBoxLayout(left)
|
||
ll.setContentsMargins(0, 0, 0, 0)
|
||
ll.setSpacing(4)
|
||
ll.addWidget(NavigationToolbar2QT(self.canvas, left))
|
||
ll.addWidget(self.canvas, stretch=1)
|
||
self.lbl_overlap = _wrap_label("", _CSS_INFO)
|
||
ll.addWidget(self.lbl_overlap)
|
||
self.lbl_warn = _wrap_label("", _CSS_WARN)
|
||
ll.addWidget(self.lbl_warn)
|
||
root.addWidget(left, stretch=1)
|
||
|
||
panel = QWidget()
|
||
pl = QVBoxLayout(panel)
|
||
pl.setContentsMargins(0, 0, 0, 0)
|
||
pl.setSpacing(8)
|
||
pl.addWidget(self._build_view_group())
|
||
pl.addWidget(self._build_reg_group())
|
||
pl.addWidget(self._build_run_group())
|
||
pl.addWidget(self._build_nudge_group())
|
||
pl.addWidget(self._build_table_group())
|
||
self.lbl_status = _wrap_label("", _CSS_MUTED)
|
||
pl.addWidget(self.lbl_status)
|
||
pl.addStretch()
|
||
root.addWidget(_scroll_panel(panel, _PANEL_W))
|
||
|
||
self._connect()
|
||
|
||
def _build_view_group(self) -> QWidget:
|
||
grp, lay = _group("View")
|
||
self.combo_view = _combo(["Overlap count", "Per-angle colours"])
|
||
lay.addWidget(self.combo_view)
|
||
return grp
|
||
|
||
def _build_reg_group(self) -> QWidget:
|
||
grp, lay = _group("Registration")
|
||
form = _form()
|
||
|
||
self.combo_ref = _combo(
|
||
f"Angle {a} ({self._sras.angles_deg[a]:.1f}°)"
|
||
for a in range(self._sras.n_angles))
|
||
self.combo_ref.setCurrentIndex(self._wiz.state.ref_angle_idx)
|
||
form.addRow("Reference:", self.combo_ref)
|
||
|
||
self.spin_threshold = _make_dspin(-500.0, 500.0, 3, suffix=" mV",
|
||
value=self._wiz.state.threshold_mv)
|
||
form.addRow("DC threshold:", self.spin_threshold)
|
||
|
||
self.combo_source = _combo(
|
||
(label, sources) for label, sources in _CORRELATE_SOURCES)
|
||
form.addRow("Correlate on:", self.combo_source)
|
||
|
||
self.combo_prerotate = _combo(
|
||
(label, (seed, signs, disp))
|
||
for label, seed, signs, disp in _PREROTATE_MODES)
|
||
form.addRow("Pre-rotate:", self.combo_prerotate)
|
||
|
||
self.chk_lock_rotation = QCheckBox("Lock rotation to the seed")
|
||
form.addRow("", self.chk_lock_rotation)
|
||
|
||
self.spin_search_deg = _make_dspin(0.0, 180.0, 1, suffix=" °",
|
||
value=6.0, step=1.0)
|
||
form.addRow("Rotation search (±):", self.spin_search_deg)
|
||
|
||
self.spin_coarse_step = _make_dspin(0.1, 30.0, 2, suffix=" °", value=2.0)
|
||
form.addRow("Coarse step:", self.spin_coarse_step)
|
||
|
||
self.combo_fine_dim = _combo((f"{dim} px", dim) for dim in (320, 640, 1024))
|
||
self.combo_fine_dim.setCurrentIndex(1)
|
||
form.addRow("Fine grid:", self.combo_fine_dim)
|
||
|
||
lay.addLayout(form)
|
||
lay.addWidget(_wrap_label(
|
||
"Pre-rotation only seeds the search — the rotation is still found "
|
||
"from image content, and both signs of the stage's reported angle "
|
||
"are tried unless you narrow it. Locking instead pins rotation to "
|
||
"the stage angle and searches translation only.", _CSS_HINT))
|
||
return grp
|
||
|
||
def _build_run_group(self) -> QWidget:
|
||
grp, lay = _group("Run")
|
||
self.btn_correlate = QPushButton("Run / Re-run Correlation")
|
||
lay.addWidget(self.btn_correlate)
|
||
self.progress = QProgressBar()
|
||
self.progress.setRange(0, 100)
|
||
self.progress.setValue(0)
|
||
lay.addWidget(self.progress)
|
||
self.btn_reset = QPushButton("Reset to pre-rotation only")
|
||
lay.addWidget(self.btn_reset)
|
||
return grp
|
||
|
||
def _build_nudge_group(self) -> QWidget:
|
||
self.grp_nudge, lay = _group("Manual Correction")
|
||
form = _form()
|
||
self.combo_active = _combo(
|
||
f"Angle {a} ({self._sras.angles_deg[a]:.1f}°)"
|
||
for a in range(self._sras.n_angles))
|
||
form.addRow("Active angle:", self.combo_active)
|
||
|
||
self.spin_rot = _make_dspin(-3600.0, 3600.0, 3, suffix=" °")
|
||
form.addRow("Rotation:", self.spin_rot)
|
||
self.spin_shift_x = _make_dspin(-1e5, 1e5, 4, suffix=" mm")
|
||
form.addRow("Shift X:", self.spin_shift_x)
|
||
self.spin_shift_y = _make_dspin(-1e5, 1e5, 4, suffix=" mm")
|
||
form.addRow("Shift Y:", self.spin_shift_y)
|
||
|
||
self.spin_step_translate = _make_dspin(0.0001, 1000.0, 4, suffix=" mm",
|
||
value=0.01)
|
||
form.addRow("Translate step:", self.spin_step_translate)
|
||
self.spin_step_rotate = _make_dspin(0.001, 90.0, 3, suffix=" °", value=0.1)
|
||
form.addRow("Rotate step:", self.spin_step_rotate)
|
||
self.spin_step_mult = _make_dspin(1.0, 1000.0, 1, value=10.0)
|
||
form.addRow("Coarse × (Shift):", self.spin_step_mult)
|
||
lay.addLayout(form)
|
||
lay.addWidget(_wrap_label(
|
||
"Arrow keys nudge translation; Q/E nudge rotation (CCW/CW). Hold "
|
||
"Shift for the coarse step. Click the image once to give it "
|
||
"keyboard focus. Switch to per-angle colours to see which layer "
|
||
"you are moving.", _CSS_HINT))
|
||
self.lbl_active_note = _wrap_label("", _CSS_WARN)
|
||
lay.addWidget(self.lbl_active_note)
|
||
return self.grp_nudge
|
||
|
||
def _build_table_group(self) -> QWidget:
|
||
grp, lay = _group("Fit per Angle")
|
||
self.table = QTableWidget(self._sras.n_angles, 5)
|
||
self.table.setHorizontalHeaderLabels(
|
||
["Angle", "Rot °", "Δ stage °", "Score", "On"])
|
||
self.table.verticalHeader().setVisible(False)
|
||
self.table.setEditTriggers(QTableWidget.EditTrigger.NoEditTriggers)
|
||
self.table.setSelectionMode(QTableWidget.SelectionMode.NoSelection)
|
||
self.table.horizontalHeader().setSectionResizeMode(
|
||
QHeaderView.ResizeMode.ResizeToContents)
|
||
self.table.setMinimumHeight(150)
|
||
lay.addWidget(self.table)
|
||
return grp
|
||
|
||
def _connect(self):
|
||
self.combo_view.currentIndexChanged.connect(self._redraw)
|
||
self.combo_ref.currentIndexChanged.connect(self._on_ref_changed)
|
||
self.spin_threshold.editingFinished.connect(self._on_threshold_changed)
|
||
self.combo_prerotate.currentIndexChanged.connect(self._apply_prerotation)
|
||
self.chk_lock_rotation.toggled.connect(self._on_lock_toggled)
|
||
self.btn_correlate.clicked.connect(self._on_correlate)
|
||
self.btn_reset.clicked.connect(self._apply_prerotation)
|
||
self.combo_active.currentIndexChanged.connect(self._on_active_changed)
|
||
for spin in (self.spin_rot, self.spin_shift_x, self.spin_shift_y):
|
||
spin.editingFinished.connect(self._on_manual_edit)
|
||
self.canvas.nudge_translate.connect(self._on_nudge_translate)
|
||
self.canvas.nudge_rotate.connect(self._on_nudge_rotate)
|
||
|
||
# ---- QWizardPage contract ----------------------------------------
|
||
|
||
def initializePage(self):
|
||
if self._masks_ready or self._busy:
|
||
return
|
||
self._set_controls_enabled(False)
|
||
self._start_mask_prep()
|
||
|
||
def isComplete(self) -> bool:
|
||
"""Next is gated on masks being ready and nothing running.
|
||
|
||
Returning False while busy also disables Finish, which is what stops a
|
||
Back/Next/Finish race from reaching the writer with a half-built result.
|
||
"""
|
||
return self._masks_ready and not self._busy
|
||
|
||
def request_stop(self):
|
||
if self._worker is not None:
|
||
self._worker.stop()
|
||
|
||
# ---- mask preparation --------------------------------------------
|
||
|
||
def _start_mask_prep(self):
|
||
st = self._wiz.state
|
||
st.dc4_mv = dict(self._wiz._cached_dc4)
|
||
missing = [a for a in range(self._sras.n_angles) if a not in st.dc4_mv]
|
||
if not missing:
|
||
self._finish_mask_prep()
|
||
return
|
||
self._busy = True
|
||
self.completeChanged.emit()
|
||
self.lbl_status.setText(
|
||
f"Preparing masks: 0/{len(missing)} angle(s) needed…")
|
||
worker = Ch4MaskWorker(self._sras, missing)
|
||
started = self._wiz._parent._run_worker(
|
||
Jobs.ALIGN_MASKS, worker,
|
||
connect=(("angle_done", self._on_mask_done),
|
||
("error", lambda m: self.lbl_status.setText(
|
||
f"Mask preparation failed: {m}"))),
|
||
on_done=self._finish_mask_prep)
|
||
if not started:
|
||
self._busy = False
|
||
self.completeChanged.emit()
|
||
self.lbl_status.setText(
|
||
"Another alignment step is still running — close and reopen.")
|
||
return
|
||
self._worker = worker
|
||
|
||
def _on_mask_done(self, angle_idx: int, dc4_mv: np.ndarray):
|
||
st = self._wiz.state
|
||
st.dc4_mv[angle_idx] = dc4_mv
|
||
self.lbl_status.setText(
|
||
f"Preparing masks: {len(st.dc4_mv)}/{self._sras.n_angles} ready…")
|
||
|
||
def _finish_mask_prep(self):
|
||
st = self._wiz.state
|
||
self._worker = None
|
||
self._busy = False
|
||
if len(st.dc4_mv) < self._sras.n_angles:
|
||
self.completeChanged.emit()
|
||
self.lbl_status.setText(
|
||
"Mask preparation did not finish for every angle.")
|
||
self._wiz.maybe_close_after_job()
|
||
return
|
||
|
||
# Rows and columns get independent factors. A real scan is ~7500 frames
|
||
# wide but only ~750 rows tall, so one shared factor sized for the
|
||
# frames would throw away 10x more row detail than the preview needs.
|
||
max_rows = max(img.shape[0] for img in st.dc4_mv.values())
|
||
max_cols = max(img.shape[1] for img in st.dc4_mv.values())
|
||
st.downsample = (max(1, -(-max_rows // _MAX_PREVIEW_DIM)),
|
||
max(1, -(-max_cols // _MAX_PREVIEW_DIM)))
|
||
|
||
st.params = self._wiz.seed_params()
|
||
self._recompute_masks()
|
||
self._masks_ready = True
|
||
self._set_controls_enabled(True)
|
||
self._apply_prerotation()
|
||
self.lbl_status.setText("Ready.")
|
||
self.completeChanged.emit()
|
||
self._wiz.maybe_close_after_job()
|
||
|
||
def _recompute_masks(self):
|
||
"""Threshold + downsample each angle's in-memory CH4 image. Cheap (a
|
||
compare and a block-mean), so a threshold change re-runs it in full
|
||
rather than re-fetching anything."""
|
||
st = self._wiz.state
|
||
fy, fx = st.downsample
|
||
st.masks_small = {
|
||
a: compute.block_mean_2d((img >= st.threshold_mv).astype(np.float32),
|
||
fy, fx)
|
||
for a, img in st.dc4_mv.items()
|
||
}
|
||
|
||
# ---- parameter changes -------------------------------------------
|
||
|
||
def _on_ref_changed(self, idx: int):
|
||
st = self._wiz.state
|
||
st.ref_angle_idx = idx
|
||
st.fits = {}
|
||
self._apply_prerotation()
|
||
|
||
def _on_threshold_changed(self):
|
||
st = self._wiz.state
|
||
value = self.spin_threshold.value()
|
||
if value == st.threshold_mv:
|
||
return
|
||
st.threshold_mv = value
|
||
if self._masks_ready:
|
||
self._recompute_masks()
|
||
self._refresh()
|
||
|
||
def _on_lock_toggled(self, locked: bool):
|
||
self.spin_search_deg.setEnabled(not locked)
|
||
self.spin_coarse_step.setEnabled(not locked)
|
||
|
||
def _apply_prerotation(self):
|
||
"""Seed every non-reference angle's rotation from the stage angles in
|
||
the file, and redraw — before any correlation has run.
|
||
|
||
This is the cheapest useful thing the page can show: the stage angles
|
||
are usually within a degree or two of the truth, so the stack already
|
||
looks close to right, and how close is a fair first read on whether the
|
||
scan's own metadata can be trusted.
|
||
"""
|
||
if not self._masks_ready:
|
||
return
|
||
st = self._wiz.state
|
||
_seed, _signs, disp_sign = self.combo_prerotate.currentData()
|
||
for a in range(self._sras.n_angles):
|
||
rot = 0.0 if a == st.ref_angle_idx else (
|
||
disp_sign * compute.nominal_delta_deg(self._sras, a,
|
||
st.ref_angle_idx))
|
||
st.params[a] = ManualAngleParams(rot, (0.0, 0.0))
|
||
st.fits = {}
|
||
self._refresh()
|
||
|
||
def _refresh(self):
|
||
"""Rebuild result, stack and every readout from the current params."""
|
||
st = self._wiz.state
|
||
self._wiz.rebuild_result()
|
||
st.geometry_generation += 1
|
||
self._wiz.rebuild_stack()
|
||
self._sync_spins()
|
||
self._update_table()
|
||
self._redraw()
|
||
|
||
# ---- correlation --------------------------------------------------
|
||
|
||
def _reg_kwargs(self) -> dict:
|
||
seed, signs, _disp = self.combo_prerotate.currentData()
|
||
kwargs = {
|
||
"seed_deg": seed,
|
||
"seed_signs": signs,
|
||
"coarse_step_deg": self.spin_coarse_step.value(),
|
||
"fine_dim": self.combo_fine_dim.currentData(),
|
||
}
|
||
if self.chk_lock_rotation.isChecked():
|
||
# Exactly one candidate, no hill-climb: rotation is the seed and
|
||
# only the translation is searched.
|
||
kwargs["refine"] = False
|
||
if len(signs) > 1:
|
||
kwargs["seed_signs"] = (1,)
|
||
return kwargs
|
||
|
||
def _on_correlate(self):
|
||
if not self._masks_ready or self._busy:
|
||
return
|
||
st = self._wiz.state
|
||
angles = [a for a in range(self._sras.n_angles) if a != st.ref_angle_idx]
|
||
if not angles:
|
||
self.lbl_status.setText("Only one angle — nothing to correlate.")
|
||
return
|
||
|
||
search = 0.0 if self.chk_lock_rotation.isChecked() \
|
||
else self.spin_search_deg.value()
|
||
worker = CrossCorrelateWorker(
|
||
self._sras, st.ref_angle_idx, angles, st.dc4_mv,
|
||
sources=self.combo_source.currentData(),
|
||
dc_threshold_mv=st.threshold_mv, search_deg=search,
|
||
reg_kwargs=self._reg_kwargs())
|
||
|
||
# Claim busy and disable the trigger *before* _run_worker, never after:
|
||
# anything that pumps the event loop in between could deliver a second
|
||
# click that starts a thread the first assignment then drops.
|
||
self._busy = True
|
||
self._done_count, self._total = 0, len(angles)
|
||
st.fits = {}
|
||
self.completeChanged.emit()
|
||
self._set_controls_enabled(False)
|
||
self.progress.setValue(0)
|
||
self.lbl_status.setText(f"Cross-correlating: 0/{self._total} angle(s)…")
|
||
|
||
started = self._wiz._parent._run_worker(
|
||
Jobs.ALIGN_CORRELATE, worker,
|
||
connect=(("angle_done", self._on_angle_done),
|
||
("error", lambda m: self.lbl_status.setText(
|
||
f"Cross-correlation failed: {m}"))),
|
||
on_done=self._finish_correlate)
|
||
if not started:
|
||
self._busy = False
|
||
self.completeChanged.emit()
|
||
self._set_controls_enabled(True)
|
||
self.lbl_status.setText(
|
||
"Another alignment step is still running — try again shortly.")
|
||
return
|
||
self._worker = worker
|
||
|
||
def _on_angle_done(self, angle_idx: int, rot: float, sx: float, sy: float,
|
||
score: float, source: str):
|
||
st = self._wiz.state
|
||
st.params[angle_idx] = ManualAngleParams(rot, (sx, sy))
|
||
st.fits[angle_idx] = (score, source)
|
||
self._done_count += 1
|
||
self.progress.setValue(int(self._done_count / max(1, self._total) * 100))
|
||
self.lbl_status.setText(
|
||
f"Cross-correlating: {self._done_count}/{self._total} angle(s)…")
|
||
|
||
def _finish_correlate(self):
|
||
self._worker = None
|
||
self._busy = False
|
||
self.progress.setValue(100)
|
||
self._refresh()
|
||
self._set_controls_enabled(True)
|
||
self.completeChanged.emit()
|
||
self.lbl_status.setText(
|
||
f"Correlated {self._done_count} angle(s) against angle "
|
||
f"{self._wiz.state.ref_angle_idx}. " + self._fit_summary())
|
||
self._wiz.maybe_close_after_job()
|
||
|
||
def _fit_summary(self) -> str:
|
||
"""Worst fit and any angle whose rotation disagrees with its stage angle.
|
||
|
||
Surfaced rather than buried because a single bad acquisition (stage
|
||
glitch, laser dropout) registers poorly and would otherwise be fused in
|
||
silently — knowing *which* angle is what makes dropping it with
|
||
sras_edit_scans.py actionable.
|
||
"""
|
||
st = self._wiz.state
|
||
if not st.fits:
|
||
return ""
|
||
rows = sorted(st.fits.items(), key=lambda kv: kv[1][0])
|
||
worst_a, (worst_score, worst_src) = rows[0]
|
||
parts = [f"Worst fit: angle {worst_a} (score {worst_score:.3f}, "
|
||
f"{worst_src or 'n/a'})."]
|
||
failed = [str(a) for a, (score, src) in rows if score < 0 or src == "none"]
|
||
if failed:
|
||
parts.append(
|
||
"Angle(s) " + ", ".join(failed) + " did not register at all and "
|
||
"are being treated as unrotated — lower the DC threshold, try "
|
||
"Raw signal, nudge them by hand, or drop them with "
|
||
"sras_edit_scans.py.")
|
||
drifted = []
|
||
for a, _ in rows:
|
||
nominal = compute.nominal_delta_deg(self._sras, a, st.ref_angle_idx)
|
||
got = st.params[a].rotation_deg
|
||
dev = min(abs(got - nominal), abs(got + nominal))
|
||
if dev > 1.0:
|
||
drifted.append(f"{a} ({dev:.2f}°)")
|
||
if drifted:
|
||
parts.append("Rotation differs from the stage angle by >1° for "
|
||
"angle(s) " + ", ".join(drifted) + ".")
|
||
return " ".join(parts)
|
||
|
||
# ---- manual nudging ----------------------------------------------
|
||
|
||
def _on_active_changed(self, idx: int):
|
||
self._active_angle = idx
|
||
is_ref = idx == self._wiz.state.ref_angle_idx
|
||
for spin in (self.spin_rot, self.spin_shift_x, self.spin_shift_y):
|
||
spin.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_spins()
|
||
self._redraw()
|
||
|
||
def _sync_spins(self):
|
||
p = self._wiz.state.params.get(self._active_angle, ManualAngleParams())
|
||
for spin, val in ((self.spin_rot, p.rotation_deg),
|
||
(self.spin_shift_x, p.shift_mm[0]),
|
||
(self.spin_shift_y, p.shift_mm[1])):
|
||
with QSignalBlocker(spin):
|
||
spin.setValue(val)
|
||
|
||
def _on_manual_edit(self):
|
||
if not self._masks_ready or self._active_angle == self._wiz.state.ref_angle_idx:
|
||
return
|
||
self._wiz.state.params[self._active_angle] = ManualAngleParams(
|
||
self.spin_rot.value(),
|
||
(self.spin_shift_x.value(), self.spin_shift_y.value()))
|
||
self._refresh()
|
||
|
||
def _on_nudge_translate(self, dir_x: int, dir_y: int, coarse: bool):
|
||
if not self._masks_ready or self._active_angle == self._wiz.state.ref_angle_idx:
|
||
return
|
||
step = self.spin_step_translate.value()
|
||
if coarse:
|
||
step *= self.spin_step_mult.value()
|
||
p = self._wiz.state.params[self._active_angle]
|
||
self._wiz.state.params[self._active_angle] = ManualAngleParams(
|
||
p.rotation_deg,
|
||
(p.shift_mm[0] + dir_x * step, p.shift_mm[1] + dir_y * step))
|
||
self._refresh()
|
||
|
||
def _on_nudge_rotate(self, direction: int, coarse: bool):
|
||
if not self._masks_ready or self._active_angle == self._wiz.state.ref_angle_idx:
|
||
return
|
||
step = self.spin_step_rotate.value()
|
||
if coarse:
|
||
step *= self.spin_step_mult.value()
|
||
p = self._wiz.state.params[self._active_angle]
|
||
self._wiz.state.params[self._active_angle] = ManualAngleParams(
|
||
p.rotation_deg + direction * step, p.shift_mm)
|
||
self._refresh()
|
||
|
||
# ---- drawing ------------------------------------------------------
|
||
|
||
def _set_controls_enabled(self, enabled: bool):
|
||
for w in (self.combo_ref, self.spin_threshold, self.combo_source,
|
||
self.combo_prerotate, self.chk_lock_rotation,
|
||
self.spin_search_deg, self.spin_coarse_step,
|
||
self.combo_fine_dim, self.btn_correlate, self.btn_reset,
|
||
self.grp_nudge):
|
||
w.setEnabled(enabled)
|
||
if enabled:
|
||
self._on_lock_toggled(self.chk_lock_rotation.isChecked())
|
||
self._on_active_changed(self._active_angle)
|
||
|
||
def _redraw(self):
|
||
st = self._wiz.state
|
||
if st.counts is None or st.result is None:
|
||
return
|
||
extent = self._wiz.preview_extent()
|
||
if self.combo_view.currentIndex() == 0:
|
||
self.canvas.show_counts(
|
||
st.counts, self._sras.n_angles, extent,
|
||
f"Mask stack — {self._sras.n_angles} angles, "
|
||
f"ref angle {st.ref_angle_idx}")
|
||
else:
|
||
self.canvas.show_overlay(
|
||
self._rgba(), extent,
|
||
f"Angle {self._active_angle} active "
|
||
f"({self._sras.angles_deg[self._active_angle]:.1f}°)")
|
||
self._update_overlap_text()
|
||
|
||
def _rgba(self) -> np.ndarray:
|
||
"""Alpha-composite each angle's mask in its own colour, active on top."""
|
||
st = self._wiz.state
|
||
n_rows, n_cols = st.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 = st.layers.get(a)
|
||
if layer is None:
|
||
continue
|
||
alpha = _ACTIVE_ALPHA if a == self._active_angle else _BASE_ALPHA
|
||
color = self._wiz.angle_colors[a]
|
||
fg = layer * alpha
|
||
for c in range(3):
|
||
rgba[..., c] = color[c] * fg + rgba[..., c] * rgba[..., 3] * (1 - fg)
|
||
rgba[..., 3] = fg + rgba[..., 3] * (1 - fg)
|
||
return rgba
|
||
|
||
def _update_overlap_text(self):
|
||
st = self._wiz.state
|
||
n = self._sras.n_angles
|
||
stats = compute.overlap_stats(st.counts, n)
|
||
px = st.preview_pitch_mm
|
||
area = abs(px[0] * px[1])
|
||
self.lbl_overlap.setText(
|
||
f"Covered by any angle: {stats['union_px']:,} px "
|
||
f"({stats['union_px'] * area:.2f} mm²) · "
|
||
f"by all {n}: {stats['full_px']:,} px "
|
||
f"({stats['full_frac'] * 100:.1f}% of that) · "
|
||
f"mean overlap {stats['mean_count']:.2f} angles")
|
||
if stats["empty"]:
|
||
self.lbl_warn.setText(
|
||
"No angle covers any pixel — check the DC threshold.")
|
||
elif stats["max_count"] < n:
|
||
self.lbl_warn.setText(
|
||
f"No pixel is covered by all {n} angles (best is "
|
||
f"{stats['max_count']}). Either the alignment is wrong or these "
|
||
f"scans genuinely do not all overlap. You can still export.")
|
||
else:
|
||
self.lbl_warn.setText("")
|
||
|
||
def _update_table(self):
|
||
st = self._wiz.state
|
||
for a in range(self._sras.n_angles):
|
||
score, source = st.fits.get(a, (float("nan"), ""))
|
||
nominal = compute.nominal_delta_deg(self._sras, a, st.ref_angle_idx)
|
||
got = st.params[a].rotation_deg
|
||
dev = min(abs(got - nominal), abs(got + nominal))
|
||
is_ref = a == st.ref_angle_idx
|
||
cells = [
|
||
f"{a}" + (" [ref]" if is_ref else ""),
|
||
f"{got:.3f}",
|
||
"—" if is_ref else f"{dev:.2f}",
|
||
"—" if is_ref or np.isnan(score) else f"{score:.3f}",
|
||
"ref" if is_ref else (source or "—"),
|
||
]
|
||
bad = (not is_ref) and (score < 0 or source == "none")
|
||
weak = (not is_ref) and (not np.isnan(score)) and 0 <= score < 0.3
|
||
for col, text in enumerate(cells):
|
||
item = QTableWidgetItem(text)
|
||
item.setTextAlignment(Qt.AlignmentFlag.AlignCenter)
|
||
if bad:
|
||
item.setForeground(Qt.GlobalColor.red)
|
||
elif weak or (not is_ref and dev > 1.0):
|
||
item.setForeground(Qt.GlobalColor.darkYellow)
|
||
self.table.setItem(a, col, item)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Page 2 — ROI crop
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class RoiPage(QWizardPage):
|
||
"""Pick the rectangle of the shared canvas to keep.
|
||
|
||
Axis-aligned, because v6 geometry is (x_start, x_delta, n_frames, n_rows)
|
||
plus a Y row table — a rectangle on the canvas grid is the only crop the
|
||
file can express, and drawing something else would have to be squared off
|
||
without saying so.
|
||
"""
|
||
|
||
def __init__(self, wizard: AlignmentWizard):
|
||
super().__init__(wizard)
|
||
self._wiz = wizard
|
||
self._sras = wizard.sras
|
||
self._syncing = False
|
||
self._generation = -1
|
||
|
||
self.setTitle("Step 2 — Choose the region to keep")
|
||
self.setSubTitle(
|
||
"Drag a rectangle on the stack, or type canvas pixels directly. "
|
||
"Cropping to the region the angles actually share is usually a big "
|
||
"size saving over the full canvas.")
|
||
self._build_ui()
|
||
|
||
def _build_ui(self):
|
||
root = QHBoxLayout(self)
|
||
|
||
self.canvas = ImageCanvas(rect_only=True)
|
||
left = QWidget()
|
||
ll = QVBoxLayout(left)
|
||
ll.setContentsMargins(0, 0, 0, 0)
|
||
ll.setSpacing(4)
|
||
ll.addWidget(NavigationToolbar2QT(self.canvas, left))
|
||
ll.addWidget(self.canvas, stretch=1)
|
||
root.addWidget(left, stretch=1)
|
||
|
||
panel = QWidget()
|
||
pl = QVBoxLayout(panel)
|
||
pl.setContentsMargins(0, 0, 0, 0)
|
||
pl.setSpacing(8)
|
||
|
||
grp_fit, fl = _group("Preset")
|
||
self.btn_fit_overlap = QPushButton("Fit to full overlap")
|
||
self.btn_fit_union = QPushButton("Fit to any coverage")
|
||
self.btn_whole = QPushButton("Whole canvas (no crop)")
|
||
for b in (self.btn_fit_overlap, self.btn_fit_union, self.btn_whole):
|
||
fl.addWidget(b)
|
||
fl.addWidget(_wrap_label(
|
||
"“Full overlap” finds the largest rectangle lying entirely inside "
|
||
"the region every angle covers — not its bounding box, which would "
|
||
"include corners no angle reaches.", _CSS_HINT))
|
||
pl.addWidget(grp_fit)
|
||
|
||
grp_rect, rl = _group("Crop (canvas pixels)")
|
||
form = _form()
|
||
self.spin_col0 = QSpinBox()
|
||
self.spin_row0 = QSpinBox()
|
||
self.spin_cols = QSpinBox()
|
||
self.spin_rows = QSpinBox()
|
||
for spin, label in ((self.spin_col0, "First column:"),
|
||
(self.spin_row0, "First row:"),
|
||
(self.spin_cols, "Columns:"),
|
||
(self.spin_rows, "Rows:")):
|
||
spin.setRange(0, 1)
|
||
spin.setMinimumWidth(96)
|
||
form.addRow(label, spin)
|
||
rl.addLayout(form)
|
||
self.btn_draw = QPushButton("Draw a new rectangle")
|
||
rl.addWidget(self.btn_draw)
|
||
pl.addWidget(grp_rect)
|
||
|
||
grp_info, il = _group("Result")
|
||
self.lbl_extent = _wrap_label("", _CSS_INFO)
|
||
self.lbl_size = _wrap_label("", _CSS_INFO)
|
||
self.lbl_coverage = _wrap_label("", _CSS_MUTED)
|
||
self.lbl_warn = _wrap_label("", _CSS_WARN)
|
||
for w in (self.lbl_extent, self.lbl_size, self.lbl_coverage, self.lbl_warn):
|
||
il.addWidget(w)
|
||
pl.addWidget(grp_info)
|
||
|
||
pl.addStretch()
|
||
root.addWidget(_scroll_panel(panel, _PANEL_W))
|
||
|
||
self.btn_fit_overlap.clicked.connect(self._on_fit_overlap)
|
||
self.btn_fit_union.clicked.connect(self._on_fit_union)
|
||
self.btn_whole.clicked.connect(self._on_whole)
|
||
self.btn_draw.clicked.connect(self.canvas.start_drawing)
|
||
self.canvas.roi_changed.connect(self._on_roi_changed)
|
||
for spin in (self.spin_col0, self.spin_row0, self.spin_cols, self.spin_rows):
|
||
spin.valueChanged.connect(self._on_spin_changed)
|
||
|
||
# ---- QWizardPage contract ----------------------------------------
|
||
|
||
def initializePage(self):
|
||
st = self._wiz.state
|
||
# A crop is indexed in canvas pixels, so it is meaningless against a
|
||
# canvas built from different rotations. Drop a stale one rather than
|
||
# silently reinterpreting its indices.
|
||
if st.crop is not None and self._generation != st.geometry_generation:
|
||
st.crop = None
|
||
self._generation = st.geometry_generation
|
||
|
||
n_rows, n_cols = st.result.canvas_shape
|
||
# Blocked: setRange clamps the current value into the new range and
|
||
# emits valueChanged, which would run _on_spin_changed and commit a
|
||
# crop built from three not-yet-set spin boxes — landing on 1x1 and
|
||
# making the crop look already-chosen, so the preset below is skipped.
|
||
for spin, lo, hi in ((self.spin_col0, 0, n_cols - 1),
|
||
(self.spin_row0, 0, n_rows - 1),
|
||
(self.spin_cols, 1, n_cols),
|
||
(self.spin_rows, 1, n_rows)):
|
||
with QSignalBlocker(spin):
|
||
spin.setRange(lo, max(lo, hi))
|
||
self._draw_counts()
|
||
self.btn_fit_overlap.setEnabled(
|
||
compute.largest_rect_at_least(st.counts, self._sras.n_angles)
|
||
is not None)
|
||
if st.crop is None:
|
||
self._on_fit_union()
|
||
else:
|
||
self._set_crop(*st.crop)
|
||
|
||
def isComplete(self) -> bool:
|
||
crop = self._wiz.state.crop
|
||
return crop is not None and crop[2] >= 1 and crop[3] >= 1
|
||
|
||
def validatePage(self) -> bool:
|
||
st = self._wiz.state
|
||
cropped = compute.crop_alignment_result(st.result, *st.crop)
|
||
plan = export.plan_export(self._sras, cropped)
|
||
|
||
empty = [a for a in range(self._sras.n_angles) if plan.valid_px[a] == 0]
|
||
if len(empty) == self._sras.n_angles:
|
||
QMessageBox.warning(
|
||
self, "Empty crop",
|
||
"This rectangle contains no data from any angle. Move or grow "
|
||
"it before continuing.")
|
||
return False
|
||
if empty:
|
||
answer = QMessageBox.question(
|
||
self, "Some angles are empty",
|
||
f"Angle(s) {', '.join(map(str, empty))} have no data inside "
|
||
f"this crop and would be written as padding.\n\nContinue anyway?")
|
||
if answer != QMessageBox.StandardButton.Yes:
|
||
return False
|
||
|
||
st.cropped_result = cropped
|
||
return True
|
||
|
||
def cleanupPage(self):
|
||
"""Going back means the canvas geometry may change under this crop."""
|
||
self._wiz.state.crop = None
|
||
self._wiz.state.cropped_result = None
|
||
|
||
# ---- drawing / presets --------------------------------------------
|
||
|
||
def _draw_counts(self):
|
||
st = self._wiz.state
|
||
n = self._sras.n_angles
|
||
# Same colormap the previous page used, so a count keeps its colour
|
||
# across the two pages that show this image on different canvas classes.
|
||
cmap, norm, ticks = count_colormap(n)
|
||
self.canvas.show_image(
|
||
st.counts, self._wiz.preview_extent(), cmap, 0.0, 0.0,
|
||
"X (mm)", "Y (mm)",
|
||
f"Overlap count — drag the crop rectangle ({n} angles)",
|
||
colorbar_label="angles overlapping", cb_ticks=ticks, norm=norm)
|
||
|
||
def _coarse_rect_to_canvas(self, rect) -> tuple[int, int, int, int]:
|
||
"""A rectangle in coarse preview indices, as final canvas pixels.
|
||
|
||
Inset by one coarse block on each side. The coarse grid samples canvas
|
||
pixels 0, fx, 2fx, …, so a coarse pixel reported as fully covered
|
||
stands for a block whose far edge may not be; shrinking by a block keeps
|
||
a "fit to full overlap" rectangle honestly inside the overlap region.
|
||
"""
|
||
st = self._wiz.state
|
||
fy, fx = st.downsample
|
||
row0, col0, nr, nc = rect
|
||
n_rows, n_cols = st.result.canvas_shape
|
||
r0 = min(n_rows - 1, row0 * fy + fy)
|
||
c0 = min(n_cols - 1, col0 * fx + fx)
|
||
nr_c = max(1, min(n_rows - r0, nr * fy - 2 * fy))
|
||
nc_c = max(1, min(n_cols - c0, nc * fx - 2 * fx))
|
||
return r0, c0, nr_c, nc_c
|
||
|
||
def _on_fit_overlap(self):
|
||
st = self._wiz.state
|
||
rect = compute.largest_rect_at_least(st.counts, self._sras.n_angles)
|
||
if rect is None:
|
||
return
|
||
self._set_crop(*self._coarse_rect_to_canvas(rect))
|
||
|
||
def _on_fit_union(self):
|
||
st = self._wiz.state
|
||
rows, cols = np.nonzero(st.counts > 0)
|
||
if rows.size == 0:
|
||
self._on_whole()
|
||
return
|
||
fy, fx = st.downsample
|
||
n_rows, n_cols = st.result.canvas_shape
|
||
r0 = int(rows.min()) * fy
|
||
c0 = int(cols.min()) * fx
|
||
r1 = min(n_rows, (int(rows.max()) + 1) * fy)
|
||
c1 = min(n_cols, (int(cols.max()) + 1) * fx)
|
||
self._set_crop(r0, c0, max(1, r1 - r0), max(1, c1 - c0))
|
||
|
||
def _on_whole(self):
|
||
n_rows, n_cols = self._wiz.state.result.canvas_shape
|
||
self._set_crop(0, 0, n_rows, n_cols)
|
||
|
||
# ---- two-way sync -------------------------------------------------
|
||
|
||
def _set_crop(self, row0: int, col0: int, n_rows: int, n_cols: int):
|
||
st = self._wiz.state
|
||
cr, cc = st.result.canvas_shape
|
||
row0 = int(np.clip(row0, 0, cr - 1))
|
||
col0 = int(np.clip(col0, 0, cc - 1))
|
||
n_rows = int(np.clip(n_rows, 1, cr - row0))
|
||
n_cols = int(np.clip(n_cols, 1, cc - col0))
|
||
st.crop = (row0, col0, n_rows, n_cols)
|
||
|
||
self._syncing = True
|
||
try:
|
||
for spin, val in ((self.spin_row0, row0), (self.spin_col0, col0),
|
||
(self.spin_rows, n_rows), (self.spin_cols, n_cols)):
|
||
with QSignalBlocker(spin):
|
||
spin.setValue(val)
|
||
x0, y0 = self._wiz.canvas_to_mm(col0 - 0.5, row0 - 0.5)
|
||
x1, y1 = self._wiz.canvas_to_mm(col0 + n_cols - 0.5,
|
||
row0 + n_rows - 0.5)
|
||
self.canvas.set_roi(RoiQuad.from_bbox(min(x0, x1), min(y0, y1),
|
||
max(x0, x1), max(y0, y1)))
|
||
finally:
|
||
self._syncing = False
|
||
self._update_readout()
|
||
self.completeChanged.emit()
|
||
|
||
def _on_spin_changed(self):
|
||
if self._syncing:
|
||
return
|
||
self._set_crop(self.spin_row0.value(), self.spin_col0.value(),
|
||
self.spin_rows.value(), self.spin_cols.value())
|
||
|
||
def _on_roi_changed(self):
|
||
# set_roi itself emits roi_changed, so without the guard _set_crop
|
||
# would re-enter through its own canvas update.
|
||
if self._syncing:
|
||
return
|
||
roi = self.canvas.get_roi()
|
||
if roi is None:
|
||
return
|
||
pts = roi.corners()
|
||
(x0, y0), (x1, y1) = pts.min(axis=0), pts.max(axis=0)
|
||
c_a, r_a = self._wiz.mm_to_canvas(x0, y0)
|
||
c_b, r_b = self._wiz.mm_to_canvas(x1, y1)
|
||
col0, col1 = sorted((c_a, c_b))
|
||
row0, row1 = sorted((r_a, r_b))
|
||
# Both edges snap to the nearest pixel boundary with the same rule, so
|
||
# the numeric boxes -> rectangle -> numeric boxes round trip is exact.
|
||
# _SNAP_TOL absorbs the float noise of the mm round trip: a boundary
|
||
# that should land on 11.0 arrives as 11.000000000000002, and a bare
|
||
# ceil() turns that into 12 — one spurious column per edit.
|
||
r0, r1 = _snap_edge(row0), _snap_edge(row1)
|
||
c0, c1 = _snap_edge(col0), _snap_edge(col1)
|
||
self._set_crop(r0, c0, max(1, r1 - r0), max(1, c1 - c0))
|
||
|
||
def _update_readout(self):
|
||
st = self._wiz.state
|
||
row0, col0, n_rows, n_cols = st.crop
|
||
x0, y0 = self._wiz.canvas_to_mm(col0, row0)
|
||
x1, y1 = self._wiz.canvas_to_mm(col0 + n_cols - 1, row0 + n_rows - 1)
|
||
self.lbl_extent.setText(
|
||
f"Output: {n_cols:,} × {n_rows:,} px "
|
||
f"X {min(x0, x1):.3f} … {max(x0, x1):.3f} mm "
|
||
f"Y {min(y0, y1):.3f} … {max(y0, y1):.3f} mm")
|
||
|
||
cropped = compute.crop_alignment_result(st.result, *st.crop)
|
||
plan = export.plan_export(self._sras, cropped)
|
||
self.lbl_size.setText(f"Estimated file size: {_humanize(plan.total_bytes)}"
|
||
f" ({self._sras.n_angles} angles)")
|
||
self.lbl_coverage.setText("Real data per angle: " + ", ".join(
|
||
f"{a}: {plan.coverage_frac(a) * 100:.0f}%"
|
||
for a in range(self._sras.n_angles)))
|
||
empty = [a for a in range(self._sras.n_angles) if plan.valid_px[a] == 0]
|
||
self.lbl_warn.setText(
|
||
f"Angle(s) {', '.join(map(str, empty))} have no data here and would "
|
||
f"be written as padding." if empty else "")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Page 3 — save
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class SavePage(QWizardPage):
|
||
"""Choose a destination and write the aligned, cropped scan.
|
||
|
||
The write runs on a worker and is started by a button rather than from
|
||
validatePage, which must not block the GUI thread for what can be minutes of
|
||
I/O. Finish only becomes available once a file has actually been written, so
|
||
the wizard cannot be completed on a failed export.
|
||
"""
|
||
|
||
def __init__(self, wizard: AlignmentWizard):
|
||
super().__init__(wizard)
|
||
self._wiz = wizard
|
||
self._sras = wizard.sras
|
||
self._busy = False
|
||
self._worker = None
|
||
|
||
self.setTitle("Step 3 — Save the aligned scan")
|
||
self.setSubTitle(
|
||
"Writes a new v6 .sras in which every angle shares the cropped "
|
||
"grid, so it opens already aligned. The original is not modified.")
|
||
self._build_ui()
|
||
|
||
def _build_ui(self):
|
||
root = QVBoxLayout(self)
|
||
|
||
row = QHBoxLayout()
|
||
row.addWidget(QLabel("Save to:"))
|
||
self.edit_path = QLineEdit()
|
||
self.edit_path.setReadOnly(True)
|
||
row.addWidget(self.edit_path, stretch=1)
|
||
self.btn_browse = QPushButton("Browse…")
|
||
row.addWidget(self.btn_browse)
|
||
root.addLayout(row)
|
||
|
||
grp, gl = _group("What will be written")
|
||
self.lbl_summary = _wrap_label("", _CSS_INFO)
|
||
gl.addWidget(self.lbl_summary)
|
||
self.lbl_notes = _wrap_label("", _CSS_WARN)
|
||
gl.addWidget(self.lbl_notes)
|
||
root.addWidget(grp)
|
||
|
||
run = QHBoxLayout()
|
||
self.btn_export = QPushButton("Write .sras")
|
||
run.addWidget(self.btn_export)
|
||
self.btn_cancel = QPushButton("Cancel write")
|
||
self.btn_cancel.setEnabled(False)
|
||
run.addWidget(self.btn_cancel)
|
||
run.addStretch()
|
||
root.addLayout(run)
|
||
|
||
self.progress = QProgressBar()
|
||
self.progress.setRange(0, 100)
|
||
root.addWidget(self.progress)
|
||
|
||
self.lbl_status = _wrap_label("", _CSS_MUTED)
|
||
root.addWidget(self.lbl_status)
|
||
root.addStretch()
|
||
|
||
self.btn_browse.clicked.connect(self._on_browse)
|
||
self.btn_export.clicked.connect(self._on_export)
|
||
self.btn_cancel.clicked.connect(self.request_stop)
|
||
|
||
# ---- QWizardPage contract ----------------------------------------
|
||
|
||
def initializePage(self):
|
||
st = self._wiz.state
|
||
if not st.out_path:
|
||
src = Path(self._sras.path)
|
||
st.out_path = str(src.with_name(f"{src.stem}_aligned.sras"))
|
||
self.edit_path.setText(st.out_path)
|
||
st.exported_path = ""
|
||
self.progress.setValue(0)
|
||
self.lbl_status.setText("")
|
||
self._update_summary()
|
||
self._wiz.setButtonText(QWizard.WizardButton.FinishButton, "Done")
|
||
self.completeChanged.emit()
|
||
|
||
def isComplete(self) -> bool:
|
||
return bool(self._wiz.state.exported_path) and not self._busy
|
||
|
||
def request_stop(self):
|
||
if self._worker is not None:
|
||
self._worker.stop()
|
||
self.lbl_status.setText("Cancelling…")
|
||
|
||
# ---- summary ------------------------------------------------------
|
||
|
||
def _update_summary(self):
|
||
st = self._wiz.state
|
||
plan = export.plan_export(self._sras, st.cropped_result)
|
||
x0, y0 = st.cropped_result.canvas_origin_mm
|
||
self.lbl_summary.setText(
|
||
f"Format: .sras v6, no precomputed cache (the viewer recomputes "
|
||
f"DC/FFT on first open).\n"
|
||
f"Angles: {plan.n_angles}, all sharing one grid of "
|
||
f"{plan.n_frames:,} × {plan.n_rows:,} px.\n"
|
||
f"Origin: X {x0:.4f} mm, Y {y0:.4f} mm · "
|
||
f"pitch {st.cropped_result.canvas_dx_mm * 1000:.2f} × "
|
||
f"{abs(st.cropped_result.canvas_dy_mm) * 1000:.2f} µm.\n"
|
||
f"Stage angles, calibration preambles and the background waveform "
|
||
f"are carried over unchanged.\n"
|
||
f"Size: {_humanize(plan.bytes_per_angle)} per angle, "
|
||
f"{_humanize(plan.total_bytes)} total.\n"
|
||
f"Real data per angle: " + ", ".join(
|
||
f"{a}: {plan.coverage_frac(a) * 100:.0f}%"
|
||
for a in range(plan.n_angles)))
|
||
self.lbl_notes.setText("\n".join(plan.warnings))
|
||
|
||
# ---- export -------------------------------------------------------
|
||
|
||
def _on_browse(self):
|
||
st = self._wiz.state
|
||
path, _ = QFileDialog.getSaveFileName(
|
||
self, "Save aligned .sras", st.out_path,
|
||
"SRAS scans (*.sras);;All files (*)")
|
||
if not path:
|
||
return
|
||
if not path.lower().endswith(".sras"):
|
||
path += ".sras"
|
||
st.out_path = path
|
||
self.edit_path.setText(path)
|
||
st.exported_path = ""
|
||
self.completeChanged.emit()
|
||
|
||
def _on_export(self):
|
||
st = self._wiz.state
|
||
if self._busy or not st.out_path:
|
||
return
|
||
|
||
worker = AlignedExportWorker(self._sras, st.cropped_result, st.out_path)
|
||
self._busy = True
|
||
self.completeChanged.emit()
|
||
self.btn_export.setEnabled(False)
|
||
self.btn_browse.setEnabled(False)
|
||
self.progress.setValue(0)
|
||
self.lbl_status.setText(f"Writing {Path(st.out_path).name}…")
|
||
|
||
started = self._wiz._parent._run_worker(
|
||
Jobs.ALIGN_EXPORT, worker,
|
||
connect=(("progress", self.progress.setValue),
|
||
("finished", self._on_export_finished)))
|
||
if not started:
|
||
self._busy = False
|
||
self.completeChanged.emit()
|
||
self.btn_export.setEnabled(True)
|
||
self.btn_browse.setEnabled(True)
|
||
self.lbl_status.setText(
|
||
"Another alignment step is still running — try again shortly.")
|
||
return
|
||
self._worker = worker
|
||
self.btn_cancel.setEnabled(True)
|
||
|
||
def _on_export_finished(self, written: str, error: str):
|
||
st = self._wiz.state
|
||
self._worker = None
|
||
self._busy = False
|
||
self.btn_export.setEnabled(True)
|
||
self.btn_browse.setEnabled(True)
|
||
self.btn_cancel.setEnabled(False)
|
||
|
||
if error:
|
||
self.progress.setValue(0)
|
||
self.lbl_status.setText(f"Export failed: {error}")
|
||
elif not written:
|
||
self.progress.setValue(0)
|
||
self.lbl_status.setText("Export cancelled; no file was written.")
|
||
else:
|
||
st.exported_path = written
|
||
self.progress.setValue(100)
|
||
self.lbl_status.setText(
|
||
f"Wrote {Path(written).name}. Choose Done to apply this "
|
||
f"alignment to the open scan as well.")
|
||
self.completeChanged.emit()
|
||
self._wiz.maybe_close_after_job()
|
||
|
||
|
||
# Slack when snapping a rectangle edge, in canvas pixels. The mm round trip is
|
||
# two multiplications and a subtraction, so an exact boundary can come back a
|
||
# few ULPs either side of the integer.
|
||
_SNAP_TOL = 1e-6
|
||
|
||
|
||
def _combo(items) -> QComboBox:
|
||
"""A combo box whose size hint does not depend on its longest entry.
|
||
|
||
By default a QComboBox asks for enough width to show its widest item. These
|
||
hold descriptive phrases, and the panel lives in a fixed-width scroll area
|
||
with the horizontal scrollbar off (`_scroll_panel`) — so an unconstrained
|
||
hint pushes the inner widget past the panel and everything on the right,
|
||
including the hint text, is silently clipped instead of scrolling.
|
||
|
||
*items* is a sequence of (label, data) pairs, or of plain labels.
|
||
"""
|
||
combo = QComboBox()
|
||
combo.setSizeAdjustPolicy(
|
||
QComboBox.SizeAdjustPolicy.AdjustToMinimumContentsLengthWithIcon)
|
||
combo.setMinimumContentsLength(10)
|
||
for item in items:
|
||
if isinstance(item, tuple):
|
||
combo.addItem(item[0], item[1])
|
||
else:
|
||
combo.addItem(item)
|
||
return combo
|
||
|
||
|
||
def _snap_edge(coord: float) -> int:
|
||
"""A fractional pixel-index edge, as the nearest pixel boundary index.
|
||
|
||
Pixel k spans [k - 0.5, k + 0.5), so a boundary at fractional coordinate
|
||
*coord* is boundary index coord + 0.5.
|
||
"""
|
||
return int(np.floor(coord + 0.5 + _SNAP_TOL))
|
||
|
||
|
||
def _humanize(n_bytes: int) -> str:
|
||
value = float(n_bytes)
|
||
for unit in ("B", "KB", "MB", "GB", "TB"):
|
||
if value < 1024 or unit == "TB":
|
||
return f"{value:.1f} {unit}" if unit != "B" else f"{int(value)} B"
|
||
value /= 1024
|
||
return f"{value:.1f} TB"
|