diff --git a/sras_compute.py b/sras_compute.py index 6b49a6f..ee34166 100644 --- a/sras_compute.py +++ b/sras_compute.py @@ -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))) diff --git a/sras_viewer.py b/sras_viewer.py index df1bb0b..fb7e8b5 100644 --- a/sras_viewer.py +++ b/sras_viewer.py @@ -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) diff --git a/sras_workers.py b/sras_workers.py index 0068136..5696b0e 100644 --- a/sras_workers.py +++ b/sras_workers.py @@ -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)) diff --git a/tools/test_gui.py b/tools/test_gui.py index 4fd705f..482e11a 100644 --- a/tools/test_gui.py +++ b/tools/test_gui.py @@ -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)