tka/sras-viewer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add scan-editing CLI and alignment tests; extend Manual Alignment correlation

Browse Source 
Continues the Manual Alignment work: refines the FFT cross-correlation and
mask handling, adds sras_edit_scans.py (drop/renumber bad angle scans),
tools/test_alignment.py (registration ground-truth suite), and a rotating
test fixture in make_test_sras.py.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Thomas Ales
2026-08-06 09:35:30 -05:00
parent bbb075ff34
commit d5028db445
8 changed files with 1619 additions and 597 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sras_viewer.py
+131 -73
View File
@@ -821,18 +821,18 @@ class ManualAlignmentDialog(QDialog):
misalignment is visible by eye. Reference angle (always index 0) is
ground truth and never moves; every other angle is aligned to it. The
user picks an "active" angle and nudges its rotation+translation with
the keyboard; Auto De-rotate sets every non-reference angle's rotation to
the known, analytic scan-angle delta without touching any translation;
Auto Cross-Correlate does the same rotation and additionally sets
translation to the FFT-phase-correlation best fit against the reference
(see compute.correlate_translation_mm) — meant to get every angle roughly
stacked on top of each other so keyboard nudging only has to make small
corrections, not find a coarse alignment from scratch. Save writes a
JSON sidecar next to the .sras file and hands a freshly-built, full-
resolution AlignmentResult back to the main window — the exact same
object shape compute_angle_alignment produces, so every existing
Aligned-View code path (apply_alignment, _aligned_canvas_axes, the
pixel-inspector inverse-transform) works completely unmodified.
the keyboard; Auto Cross-Correlate finds every non-reference angle's
rotation *and* translation by registering its image against the
reference's (see compute.register_angle_to_reference) — meant to get every
angle stacked on top of each other so keyboard nudging only has to make
small corrections, not find an alignment from scratch; Auto De-rotate is
the weaker fallback that just seeds rotation from the stage's reported
angle, leaving translation alone. Save writes a JSON sidecar next to the
.sras file and hands a freshly-built, full-resolution AlignmentResult back
to the main window — the exact same object shape compute_angle_alignment
produces, so every existing Aligned-View code path (apply_alignment,
_aligned_canvas_axes, the pixel-inspector inverse-transform) works
completely unmodified.
Non-modal by design (shown via .show(), never .exec() or setModal(True))
so the user can still interact with the main window. Talks back to
@@ -853,6 +853,16 @@ class ManualAlignmentDialog(QDialog):
_ACTIVE_ALPHA = 0.75
_MAX_PREVIEW_DIM = 1024
# (label, sources passed to compute.register_angle_to_reference). "Both"
# registers on each and keeps whichever scores higher per angle, which
# costs roughly double but removes the failure mode where the single
# chosen source is the one that happens to be uninformative for one angle.
_CORRELATE_SOURCES = (
("Both, keep best (recommended)", ("signal", "mask")),
("Raw signal", ("signal",)),
("Thresholded mask", ("mask",)),
)
def __init__(self, parent: "SrasViewerWindow", sras: SrasFile, *,
ref_angle_idx: int, dc_threshold_mv: float,
seed_per_angle: dict[int, ManualAngleParams] | None,
@@ -861,15 +871,16 @@ class ManualAlignmentDialog(QDialog):
self._parent = parent
self._sras = sras
self._ref_angle_idx = ref_angle_idx
self._downsample_factor = 1
self._downsample = (1, 1) # (rows, cols) block-mean factors
self._dc4_mv: dict[int, np.ndarray] = {}
self._masks_small: dict[int, np.ndarray] = {}
self._pivot_mm: dict[int, tuple[float, float]] = {}
self._preview_layers: dict[int, np.ndarray] = {}
self._preview_origin_mm = (0.0, 0.0)
self._preview_shape = (1, 1)
self._preview_dx_mm = self._preview_dy_mm = 1.0
self._preview_pitch_mm = (1.0, 1.0)
self._masks_ready = False
self._fit_notes: dict[int, tuple[float, str]] = {}
self._derotate_sign_flipped = False
self.setWindowTitle(f"Manual Alignment — {sras.path.name}")
self.resize(1150, 760)
@@ -1006,25 +1017,27 @@ class ManualAlignmentDialog(QDialog):
self.grp_correlate, cl = _group("Cross-Correlate (FFT)")
cform = _form()
self.combo_correlate_source = QComboBox()
self.combo_correlate_source.addItems(
["Raw signal (recommended)", "Thresholded mask"])
for label, sources in self._CORRELATE_SOURCES:
self.combo_correlate_source.addItem(label, sources)
cform.addRow("Correlate on:", self.combo_correlate_source)
self.spin_correlate_margin = QDoubleSpinBox()
self.spin_correlate_margin.setRange(0.05, 2.0)
self.spin_correlate_margin.setSingleStep(0.05)
self.spin_correlate_margin.setDecimals(2)
self.spin_correlate_margin.setValue(0.30)
self.spin_correlate_margin.setMinimumWidth(_SPIN_MIN_W)
cform.addRow("Search margin (× extent):", self.spin_correlate_margin)
self.spin_correlate_search_deg = QDoubleSpinBox()
self.spin_correlate_search_deg.setRange(0.0, 180.0)
self.spin_correlate_search_deg.setSingleStep(1.0)
self.spin_correlate_search_deg.setDecimals(1)
self.spin_correlate_search_deg.setSuffix(" °")
self.spin_correlate_search_deg.setValue(6.0)
self.spin_correlate_search_deg.setMinimumWidth(_SPIN_MIN_W)
cform.addRow("Rotation search (±):", self.spin_correlate_search_deg)
cl.addLayout(cform)
self.btn_auto_correlate = QPushButton("Auto Cross-Correlate (vs Reference)")
cl.addWidget(self.btn_auto_correlate)
cl.addWidget(_wrap_label(
"Sets rotation to the known scan angle and translation to the "
"FFT-correlated best fit for every non-reference angle. Run this "
"first, then use manual nudging only for small corrections.",
_CSS_HINT))
"Finds each non-reference angle's rotation *and* translation by "
"cross-correlating its image against the reference's — the stage's "
"reported angle is only the starting point of the search, and both "
"of its signs are tried. Run this first, then nudge only for small "
"corrections.", _CSS_HINT))
panel_l.addWidget(self.grp_correlate)
# ---- Actions ------------------------------------------------------
@@ -1091,15 +1104,16 @@ class ManualAlignmentDialog(QDialog):
def _finish_mask_prep(self):
if len(self._dc4_mv) < self._sras.n_angles:
return # a mask-worker error left some angles unfetched
max_dim = max(max(img.shape) for img in self._dc4_mv.values())
self._downsample_factor = max(1, int(np.ceil(max_dim / self._MAX_PREVIEW_DIM)))
# Rows and columns get their own factor. A real scan is ~7500 frames
# wide but only ~750 rows tall, so one shared factor sized for the
# frames would throw away 8x more row detail than the preview needs and
# leave the overlay too coarse in y to judge alignment by eye.
max_rows = max(img.shape[0] for img in self._dc4_mv.values())
max_cols = max(img.shape[1] for img in self._dc4_mv.values())
self._downsample = (
max(1, int(np.ceil(max_rows / self._MAX_PREVIEW_DIM))),
max(1, int(np.ceil(max_cols / self._MAX_PREVIEW_DIM))))
self._recompute_masks_small()
# Alignment pivot: the CH4-signal-weighted centroid of each angle's
# own footprint (see compute.compute_pivot_points_mm) — computed once
# from the full-res CH4 images and deliberately independent of the
# mask threshold, so it never needs recomputing when that changes
# (unlike _masks_small, which is purely for the overlay's visuals).
self._pivot_mm = compute.compute_pivot_points_mm(self._sras, self._dc4_mv)
self._rebuild_preview_canvas()
self._set_controls_enabled(True)
self.lbl_status.setText("Ready.")
@@ -1108,13 +1122,12 @@ class ManualAlignmentDialog(QDialog):
"""Threshold + downsample every angle's already-in-memory full-res
CH4 mV image. Cheap (a compare + block-mean), so this re-runs in
full whenever the mask-threshold spin box changes — no re-fetch.
Purely for the overlay's visuals — the alignment pivot does not
depend on this threshold (see _pivot_mm / compute_pivot_points_mm)."""
Purely for the overlay's visuals: no alignment geometry depends on this
threshold, only which pixels the overlay paints."""
threshold = self.spin_mask_threshold_mv.value()
factor = self._downsample_factor
fy, fx = self._downsample
self._masks_small = {
a: compute._block_mean_downsample(
(img >= threshold).astype(np.float32), factor)
a: compute._block_mean_2d((img >= threshold).astype(np.float32), fy, fx)
for a, img in self._dc4_mv.items()
}
@@ -1131,32 +1144,35 @@ class ManualAlignmentDialog(QDialog):
active angle. NOT triggered by a translation-only nudge — see
_refresh_active_preview_layer."""
dx_ref, dy_ref = compute._pixel_pitch_mm(self._sras, self._ref_angle_idx)
factor = self._downsample_factor
dx_c, dy_c = dx_ref * factor, dy_ref * factor
origin, shape = compute.union_canvas_mm(
self._sras, self._ref_angle_idx, dx_c, dy_c, self._angle_params,
self._pivot_mm, margin_frac=self._PREVIEW_MARGIN_FRAC)
fy, fx = self._downsample
pitch = (dx_ref * fx, dy_ref * fy)
origin, shape = compute.canvas_for_params(
self._sras, self._ref_angle_idx, pitch, self._angle_params,
margin_frac=self._PREVIEW_MARGIN_FRAC, snap=False)
self._preview_origin_mm, self._preview_shape = origin, shape
self._preview_dx_mm, self._preview_dy_mm = dx_c, dy_c
self._preview_pitch_mm = pitch
self._preview_layers = {
a: compute.reproject_mask(
self._sras, a, self._ref_angle_idx, self._masks_small[a],
self._angle_params[a].rotation_deg, self._angle_params[a].shift_mm,
dx_c, dy_c, origin, shape, self._pivot_mm)
for a in range(self._sras.n_angles)
a: self._reproject(a) for a in range(self._sras.n_angles)
}
self._redraw_overlay()
def _reproject(self, angle_idx: int) -> np.ndarray:
"""One angle's downsampled mask on the current preview canvas.
src_downsample must match _masks_small's block-mean factors, or the
layer lands magnified and offset instead of where the alignment
actually puts it."""
p = self._angle_params[angle_idx]
return compute.reproject_mask(
self._sras, angle_idx, self._ref_angle_idx,
self._masks_small[angle_idx], p.rotation_deg, p.shift_mm,
self._preview_pitch_mm, self._preview_origin_mm, self._preview_shape,
src_downsample=self._downsample)
def _refresh_active_preview_layer(self):
"""Cheap path for a translation-only nudge/edit of the active angle:
reproject just that one angle's downsampled mask onto the *existing*
preview canvas — every other angle's cached layer is untouched."""
a = self._active_angle
self._preview_layers[a] = compute.reproject_mask(
self._sras, a, self._ref_angle_idx, self._masks_small[a],
self._angle_params[a].rotation_deg, self._angle_params[a].shift_mm,
self._preview_dx_mm, self._preview_dy_mm,
self._preview_origin_mm, self._preview_shape, self._pivot_mm)
self._preview_layers[self._active_angle] = self._reproject(self._active_angle)
self._redraw_overlay()
def _redraw_overlay(self):
@@ -1182,7 +1198,7 @@ class ManualAlignmentDialog(QDialog):
rgba[..., 3] = fg_a + rgba[..., 3] * (1 - fg_a)
x0, y0 = self._preview_origin_mm
dx, dy = self._preview_dx_mm, self._preview_dy_mm
dx, dy = self._preview_pitch_mm
x_axis = x0 + np.arange(n_cols) * dx
y_axis = y0 + np.arange(n_rows) * dy
extent = [x_axis[0] - dx / 2, x_axis[-1] + dx / 2,
@@ -1258,18 +1274,29 @@ class ManualAlignmentDialog(QDialog):
# ------------------------------------------------------------------
def _on_auto_derotate(self):
"""Seed every angle's rotation from the stage's reported angle.
A starting point for nudging by eye, not an alignment: the stage's
sign convention relative to this module's is not knowable from the
file, so the sign that lines the scans up is whichever of the two looks
right in the overlay. Auto Cross-Correlate decides that from the images
instead, and is the button to reach for first.
"""
sign = -1.0 if self._derotate_sign_flipped else 1.0
self._derotate_sign_flipped = not self._derotate_sign_flipped
n_changed = 0
for a in range(self._sras.n_angles):
if a == self._ref_angle_idx:
continue
self._angle_params[a].rotation_deg = compute._theta_deg(
self._angle_params[a].rotation_deg = sign * compute._nominal_delta_deg(
self._sras, a, self._ref_angle_idx)
n_changed += 1
self._sync_active_spinboxes()
self._rebuild_preview_canvas()
self.lbl_status.setText(
f"Rotation set to the known scan angle for {n_changed} angle(s) "
"(translation left untouched).")
f"Rotation set to the stage angle ({'−' if sign < 0 else '+'}delta) "
f"for {n_changed} angle(s); translation untouched. Click again to "
"try the opposite sign.")
def _on_auto_correlate(self):
if not self._masks_ready:
@@ -1277,13 +1304,14 @@ class ManualAlignmentDialog(QDialog):
angles = [a for a in range(self._sras.n_angles) if a != self._ref_angle_idx]
if not angles:
return
use_mask = self.combo_correlate_source.currentIndex() == 1
worker = CrossCorrelateWorker(
self._sras, self._ref_angle_idx, angles, self._dc4_mv, self._pivot_mm,
use_mask=use_mask, dc_threshold_mv=self.spin_mask_threshold_mv.value(),
margin_frac=self.spin_correlate_margin.value())
self._sras, self._ref_angle_idx, angles, self._dc4_mv,
sources=self.combo_correlate_source.currentData(),
dc_threshold_mv=self.spin_mask_threshold_mv.value(),
search_deg=self.spin_correlate_search_deg.value())
self._correlate_done_count = 0
self._correlate_total = len(angles)
self._fit_notes = {}
self._set_controls_enabled(False)
self.lbl_status.setText(f"Cross-correlating: 0/{self._correlate_total} angle(s)…")
started = self._parent._run_worker(
@@ -1298,8 +1326,10 @@ class ManualAlignmentDialog(QDialog):
self.lbl_status.setText("Could not start cross-correlation (busy) — try again.")
def _on_correlate_angle_done(self, angle_idx: int, rotation_deg: float,
shift_x_mm: float, shift_y_mm: float):
shift_x_mm: float, shift_y_mm: float,
score: float, source: str):
self._angle_params[angle_idx] = ManualAngleParams(rotation_deg, (shift_x_mm, shift_y_mm))
self._fit_notes[angle_idx] = (score, source)
self._correlate_done_count += 1
self.lbl_status.setText(
f"Cross-correlating: {self._correlate_done_count}/{self._correlate_total} angle(s)…")
@@ -1311,20 +1341,47 @@ class ManualAlignmentDialog(QDialog):
self._sync_active_spinboxes()
self._rebuild_preview_canvas()
self._set_controls_enabled(True)
source = "thresholded mask" if self.combo_correlate_source.currentIndex() == 1 \
else "raw signal"
self.lbl_status.setText(
f"Cross-correlated {self._correlate_done_count} angle(s) against "
f"Angle {self._ref_angle_idx} using the {source}. Nudge from here "
"for any remaining fine correction.")
f"Angle {self._ref_angle_idx}.\n" + self._fit_report())
def _fit_report(self) -> str:
"""Per-angle registration quality, worst first.
Surfaced rather than buried because a single bad acquisition (stage
glitch, laser dropout) registers poorly and would otherwise be fused in
silently — seeing which angle it is, is what makes dropping it with
sras_edit_scans.py actionable. The deviation from the stage's own
reported angle is shown alongside: a large one means the search and the
stage disagree, which is either a genuine mechanical error or a sign
that this angle's fit is not to be trusted.
"""
if not self._fit_notes:
return ""
rows = sorted(self._fit_notes.items(), key=lambda kv: kv[1][0])
worst = rows[0]
lines = [f"Worst fit: angle {worst[0]} (score {worst[1][0]:.3f}, "
f"{worst[1][1]})."]
drifted = []
for a, _note in rows:
nominal = compute._nominal_delta_deg(self._sras, a, self._ref_angle_idx)
got = self._angle_params[a].rotation_deg
dev = min(abs(got - nominal), abs(got + nominal))
if dev > 1.0:
drifted.append(f"{a} ({dev:.2f}°)")
if drifted:
lines.append("Rotation differs from the stage angle by >1° for "
"angle(s) " + ", ".join(drifted) + ".")
lines.append("Nudge from here for any remaining fine correction.")
return " ".join(lines)
def _on_save(self):
threshold = self.spin_mask_threshold_mv.value()
resolved = dict(self._angle_params) # already concrete floats
try:
path = save_manual_alignment(self._sras, self._ref_angle_idx, threshold, resolved)
result = build_manual_alignment(self._sras, self._ref_angle_idx, threshold,
resolved, self._pivot_mm)
result = build_manual_alignment(self._sras, self._ref_angle_idx,
threshold, resolved)
except OSError as exc:
QMessageBox.warning(self, "Save Alignment Failed", str(exc))
return
@@ -1348,6 +1405,7 @@ class ManualAlignmentDialog(QDialog):
f"Could not delete the saved alignment file: {exc}")
return
self._angle_params = {a: ManualAngleParams() for a in range(self._sras.n_angles)}
self._fit_notes = {}
self._sync_active_spinboxes()
self._rebuild_preview_canvas()
self.lbl_status.setText(