Thomas Ales
c30c8b1815
- Extract cache mismatch logic into reusable cache_mismatch_reasons() function for cleaner validation and consistent miss reporting - Expose memory_budget_bytes() for external callers (aligned exporter) - Optimize rebuild_stack() with incremental layer updates when only one angle parameters change, avoiding full reprojection overhead - Remove unused seed_per_angle parameter from alignment wizard - Simplify cached_rf_image() DC masking with cleaner early exit logic - Clean up internal state tracking with explicit origin caching Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>
1447 lines
61 KiB
Python
1447 lines
61 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, 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, _combo,
|
||
_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
|
||
|
||
# Rotation this far from the file's stage angle is worth calling out: the stage
|
||
# angles are usually good to within a degree, so more than that means either the
|
||
# metadata or the fit is wrong.
|
||
_DRIFT_WARN_DEG = 1.0
|
||
|
||
# (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)
|
||
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 self.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,
|
||
cached_dc4_mv: dict[int, np.ndarray]):
|
||
super().__init__(parent)
|
||
self._parent = parent
|
||
self._sras = sras
|
||
self._cached_dc4 = dict(cached_dc4_mv)
|
||
self._closing = False
|
||
# (result, crop, cropped result, plan) — see cropped_plan.
|
||
self._plan_cache: tuple | None = None
|
||
# Canvas origin the current st.layers were reprojected against; a
|
||
# change in it invalidates every layer. See rebuild_stack.
|
||
self._stack_origin_mm: tuple[float, float] | None = None
|
||
|
||
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 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 _reproject(self, angle_idx: int) -> np.ndarray:
|
||
st = self.state
|
||
p = st.params[angle_idx]
|
||
return compute.reproject_mask(
|
||
self._sras, angle_idx, st.ref_angle_idx, st.masks_small[angle_idx],
|
||
p.rotation_deg, p.shift_mm, st.preview_pitch_mm,
|
||
st.result.canvas_origin_mm, st.preview_shape,
|
||
src_downsample=st.downsample)
|
||
|
||
def rebuild_stack(self, only_angle: int | None = None):
|
||
"""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.
|
||
|
||
*only_angle* says that angle is the only one whose parameters changed.
|
||
Reprojection is the one genuinely per-pixel operation on the nudge path
|
||
and it costs the same for every angle, so a full rebuild makes an arrow
|
||
keypress N times more expensive than it needs to be. Each angle's affine
|
||
depends on its own parameters plus the shared canvas, so if the canvas
|
||
came back identical the other layers provably did not move and only the
|
||
nudged one is redrawn; if the canvas did shift, every angle's affine
|
||
changed with it and the hint is ignored.
|
||
"""
|
||
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
|
||
pitch = (st.result.canvas_dx_mm * fx, st.result.canvas_dy_mm * fy)
|
||
shape = (max(1, -(-n_rows // fy)), max(1, -(-n_cols // fx)))
|
||
|
||
incremental = (only_angle is not None
|
||
and st.counts is not None
|
||
and len(st.layers) == self._sras.n_angles
|
||
and pitch == st.preview_pitch_mm
|
||
and shape == st.preview_shape
|
||
and st.result.canvas_origin_mm == self._stack_origin_mm)
|
||
st.preview_pitch_mm, st.preview_shape = pitch, shape
|
||
self._stack_origin_mm = st.result.canvas_origin_mm
|
||
|
||
if incremental:
|
||
before = st.layers[only_angle] > 0.5
|
||
layer = self._reproject(only_angle)
|
||
st.layers[only_angle] = layer
|
||
# Exact integer arithmetic on booleans, so this cannot drift from
|
||
# what a full rebuild would have produced.
|
||
st.counts += (layer > 0.5).astype(np.int16) - before.astype(np.int16)
|
||
return
|
||
|
||
st.layers = {}
|
||
counts = np.zeros(shape, dtype=np.int16)
|
||
for a in range(self._sras.n_angles):
|
||
layer = self._reproject(a)
|
||
st.layers[a] = layer
|
||
counts += (layer > 0.5).astype(np.int16)
|
||
st.counts = counts
|
||
|
||
def cropped_plan(self) -> tuple[compute.AlignmentResult | None,
|
||
export.ExportPlan | None]:
|
||
"""The current crop's AlignmentResult and ExportPlan, or (None, None)
|
||
before a crop exists.
|
||
|
||
Both are pure functions of state.result and state.crop, so neither is
|
||
stored as its own piece of state. They are memoized instead, because
|
||
plan_export counts every output pixel of every angle (hundreds of
|
||
milliseconds on a full-size scan) and the ROI readout, that page's
|
||
validate step and the save page's summary all ask for the same crop.
|
||
"""
|
||
st = self.state
|
||
if st.result is None or st.crop is None:
|
||
return None, None
|
||
cached = self._plan_cache
|
||
if cached is not None and cached[0] is st.result and cached[1] == st.crop:
|
||
return cached[2], cached[3]
|
||
cropped = compute.crop_alignment_result(st.result, *st.crop)
|
||
plan = export.plan_export(self._sras, cropped)
|
||
self._plan_cache = (st.result, st.crop, cropped, plan)
|
||
return cropped, plan
|
||
|
||
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 run_worker(self, key: str, worker, **kwargs) -> bool:
|
||
"""Start a background job on the main window's registry. Returns False
|
||
if another job already holds *key*, which callers must not ignore.
|
||
|
||
The pages go through here rather than reaching into the parent window
|
||
themselves, so what the wizard actually needs from its parent — a job
|
||
runner — is stated in one place instead of at every launch site."""
|
||
return self._parent._run_worker(key, worker, **kwargs)
|
||
|
||
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):
|
||
cropped, _ = self.cropped_plan()
|
||
self.alignment_ready.emit(cropped, self.state.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.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)))
|
||
|
||
self._recompute_masks()
|
||
self._masks_ready = True
|
||
self._set_controls_enabled(True)
|
||
# Sets st.params for every angle, so it is the single source of the
|
||
# starting parameters — nothing needs to seed them beforehand.
|
||
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 _stage_drift_deg(self, angle_idx: int) -> float:
|
||
"""How far this angle's rotation has ended up from the stage angle the
|
||
file reports. Both signs are allowed: the stage's rotational sense
|
||
relative to this module's convention is not knowable from the file,
|
||
so the closer of the two is the honest comparison."""
|
||
st = self._wiz.state
|
||
nominal = compute.nominal_delta_deg(self._sras, angle_idx,
|
||
st.ref_angle_idx)
|
||
got = st.params[angle_idx].rotation_deg
|
||
return min(abs(got - nominal), abs(got + nominal))
|
||
|
||
def _refresh(self, only_angle: int | None = None):
|
||
"""Rebuild result, stack and every readout from the current params.
|
||
|
||
*only_angle* is passed through to rebuild_stack, which uses it to skip
|
||
reprojecting angles that cannot have moved; callers that changed more
|
||
than one angle's parameters must leave it None."""
|
||
st = self._wiz.state
|
||
self._wiz.rebuild_result()
|
||
st.geometry_generation += 1
|
||
self._wiz.rebuild_stack(only_angle)
|
||
self._sync_spins()
|
||
self._update_table()
|
||
self._redraw()
|
||
|
||
# ---- correlation --------------------------------------------------
|
||
|
||
def _reg_kwargs(self) -> dict:
|
||
"""Every argument register_angle_to_reference takes from this page, in
|
||
one dict — so "lock rotation" can express all of what it means in one
|
||
place instead of half here and half at the call site."""
|
||
seed, signs, _disp = self.combo_prerotate.currentData()
|
||
kwargs = {
|
||
"sources": self.combo_source.currentData(),
|
||
"dc_threshold_mv": self._wiz.state.threshold_mv,
|
||
"search_deg": self.spin_search_deg.value(),
|
||
"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
|
||
kwargs["search_deg"] = 0.0
|
||
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
|
||
|
||
worker = CrossCorrelateWorker(
|
||
self._sras, st.ref_angle_idx, angles, st.dc4_mv,
|
||
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.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 = [f"{a} ({dev:.2f}°)" for a, _ in rows
|
||
if (dev := self._stage_drift_deg(a)) > _DRIFT_WARN_DEG]
|
||
if drifted:
|
||
parts.append(f"Rotation differs from the stage angle by "
|
||
f">{_DRIFT_WARN_DEG:g}° 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(self._active_angle)
|
||
|
||
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(self._active_angle)
|
||
|
||
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(self._active_angle)
|
||
|
||
# ---- 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"), ""))
|
||
got = st.params[a].rotation_deg
|
||
dev = self._stage_drift_deg(a)
|
||
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)
|
||
# Each committed edit recounts coverage over the whole crop, so
|
||
# only commit on Enter/focus-out rather than on every digit typed
|
||
# into a four-figure column count.
|
||
spin.setKeyboardTracking(False)
|
||
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()
|
||
# largest_rect_at_least returns None iff no pixel qualifies, so the
|
||
# .any() answers the enable question without running its O(rows*cols)
|
||
# Python sweep on the GUI thread just to throw the rectangle away.
|
||
self.btn_fit_overlap.setEnabled(
|
||
bool((st.counts >= self._sras.n_angles).any()))
|
||
if st.crop is None:
|
||
self._on_fit_union()
|
||
else:
|
||
self._set_crop(*st.crop)
|
||
|
||
def isComplete(self) -> bool:
|
||
# _set_crop is the only writer and clamps both extents to >= 1.
|
||
return self._wiz.state.crop is not None
|
||
|
||
def validatePage(self) -> bool:
|
||
_, plan = self._wiz.cropped_plan()
|
||
empty = plan.empty_angles()
|
||
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
|
||
return True
|
||
|
||
def cleanupPage(self):
|
||
"""Going back means the canvas geometry may change under this crop."""
|
||
self._wiz.state.crop = 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, *, inset_blocks: int
|
||
) -> tuple[int, int, int, int]:
|
||
st = self._wiz.state
|
||
return compute.coarse_rect_to_canvas(rect, st.downsample,
|
||
st.result.canvas_shape,
|
||
inset_blocks=inset_blocks)
|
||
|
||
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, inset_blocks=1))
|
||
|
||
def _on_fit_union(self):
|
||
st = self._wiz.state
|
||
rows, cols = np.nonzero(st.counts > 0)
|
||
if rows.size == 0:
|
||
self._on_whole()
|
||
return
|
||
r0, c0 = int(rows.min()), int(cols.min())
|
||
rect = (r0, c0, int(rows.max()) - r0 + 1, int(cols.max()) - c0 + 1)
|
||
self._set_crop(*self._coarse_rect_to_canvas(rect, inset_blocks=0))
|
||
|
||
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")
|
||
|
||
_, plan = self._wiz.cropped_plan()
|
||
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 = plan.empty_angles()
|
||
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):
|
||
cropped, plan = self._wiz.cropped_plan()
|
||
x0, y0 = cropped.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 {cropped.canvas_dx_mm * 1000:.2f} × "
|
||
f"{abs(cropped.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
|
||
|
||
cropped, _ = self._wiz.cropped_plan()
|
||
worker = AlignedExportWorker(self._sras, cropped, 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.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 _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"
|