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

Add FFT cross-correlation to Manual Alignment, with tunable options

Browse Source 
Manual mode previously offered only Auto De-rotate (rotation from known
scan angles) plus keyboard nudging, so every angle's translation had to
be found entirely by eye — fine for a small correction, unworkable for
the tens-of-mm scatter a real many-angle scan can have between angles
before any translation search runs at all.

Add an "Auto Cross-Correlate (vs Reference)" action that, for every
non-reference angle, sets rotation to the known analytic angle (same as
Auto De-rotate) and translation to the FFT-phase-correlation best fit
against angle #0 (always the reference/ground truth). This is meant to
land every angle's outline roughly stacked on the others so keyboard
nudging only has to make small corrections afterward, per the intended
workflow: cross-correlate first, then nudge.

Exposes two tunable options, since real data may correlate better one
way than the other: "Correlate on" (raw CH4 signal, minus its own
minimum -- the default, robust to a shared threshold not suiting every
angle's real signal level -- or the thresholded binary mask), and
"Search margin" (how far a shift the phase correlation searches before
wraparound bias becomes a risk).

Runs on a background thread (new CrossCorrelateWorker in
sras_workers.py) since correlating many high-resolution angles can take
long enough to visibly freeze the dialog otherwise.

The underlying per-angle correlation math is factored out of
compute_angle_alignment's Step 2 into a new, reusable
correlate_translation_mm, so the existing automatic Fusion -> Angle
Alignment action and the new manual button now share one implementation
instead of two copies of the same phase-correlation logic.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
Thomas Ales
2026-07-31 11:31:14 -05:00
parent 52e91d46fd
commit bbb075ff34
4 changed files with 258 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sras_compute.py
+74 -35
View File
@@ -658,6 +658,71 @@ def _working_canvas_for_pair(sras: SrasFile, ref_idx: int, a_idx: int,
return origin, (n_rows, n_cols)
def correlate_translation_mm(sras: SrasFile, angle_idx: int, ref_angle_idx: int,
signal_mv: dict[int, np.ndarray],
pivot_mm: dict[int, tuple[float, float]], *,
use_mask: bool = False,
dc_threshold_mv: float = 0.0,
margin_frac: float = 0.3,
max_corr_dim: int = 1024) -> tuple[float, float]:
"""FFT phase-correlation translation (shift_x_mm, shift_y_mm) that best
lines up angle_idx's CH4 content onto ref_angle_idx's, given each angle's
own alignment pivot (pivot_mm — see compute_pivot_points_mm) already
rotated about that pivot by the analytic scan-angle delta with zero
shift. Returns (0.0, 0.0) unconditionally for the reference angle.
By default (use_mask=False) correlates on each angle's own raw CH4
signal minus its own minimum — the same weighting _signal_centroid_mm
uses — rather than a dc_threshold_mv binary mask: a shared threshold
that doesn't suit every angle's real signal level can empty or distort
one angle's mask (the same failure mode the alignment pivot was fixed
to avoid), and correlating on the raw signal also lets phase
correlation lock onto real internal sample structure rather than just
the scan window's silhouette. Subtracting each angle's own minimum
(rather than using the signal as-is) keeps the zero-padding surrounding
the rotated image from reading as a spurious high-contrast edge against
a nonzero DC baseline. use_mask=True switches to the binary
>= dc_threshold_mv mask instead (compute_angle_alignment's original
behavior), for cases where raw-signal correlation locks onto noise.
"""
if angle_idx == ref_angle_idx:
return (0.0, 0.0)
def corr_img(a: int) -> np.ndarray:
img = signal_mv[a]
if use_mask:
return (img >= dc_threshold_mv).astype(np.float32)
return (img - img.min()).astype(np.float32)
img_a = corr_img(angle_idx)
img_ref = corr_img(ref_angle_idx)
dx_ref, dy_ref = _pixel_pitch_mm(sras, ref_angle_idx)
max_dim = max(img_a.shape + img_ref.shape)
factor = max(1, int(np.ceil(max_dim / max_corr_dim)))
dx_c, dy_c = dx_ref * factor, dy_ref * factor
small_a = _block_mean_downsample(img_a, factor)
small_ref = _block_mean_downsample(img_ref, factor)
work_origin, work_shape = _working_canvas_for_pair(
sras, ref_angle_idx, angle_idx, dx_c, dy_c, pivot_mm, margin_frac=margin_frac)
m_a, o_a = _build_affine_canvas_to_raw(
sras, angle_idx, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin, pivot_mm)
m_ref, o_ref = _build_affine_canvas_to_raw(
sras, ref_angle_idx, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin, pivot_mm)
rotated_a = scipy_ndimage.affine_transform(
small_a, m_a / factor, offset=o_a / factor,
output_shape=work_shape, order=0, mode="constant", cval=0.0)
embedded_ref = scipy_ndimage.affine_transform(
small_ref, m_ref / factor, offset=o_ref / factor,
output_shape=work_shape, order=0, mode="constant", cval=0.0)
dr, dc = _phase_correlate_shift(embedded_ref, rotated_a)
return (dc * dx_c, dr * dy_c)
def _parallel_map(fn, items, n_workers: int) -> list:
"""fn over items, in order, threaded when it pays."""
items = list(items)
@@ -697,8 +762,6 @@ def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int,
return dc4
dc4_mv = dict(enumerate(_parallel_map(dc4_for, range(n), n_workers)))
masks = {a: (img >= dc_threshold_mv).astype(np.float32)
for a, img in dc4_mv.items()}
# Each angle's own alignment pivot — the CH4-signal-weighted centroid,
# not the raw scan-window bbox center (see compute_pivot_points_mm).
@@ -708,41 +771,17 @@ def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int,
# mask empty can't silently degrade the pivot back to the bbox center.
pivot_mm = compute_pivot_points_mm(sras, dc4_mv=dc4_mv)
# ---- Step 2: coarse correlation stage (downsample first, then rotate)
# Downsampling before affine_transform (not after) is what keeps this
# tractable for a v6 scan with thousands of rows/frames per angle.
max_dim = max(max(m.shape) for m in masks.values())
factor = max(1, int(np.ceil(max_dim / 1024)))
dx_c, dy_c = dx_ref * factor, dy_ref * factor
masks_small = {a: _block_mean_downsample(m, factor) for a, m in masks.items()}
_ticks.clear()
# ---- Step 2: coarse translation via FFT phase correlation, on each
# angle's binarized CH4 mask against the reference's (see
# correlate_translation_mm — also the engine behind ManualAlignmentDialog's
# Auto Cross-Correlate button, there defaulting to the raw signal instead
# of a mask).
def shift_for(a: int) -> tuple[float, float]:
if a == ref_angle_idx:
tick(25, 50)
return (0.0, 0.0)
work_origin, work_shape = _working_canvas_for_pair(
sras, ref_angle_idx, a, dx_c, dy_c, pivot_mm, margin_frac=0.3)
# Coarse canvas->raw affine at native pitch, then rescale by /factor
# so it maps coarse-canvas idx -> coarse (downsampled) raw idx —
# exact for block-mean downsampling (up to the trimmed remainder).
m_a, o_a = _build_affine_canvas_to_raw(
sras, a, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin, pivot_mm)
m_ref, o_ref = _build_affine_canvas_to_raw(
sras, ref_angle_idx, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin,
pivot_mm)
rotated_a = scipy_ndimage.affine_transform(
masks_small[a], m_a / factor, offset=o_a / factor,
output_shape=work_shape, order=0, mode="constant", cval=0.0)
embedded_ref = scipy_ndimage.affine_transform(
masks_small[ref_angle_idx], m_ref / factor, offset=o_ref / factor,
output_shape=work_shape, order=0, mode="constant", cval=0.0)
dr, dc = _phase_correlate_shift(embedded_ref, rotated_a)
shift = correlate_translation_mm(
sras, a, ref_angle_idx, dc4_mv, pivot_mm,
use_mask=True, dc_threshold_mv=dc_threshold_mv)
tick(25, 50)
return (dc * dx_c, dr * dy_c)
return shift
shifts_mm = dict(enumerate(_parallel_map(shift_for, range(n), n_workers)))
sras_viewer.py
+97 -19
View File
@@ -46,7 +46,7 @@ from sras_format import (
)
from sras_workers import (
AngleAlignmentWorker, BatchCacheWorker, Ch4MaskWorker, ComputeWorker,
DcPrecomputeWorker, LoadWorker,
CrossCorrelateWorker, DcPrecomputeWorker, LoadWorker,
)
faulthandler.enable() # print a native stack trace on SIGSEGV/SIGABRT/etc.
@@ -818,27 +818,31 @@ class ManualAlignmentDialog(QDialog):
Shows every angle's binarized CH4 (Bias B) mask overlaid in a distinct
color at partial opacity on one shared canvas, so translation/rotation
misalignment is visible by eye — the thing the automatic phase-
correlation step (compute_angle_alignment) sometimes gets wrong. The user
picks an "active" angle and nudges its rotation+translation with the
keyboard; Auto De-rotate sets every non-reference angle's rotation to the
known, analytic scan-angle delta without touching any translation the
user has already dialed in. Save writes a JSON sidecar next to the .sras
file and hands a freshly-built, full-resolution AlignmentResult back to
the main window — the exact same object shape compute_angle_alignment
produces, so every existing Aligned-View code path (apply_alignment,
_aligned_canvas_axes, the pixel-inspector inverse-transform) works
completely unmodified.
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.
Non-modal by design (shown via .show(), never .exec() or setModal(True))
so the user can still interact with the main window. Talks back to
SrasViewerWindow two ways: it reuses parent._run_worker/_jobs directly
for its own (rare) background mask-fetch step, so the main window's
existing shutdown/lifecycle plumbing covers it for free, and it emits
alignment_saved / alignment_cleared signals for the two moments that
should actually mutate the main window's persistent state — everything
else (nudging, Auto De-rotate, threshold edits) stays purely local to
this dialog until Save.
for its background mask-fetch and cross-correlate steps, so the main
window's existing shutdown/lifecycle plumbing covers both for free, and
it emits alignment_saved / alignment_cleared signals for the two moments
that should actually mutate the main window's persistent state —
everything else (nudging, Auto De-rotate, Auto Cross-Correlate, threshold
edits) stays purely local to this dialog until Save.
"""
alignment_saved = pyqtSignal(object, str) # AlignmentResult, sidecar path (str)
@@ -998,6 +1002,31 @@ class ManualAlignmentDialog(QDialog):
tl.addLayout(tform)
panel_l.addWidget(self.grp_mask_threshold)
# ---- Cross-Correlate (FFT) -----------------------------------------
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"])
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)
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))
panel_l.addWidget(self.grp_correlate)
# ---- Actions ------------------------------------------------------
grp_actions, acl = _group("Actions")
self.btn_auto_derotate = QPushButton("Auto De-rotate (use known angles)")
@@ -1012,7 +1041,7 @@ class ManualAlignmentDialog(QDialog):
panel_l.addWidget(self.lbl_status)
panel_l.addStretch()
root.addWidget(_scroll_panel(panel, 300))
root.addWidget(_scroll_panel(panel, 320))
self.combo_active_angle.currentIndexChanged.connect(self._on_active_angle_changed)
self.spin_active_rotation_deg.editingFinished.connect(self._on_rotation_spin_edited)
@@ -1020,6 +1049,7 @@ class ManualAlignmentDialog(QDialog):
self.spin_active_shift_y_mm.editingFinished.connect(self._on_shift_spin_edited)
self.spin_mask_threshold_mv.editingFinished.connect(self._on_mask_threshold_edited)
self.btn_auto_derotate.clicked.connect(self._on_auto_derotate)
self.btn_auto_correlate.clicked.connect(self._on_auto_correlate)
self.btn_save.clicked.connect(self._on_save)
self.btn_clear.clicked.connect(self._on_clear)
self.btn_close.clicked.connect(self.close)
@@ -1241,6 +1271,53 @@ class ManualAlignmentDialog(QDialog):
f"Rotation set to the known scan angle for {n_changed} angle(s) "
"(translation left untouched).")
def _on_auto_correlate(self):
if not self._masks_ready:
return
angles = [a for a in range(self._sras.n_angles) if a != self._ref_angle_idx]
if not angles:
return
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._correlate_done_count = 0
self._correlate_total = len(angles)
self._set_controls_enabled(False)
self.lbl_status.setText(f"Cross-correlating: 0/{self._correlate_total} angle(s)…")
started = self._parent._run_worker(
"manual_align_correlate", worker,
connect=(
("angle_done", self._on_correlate_angle_done),
("error", self._on_correlate_error),
),
on_done=self._finish_auto_correlate)
if not started:
self._set_controls_enabled(True)
self.lbl_status.setText("Could not start cross-correlation (busy) — try again.")
def _on_correlate_angle_done(self, angle_idx: int, rotation_deg: float,
shift_x_mm: float, shift_y_mm: float):
self._angle_params[angle_idx] = ManualAngleParams(rotation_deg, (shift_x_mm, shift_y_mm))
self._correlate_done_count += 1
self.lbl_status.setText(
f"Cross-correlating: {self._correlate_done_count}/{self._correlate_total} angle(s)…")
def _on_correlate_error(self, msg: str):
self.lbl_status.setText(f"Cross-correlation error: {msg}")
def _finish_auto_correlate(self):
self._sync_active_spinboxes()
self._rebuild_preview_canvas()
self._set_controls_enabled(True)
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.")
def _on_save(self):
threshold = self.spin_mask_threshold_mv.value()
resolved = dict(self._angle_params) # already concrete floats
@@ -1284,6 +1361,7 @@ class ManualAlignmentDialog(QDialog):
self.grp_manual_adjust.setEnabled(enabled and self._active_angle != self._ref_angle_idx)
self.grp_step_sizes.setEnabled(enabled)
self.grp_mask_threshold.setEnabled(enabled)
self.grp_correlate.setEnabled(enabled)
self.btn_auto_derotate.setEnabled(enabled)
self.btn_save.setEnabled(enabled)
self.btn_clear.setEnabled(enabled)
sras_workers.py
+54
View File
@@ -346,3 +346,57 @@ class Ch4MaskWorker(QObject):
self.finished.emit()
except Exception as exc:
self.error.emit(str(exc))
class CrossCorrelateWorker(QObject):
"""FFT phase-correlation translation for each of *angle_indices* against
*ref_angle_idx*, for ManualAlignmentDialog's Auto Cross-Correlate button.
Runs on a background thread — a real many-angle, high-resolution scan's
correlation (even at its downsampled working resolution) can take long
enough that doing all of them on the GUI thread would visibly freeze the
dialog. Rotation is set to the same analytic scan-angle delta Auto
De-rotate uses alongside the correlated shift, since a translation
search is only meaningful once both angles' content is already oriented
the same way. dc4_mv/pivot_mm are the dialog's own already-in-memory
per-angle images/pivots — this worker does no fetching of its own.
"""
angle_done = pyqtSignal(int, float, float, float) # angle_idx, rotation_deg, shift_x_mm, shift_y_mm
finished = pyqtSignal()
error = pyqtSignal(str)
def __init__(self, sras: SrasFile, ref_angle_idx: int, angle_indices: list[int],
dc4_mv: dict[int, np.ndarray], pivot_mm: dict[int, tuple[float, float]],
*, use_mask: bool, dc_threshold_mv: float, margin_frac: float):
super().__init__()
self._sras = sras
self._ref = ref_angle_idx
self._angles = angle_indices
self._dc4_mv = dc4_mv
self._pivot_mm = pivot_mm
self._use_mask = use_mask
self._threshold = dc_threshold_mv
self._margin = margin_frac
def _one(self, a: int) -> tuple[int, float, float, float]:
theta = compute._theta_deg(self._sras, a, self._ref)
dx, dy = compute.correlate_translation_mm(
self._sras, a, self._ref, self._dc4_mv, self._pivot_mm,
use_mask=self._use_mask, dc_threshold_mv=self._threshold,
margin_frac=self._margin)
return a, theta, dx, dy
def run(self):
try:
n_workers, _budget = compute.plan_angle_level(self._sras)
pool = ThreadPoolExecutor(max_workers=max(1, n_workers))
try:
futures = [pool.submit(self._one, a) for a in self._angles]
for fut in as_completed(futures):
a, theta, dx, dy = fut.result()
self.angle_done.emit(a, theta, dx, dy)
finally:
pool.shutdown(wait=True)
self.finished.emit()
except Exception as exc:
self.error.emit(str(exc))
tools/test_gui.py
+33
View File
@@ -278,6 +278,16 @@ def main():
-(float(s.angles_deg[a]) - float(s.angles_deg[0])))
for a in range(s.n_angles)))
# --- FFT phase correlation recovers a known synthetic pixel shift ------
rng = np.random.default_rng(0)
corr_ref = np.zeros((40, 50), dtype=np.float32)
corr_ref[10:25, 15:35] = 1.0
corr_ref += 0.05 * rng.standard_normal(corr_ref.shape).astype(np.float32)
corr_mov = np.roll(corr_ref, shift=(4, -7), axis=(0, 1))
dr, dc = compute._phase_correlate_shift(corr_ref, corr_mov)
check("phase correlation recovers the shift that aligns mov onto ref",
(dr, dc) == (-4, 7), f"got (dr, dc)={(dr, dc)}")
# --- Open: must NOT seed from the still-live automatic AlignmentResult --
# The automatic result's translation comes from FFT phase correlation --
# the very thing manual mode exists to work around -- so manual mode
@@ -344,6 +354,29 @@ def main():
check("reference angle stays identity after auto de-rotate",
dlg._angle_params[dlg._ref_angle_idx].rotation_deg == 0.0)
# --- Auto Cross-Correlate: rotation + FFT-correlated shift, backgrounded -
check("cross-correlate action enabled once masks are ready",
dlg.btn_auto_correlate.isEnabled())
dlg._on_auto_correlate()
check("auto cross-correlate completed", wait_until(
lambda: not win._job_running("manual_align_correlate"), timeout_ms=30000))
check("auto cross-correlate set the known analytic angle for every angle",
all(abs(dlg._angle_params[a].rotation_deg
- compute._theta_deg(s, a, dlg._ref_angle_idx)) < 1e-6
for a in range(s.n_angles) if a != dlg._ref_angle_idx))
check("auto cross-correlate reference angle stays identity",
dlg._angle_params[dlg._ref_angle_idx] == compute.ManualAngleParams())
check("auto cross-correlate re-enabled controls when done",
dlg.grp_correlate.isEnabled() and dlg.btn_save.isEnabled())
check("preview canvas rebuilt after cross-correlate",
len(dlg._preview_layers) == s.n_angles)
# The thresholded-mask option should also work end to end.
dlg.combo_correlate_source.setCurrentIndex(1) # thresholded mask
dlg._on_auto_correlate()
check("auto cross-correlate (thresholded-mask option) completed", wait_until(
lambda: not win._job_running("manual_align_correlate"), timeout_ms=30000))
# --- Save -----------------------------------------------------------------
dlg._on_save()
sidecar = compute.sidecar_path(s.path)