From 26a34f7436fc9c4ebaec9eb09915aa4753a9b0fc Mon Sep 17 00:00:00 2001 From: Thomas Ales Date: Thu, 6 Aug 2026 10:47:34 -0500 Subject: [PATCH] Split sras_viewer.py into a package (pure move) The 2790-line module becomes sras_viewer/: common.py (constants + layout helpers), canvases.py (RoiQuad, ImageCanvas, WaveformCanvas, ManualAlignOverlayCanvas), dialogs.py (FftOptionsDialog, ManualAlignmentDialog), main_window.py (SrasViewerWindow + main), with __init__ re-exporting the public names and __main__ keeping `python -m sras_viewer` working. pyproject gains a `sras-viewer` console script. Code moved verbatim; only import headers are new (pyflakes-clean). tests/test_gui.py patch targets follow the classes to their new modules. Co-Authored-By: Claude Fable 5 --- pyproject.toml | 5 +- sras_viewer.py | 2796 ------------------------------------ sras_viewer/__init__.py | 24 + sras_viewer/__main__.py | 4 + sras_viewer/canvases.py | 542 +++++++ sras_viewer/common.py | 115 ++ sras_viewer/dialogs.py | 767 ++++++++++ sras_viewer/main_window.py | 1403 ++++++++++++++++++ tests/test_gui.py | 6 +- 9 files changed, 2862 insertions(+), 2800 deletions(-) delete mode 100644 sras_viewer.py create mode 100644 sras_viewer/__init__.py create mode 100644 sras_viewer/__main__.py create mode 100644 sras_viewer/canvases.py create mode 100644 sras_viewer/common.py create mode 100644 sras_viewer/dialogs.py create mode 100644 sras_viewer/main_window.py diff --git a/pyproject.toml b/pyproject.toml index ae73ea4..edad96f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,15 +23,18 @@ dependencies = [ [project.optional-dependencies] dev = ["pytest"] +[project.scripts] +sras-viewer = "sras_viewer.main_window:main" + [tool.setuptools] py-modules = [ "sras_format", "sras_compute", "sras_workers", - "sras_viewer", "sras_average", "sras_edit_scans", ] +packages = ["sras_viewer"] [tool.pytest.ini_options] testpaths = ["tests"] diff --git a/sras_viewer.py b/sras_viewer.py deleted file mode 100644 index 23f9344..0000000 --- a/sras_viewer.py +++ /dev/null @@ -1,2796 +0,0 @@ -#!/usr/bin/env python3 -""" -SRAS Scan File Viewer -PyQt6 application for visualizing channel data from .sras binary scan files. - -Channel semantics (fixed by sc3_aui_app.py acquisition settings): - CH1 — RF Acoustic Packet (AC-coupled, 100 mV/div): FFT → peak frequency - CH3 — Bias A (DC-coupled, 50 mV/div): waveform mean - CH4 — Bias B (DC-coupled, 50 mV/div): waveform mean - -RF images are masked: pixels where CH4_dc < dc_threshold show 0. - -File parsing lives in sras_format, image/alignment math in sras_compute, and -background workers in sras_workers — the first two import neither Qt nor -matplotlib so multiprocessing children can load them cheaply. -""" - -import faulthandler -import sys -from pathlib import Path - -import matplotlib as mpl -import numpy as np -from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg, NavigationToolbar2QT -from matplotlib.figure import Figure -from matplotlib.patches import Polygon -from matplotlib.path import Path as MplPath -from PyQt6.QtCore import QObject, QSettings, Qt, QThread, pyqtSignal -from PyQt6.QtGui import QAction, QKeyEvent -from PyQt6.QtWidgets import ( - QApplication, QButtonGroup, QCheckBox, QComboBox, QDialog, QDialogButtonBox, - QDoubleSpinBox, QFileDialog, QFormLayout, QFrame, QGroupBox, QHBoxLayout, - QLabel, QMainWindow, QMessageBox, QProgressDialog, QPushButton, QRadioButton, - QScrollArea, QSizePolicy, QSpinBox, QSplitter, QVBoxLayout, QWidget, -) - -import sras_compute as compute -from sras_compute import ( - PYFFTW_AVAILABLE, ManualAngleParams, apply_alignment, build_manual_alignment, - delete_manual_alignment, load_manual_alignment, save_manual_alignment, - sidecar_path, -) -from sras_format import ( - CH1_IDX, CH3_IDX, CH4_IDX, CH_NAMES, SrasFile, adc_to_mv, mv_to_adc, - _FALLBACK_YMULT_MV, _FALLBACK_YOFF_ADC, -) -from sras_workers import ( - AngleAlignmentWorker, BatchCacheWorker, Ch4MaskWorker, ComputeWorker, - CrossCorrelateWorker, DcPrecomputeWorker, LoadWorker, -) - -faulthandler.enable() # print a native stack trace on SIGSEGV/SIGABRT/etc. - -# --------------------------------------------------------------------------- -# Display constants -# --------------------------------------------------------------------------- - -CH_LABELS = [ - "CH1 — RF (FFT peak freq)", - "CH3 — Bias A (DC mean)", - "CH4 — Bias B (DC mean)", - "CH1 — Velocity (SRAS)", -] - -# Combo index for the derived velocity mode (uses CH1_IDX data) -VELOCITY_MODE_IDX = 3 -# All modes that operate on CH1 waveforms -CH1_DERIVED_MODES = (CH1_IDX, VELOCITY_MODE_IDX) - -CMAPS = ["gray", "viridis", "plasma", "inferno", "hot", "jet", "RdBu_r", "seismic"] - -# (mode_str, status-bar unit, colorbar label) per channel index -_CHANNEL_DISPLAY = { - CH1_IDX: ("RF", "Peak frequency (MHz)", "MHz"), - CH3_IDX: ("DC", "DC mean (mV)", "mV"), - CH4_IDX: ("DC", "DC mean (mV)", "mV"), - VELOCITY_MODE_IDX: ("Velocity", "Velocity (m/s)", "m/s"), -} - -_CSS_HINT = "font-size: 11px; color: #aaa;" -_CSS_INFO = "font-size: 11px;" -_CSS_MUTED = "color: #888; font-size: 11px;" -_CSS_WARN = "color: #e07000; font-size: 11px;" -_CSS_BUSY = "color: #4a90d9; font-size: 11px;" - -# Side-panel column widths (the scroll areas that hold the controls). -_LEFT_PANEL_W = 288 -_RIGHT_PANEL_W = 272 - -# Minimum width for a spin box so its value + suffix are never clipped. -_SPIN_MIN_W = 96 - - -# --------------------------------------------------------------------------- -# Small layout helpers -# --------------------------------------------------------------------------- - -def _wrap_label(text: str = "", css: str | None = None) -> QLabel: - """A word-wrapped QLabel that reports its *wrapped* height to the layout. - - A plain word-wrapped QLabel advertises a single-line minimum height, so in a - fixed-width column the layout happily shrinks it and the extra lines get - clipped. Enabling height-for-width makes the box layout ask for the real - height at the column's width instead. - """ - lbl = QLabel(text) - lbl.setWordWrap(True) - sp = lbl.sizePolicy() - sp.setVerticalPolicy(QSizePolicy.Policy.Minimum) - sp.setHeightForWidth(True) - lbl.setSizePolicy(sp) - if css: - lbl.setStyleSheet(css) - return lbl - - -def _group(title: str) -> tuple[QGroupBox, QVBoxLayout]: - """A group box with consistent, non-cramped internal margins.""" - grp = QGroupBox(title) - lay = QVBoxLayout(grp) - lay.setContentsMargins(10, 8, 10, 10) - lay.setSpacing(6) - return grp, lay - - -def _form() -> QFormLayout: - """A label/field form layout for a narrow side panel.""" - form = QFormLayout() - form.setContentsMargins(0, 0, 0, 0) - form.setHorizontalSpacing(8) - form.setVerticalSpacing(6) - form.setLabelAlignment(Qt.AlignmentFlag.AlignRight - | Qt.AlignmentFlag.AlignVCenter) - form.setFormAlignment(Qt.AlignmentFlag.AlignLeft | Qt.AlignmentFlag.AlignTop) - form.setFieldGrowthPolicy( - QFormLayout.FieldGrowthPolicy.AllNonFixedFieldsGrow) - form.setRowWrapPolicy(QFormLayout.RowWrapPolicy.DontWrapRows) - return form - - -def _scroll_panel(inner: QWidget, width: int) -> QScrollArea: - """Put a side panel in a fixed-width scroll area. - - Without this the panels are sized by the window: a short window squeezes the - controls past their minimum heights, which is what makes text overlap the - widget below it. Scrolling keeps every control at its natural size. - """ - area = QScrollArea() - area.setWidget(inner) - area.setWidgetResizable(True) - area.setFrameShape(QFrame.Shape.NoFrame) - area.setHorizontalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAlwaysOff) - area.setVerticalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAsNeeded) - area.setFixedWidth(width) - area.viewport().setAutoFillBackground(False) - inner.setAutoFillBackground(False) - return area - - -# --------------------------------------------------------------------------- -# ROI (free quadrilateral in data coordinates) -# --------------------------------------------------------------------------- - -class RoiQuad: - """Free quadrilateral defined in data coordinates (mm). - - Stored as 4 corner points (shape (4, 2)) in CCW order: BL, BR, TR, TL. - Each corner can be positioned independently, allowing skewed / - non-orthogonal regions of interest. Because it lives in scan/data - coords it persists unchanged when the displayed channel/mode switches. - """ - - def __init__(self, pts: np.ndarray): - """pts : array-like, shape (4, 2).""" - self._pts = np.asarray(pts, dtype=np.float64).reshape(4, 2).copy() - - @classmethod - def from_bbox(cls, x0: float, y0: float, x1: float, y1: float) -> "RoiQuad": - """Create an axis-aligned rectangle from two opposite corners.""" - lx, rx = min(x0, x1), max(x0, x1) - by, ty = min(y0, y1), max(y0, y1) - return cls(np.array([[lx, by], [rx, by], [rx, ty], [lx, ty]])) - - def copy(self) -> "RoiQuad": - return RoiQuad(self._pts.copy()) - - def corners(self) -> np.ndarray: - """World-coord corners, shape (4, 2), CCW: BL, BR, TR, TL.""" - return self._pts.copy() - - def centroid(self) -> np.ndarray: - return self._pts.mean(axis=0) - - def bbox_size(self) -> np.ndarray: - """Width and height of the axis-aligned bounding box, shape (2,).""" - return self._pts.max(axis=0) - self._pts.min(axis=0) - - def contains(self, x: float, y: float) -> bool: - return bool(MplPath(self._pts).contains_point((x, y))) - - def mask_for_grid(self, x_axis: np.ndarray, - y_axis: np.ndarray) -> np.ndarray: - """Boolean mask (n_rows, n_frames) of pixels whose centres lie - inside the quadrilateral. - - Only the quad's axis-aligned bounding box is tested — meshgrid and - contains_points over the *whole* grid would be tens of millions of - point-in-polygon tests (and hundreds of MB of float64 temporaries) - on a large scan, on every ROI edit. - """ - x = np.asarray(x_axis, dtype=np.float64) - y = np.asarray(y_axis, dtype=np.float64) - mask = np.zeros((y.size, x.size), dtype=bool) - - (x0, y0), (x1, y1) = self._pts.min(axis=0), self._pts.max(axis=0) - cols = np.nonzero((x >= x0) & (x <= x1))[0] - rows = np.nonzero((y >= y0) & (y <= y1))[0] - if cols.size == 0 or rows.size == 0: - return mask - - c0, c1 = int(cols[0]), int(cols[-1]) + 1 - r0, r1 = int(rows[0]), int(rows[-1]) + 1 - X, Y = np.meshgrid(x[c0:c1], y[r0:r1]) - inside = MplPath(self._pts).contains_points( - np.column_stack([X.ravel(), Y.ravel()])) - mask[r0:r1, c0:c1] = inside.reshape(X.shape) - return mask - - -# --------------------------------------------------------------------------- -# Matplotlib canvases -# --------------------------------------------------------------------------- - -class ImageCanvas(FigureCanvasQTAgg): - pixel_clicked = pyqtSignal(int, int) # row_idx, frame_idx - roi_changed = pyqtSignal() # ROI created / edited / cleared - draw_mode_changed = pyqtSignal(bool) # "draw new ROI" arm toggled - - # Interaction state values - _IDLE = "idle" - _DRAW_NEW = "draw_new" - _MOVE = "move" - _DRAG_CORNER = "drag_corner" - - # Hit tolerance (display pixels) for handles. - _HANDLE_PX = 12 - _CLICK_THRESH_PX = 4 # releases within this of press count as a click - - def __init__(self, parent=None): - fig = Figure(figsize=(7, 5), tight_layout=True) - self.ax = fig.add_subplot(111) - super().__init__(fig) - self.setParent(parent) - self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding) - self._extent = None - self._img_shape = None - - # ROI state - self._roi: RoiQuad | None = None - self._roi_artists: list = [] - self._state = self._IDLE - self._draw_mode = False - - # Per-interaction snapshots / anchors - self._press_xy: tuple[float, float] | None = None - self._press_pixel: tuple[float, float] | None = None - self._press_button = None - self._snapshot: RoiQuad | None = None - self._drag_corner_idx: int = -1 - self._move_anchor = None # press-point in world coords - self._draw_previous: RoiQuad | None = None - - self.mpl_connect("button_press_event", self._on_press) - self.mpl_connect("motion_notify_event", self._on_motion) - self.mpl_connect("button_release_event", self._on_release) - - # ------------------------------------------------------------------ - # Public API - # ------------------------------------------------------------------ - - def show_image(self, img: np.ndarray, extent: list[float], cmap: str, - vmin: float, vmax: float, xlabel: str, ylabel: str, title: str, - colorbar_label: str = ""): - self.figure.clf() - self.ax = self.figure.add_subplot(111) - # Patches and lines are destroyed by figure.clf(); drop stale refs. - self._roi_artists = [] - - self._extent = extent - self._img_shape = img.shape - - im = self.ax.imshow( - img, aspect="auto", origin="upper", - extent=extent, cmap=cmap, vmin=vmin, vmax=vmax, - interpolation="nearest", - ) - cb = self.figure.colorbar(im, ax=self.ax, fraction=0.046, pad=0.04) - if colorbar_label: - cb.set_label(colorbar_label) - - self.ax.set_xlabel(xlabel) - self.ax.set_ylabel(ylabel) - self.ax.set_title(title) - - # Re-draw the ROI (if any) on top of the fresh image so it persists - # unchanged across mode / angle / channel switches. - self._draw_roi() - self.draw() - - def get_roi(self) -> RoiQuad | None: - return self._roi - - def set_roi(self, roi: RoiQuad | None): - self._roi = roi.copy() if roi is not None else None - self._draw_roi() - self.draw_idle() - self.roi_changed.emit() - - def clear_roi(self): - self._roi = None - self._remove_roi_artists() - self.draw_idle() - self.roi_changed.emit() - - def start_drawing(self): - """Arm the next click+drag on the image to create a new ROI, - replacing any existing one.""" - self._draw_mode = True - self.setCursor(Qt.CursorShape.CrossCursor) - self.draw_mode_changed.emit(True) - - def cancel_drawing(self): - if self._draw_mode: - self._draw_mode = False - self.setCursor(Qt.CursorShape.ArrowCursor) - self.draw_mode_changed.emit(False) - - # ------------------------------------------------------------------ - # Rendering - # ------------------------------------------------------------------ - - def _remove_roi_artists(self): - for a in self._roi_artists: - try: - a.remove() - except (ValueError, AttributeError, NotImplementedError): - pass - self._roi_artists = [] - - def _draw_roi(self): - self._remove_roi_artists() - if self._roi is None or self.ax is None: - return - corners = self._roi.corners() - - # Filled quad, then a sharp unfilled edge for visibility over bright - # images, then draggable corner handles. - for kwargs in ( - dict(fill=True, facecolor="#ffd93a", edgecolor="#e53935", - alpha=0.22, linewidth=2.0, zorder=10), - dict(fill=False, edgecolor="#e53935", linewidth=1.8, zorder=11), - ): - patch = Polygon(corners, closed=True, **kwargs) - self.ax.add_patch(patch) - self._roi_artists.append(patch) - - self._roi_artists.append(self.ax.scatter( - corners[:, 0], corners[:, 1], s=60, c="white", - edgecolors="#e53935", linewidths=1.6, zorder=13)) - - # ------------------------------------------------------------------ - # Hit testing (display pixels for handles, data coords for "inside") - # ------------------------------------------------------------------ - - def _hit_test(self, event) -> tuple[str, int | None] | None: - if self._roi is None or self.ax is None: - return None - if event.x is None or event.y is None: - return None - corners_disp = self.ax.transData.transform(self._roi.corners()) - click = np.array([event.x, event.y]) - - for i in range(4): - if np.hypot(*(corners_disp[i] - click)) <= self._HANDLE_PX: - return ("corner", i) - - if event.xdata is not None and event.ydata is not None: - if self._roi.contains(event.xdata, event.ydata): - return ("inside", None) - return None - - # ------------------------------------------------------------------ - # Mouse event handlers - # ------------------------------------------------------------------ - - def _on_press(self, event): - if event.inaxes is not self.ax or self._extent is None: - return - if event.button != 1: # only left mouse button - return - # If the matplotlib toolbar is in pan / zoom mode, let it handle - # the interaction instead of starting a ROI manipulation. - tb = getattr(self, "toolbar", None) - if tb is not None and getattr(tb, "mode", ""): - return - - self._press_xy = (event.xdata, event.ydata) - self._press_pixel = (event.x, event.y) - self._press_button = event.button - - if self._draw_mode: - self._draw_previous = self._roi.copy() if self._roi else None - self._roi = RoiQuad.from_bbox(event.xdata, event.ydata, - event.xdata, event.ydata) - self._state = self._DRAW_NEW - self._draw_roi() - self.draw_idle() - return - - hit = self._hit_test(event) - if hit is None: - self._state = self._IDLE - return - - kind, idx = hit - self._snapshot = self._roi.copy() - if kind == "corner": - self._state = self._DRAG_CORNER - self._drag_corner_idx = idx - else: - self._state = self._MOVE - self._move_anchor = (event.xdata, event.ydata) - - def _on_motion(self, event): - if self._state == self._IDLE: - return - if event.xdata is None or event.ydata is None: - return - if event.inaxes is not self.ax: - return - - if self._state == self._DRAW_NEW: - x0, y0 = self._press_xy - self._roi = RoiQuad.from_bbox(x0, y0, event.xdata, event.ydata) - elif self._state == self._MOVE: - delta = np.array([event.xdata - self._move_anchor[0], - event.ydata - self._move_anchor[1]]) - self._roi._pts = self._snapshot.corners() + delta - elif self._state == self._DRAG_CORNER: - self._roi._pts[self._drag_corner_idx] = [event.xdata, event.ydata] - - self._draw_roi() - self.draw_idle() - - def _on_release(self, event): - if event.button != 1 and self._press_button != 1: - return - prev_state = self._state - self._state = self._IDLE - try: - if prev_state == self._DRAW_NEW: - self._finish_draw() - elif prev_state in (self._MOVE, self._DRAG_CORNER): - self._draw_roi() - self.draw_idle() - self.roi_changed.emit() - else: - self._maybe_emit_pixel_click(event) - finally: - self._press_xy = self._press_pixel = None - self._press_button = None - - def _finish_draw(self): - """Commit (or reject) a freshly-dragged quad.""" - if self._extent is not None: - x0, x1, y_bot, y_top = self._extent - min_w = abs(x1 - x0) * 0.01 # minimum: 1% of each axis range - min_h = abs(y_bot - y_top) * 0.01 - else: - min_w = min_h = 1e-6 - - if self._roi is None: - too_small = True - else: - bbox = self._roi.bbox_size() - too_small = bbox[0] < min_w or bbox[1] < min_h - if too_small: - self._roi = self._draw_previous - - self._draw_previous = None - self.cancel_drawing() - self._draw_roi() - self.draw_idle() - self.roi_changed.emit() - - def _maybe_emit_pixel_click(self, event): - """A release close enough to its press counts as a pixel click.""" - if (self._press_pixel is None or event.x is None or event.y is None - or self._extent is None or event.inaxes is not self.ax - or event.xdata is None): - return - dx_px = event.x - self._press_pixel[0] - dy_px = event.y - self._press_pixel[1] - if dx_px * dx_px + dy_px * dy_px > self._CLICK_THRESH_PX ** 2: - return - - x0, x1, y_bot, y_top = self._extent - n_rows, n_frames = self._img_shape - col = int((event.xdata - x0) / (x1 - x0) * n_frames) - row = int((event.ydata - y_top) / (y_bot - y_top) * n_rows) - self.pixel_clicked.emit(max(0, min(row, n_rows - 1)), - max(0, min(col, n_frames - 1))) - - -class WaveformCanvas(FigureCanvasQTAgg): - def __init__(self, parent=None): - fig = Figure(figsize=(8, 3), tight_layout=True) - self.ax_wave = fig.add_subplot(121) - self.ax_right = fig.add_subplot(122) - super().__init__(fig) - self.setParent(parent) - self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding) - - def show_rf_waveform(self, sras: SrasFile, angle_idx: int, - row_idx: int, frame_idx: int, - apply_bg_sub: bool = True): - """CH1 RF: time-domain + FFT spectrum. - - If apply_bg_sub is True and sras.background is not None, the background - waveform is overlaid on the time-domain plot and the FFT is computed - on the subtracted signal. The unsubtracted FFT is also shown faintly - for comparison. - """ - data = sras.data[angle_idx] - waveform = data[row_idx, CH1_IDX, frame_idx, :].astype(np.float32) - t_ns = sras.time_axis_ns() - f_mhz = sras.freq_axis_mhz() - dc3_val = data[row_idx, CH3_IDX, frame_idx, :].astype(np.float32).mean() - dc4_val = data[row_idx, CH4_IDX, frame_idx, :].astype(np.float32).mean() - - bg = sras.background if (apply_bg_sub and sras.background is not None) else None - waveform_plot = waveform - bg if bg is not None else waveform - - self.ax_wave.cla() - self.ax_right.cla() - - if bg is not None: - self.ax_wave.plot(t_ns, waveform, linewidth=0.5, color="#aaaaaa", - label="raw", zorder=1) - self.ax_wave.plot(t_ns, bg, linewidth=0.5, color="#e07030", - linestyle="--", label="background", zorder=2) - self.ax_wave.plot(t_ns, waveform_plot, linewidth=0.7, color="#4488cc", - label="subtracted", zorder=3) - self.ax_wave.legend(fontsize=7, loc="upper right") - else: - self.ax_wave.plot(t_ns, waveform, linewidth=0.7, color="#4488cc") - - self.ax_wave.set_xlabel("Time (ns)") - self.ax_wave.set_ylabel("ADC counts") - bg_tag = " [bg sub]" if bg is not None else "" - dc3_mv = adc_to_mv(dc3_val, *sras.cal(CH3_IDX)) - dc4_mv = adc_to_mv(dc4_val, *sras.cal(CH4_IDX)) - self.ax_wave.set_title( - f"CH1 RF row={row_idx} frame={frame_idx}{bg_tag}\n" - f"CH3={dc3_val:.1f} CH4={dc4_val:.1f} " - f"({dc3_mv:.2f} / {dc4_mv:.2f} mV)", - fontsize=8, - ) - - # FFT of the (possibly subtracted) waveform - power_sub = np.abs(np.fft.rfft(waveform_plot)) ** 2 - power_sub[0] = 0.0 - peak_mhz = f_mhz[int(np.argmax(power_sub))] - - if bg is not None: - # Also show the unsubtracted FFT for reference - power_raw = np.abs(np.fft.rfft(waveform)) ** 2 - power_raw[0] = 0.0 - self.ax_right.plot(f_mhz, power_raw, linewidth=0.5, color="#aaaaaa", - label="raw FFT", zorder=1) - - self.ax_right.plot(f_mhz, power_sub, linewidth=0.7, color="#4488cc", - label="subtracted FFT" if bg is not None else None, zorder=2) - self.ax_right.axvline(peak_mhz, color="tomato", linestyle="--", - linewidth=1.2, label=f"peak = {peak_mhz:.1f} MHz") - self.ax_right.set_xlabel("Frequency (MHz)") - self.ax_right.set_ylabel("Power (arb.)") - self.ax_right.set_title("FFT Power Spectrum") - self.ax_right.set_xlim(0, 500) - self.ax_right.legend(fontsize=8) - - self.draw() - - def show_dc_waveform(self, sras: SrasFile, angle_idx: int, ch_idx: int, - row_idx: int, frame_idx: int): - """CH3 or CH4 DC: time-domain + mean annotation.""" - waveform = sras.data[angle_idx][row_idx, ch_idx, frame_idx, :].astype(np.float32) - mean_val = float(waveform.mean()) - mean_mv = adc_to_mv(mean_val, *sras.cal(ch_idx)) - - self.ax_wave.cla() - self.ax_right.cla() - - self.ax_wave.plot(sras.time_axis_ns(), waveform, linewidth=0.7, color="#4488cc") - self.ax_wave.axhline(mean_val, color="tomato", linestyle="--", - linewidth=1.2, label=f"mean = {mean_val:.2f} ADC") - self.ax_wave.set_xlabel("Time (ns)") - self.ax_wave.set_ylabel("ADC counts") - self.ax_wave.set_title( - f"{CH_NAMES[ch_idx]} DC row={row_idx} frame={frame_idx}") - self.ax_wave.legend(fontsize=8) - - self.ax_right.text( - 0.5, 0.5, - f"DC mode\n\nmean = {mean_val:.3f} ADC\n = {mean_mv:.3f} mV", - ha="center", va="center", - transform=self.ax_right.transAxes, fontsize=11, - ) - self.ax_right.set_axis_off() - - self.draw() - - -# --------------------------------------------------------------------------- -# FFT Options dialog -# --------------------------------------------------------------------------- - -class FftOptionsDialog(QDialog): - """Configure FFT backend and zero-padding. - - Changes take effect only when the user clicks Apply. Cancel discards - all pending edits. The live 'frequency resolution' label updates as - the user adjusts the pad factor so they can see the trade-off before - committing. - """ - - def __init__(self, parent=None, *, - current_backend: str, - current_pad_factor: int, - samples_per_frame: int | None, - sample_rate_hz: float | None, - grating_um: float): - super().__init__(parent) - self.setWindowTitle("FFT Options") - self.setModal(True) - self.setMinimumWidth(380) - - self._samples_per_frame = samples_per_frame - self._sample_rate_hz = sample_rate_hz - self._grating_um = grating_um - - layout = QVBoxLayout(self) - - # ---- Backend --------------------------------------------------- - grp_backend = QGroupBox("FFT Backend") - bl = QVBoxLayout(grp_backend) - - self._btn_scipy = QRadioButton("SciPy FFT (pocketfft) (always available)") - self._btn_pyfftw = QRadioButton( - "pyFFTW (faster for large arrays)" if PYFFTW_AVAILABLE - else "pyFFTW (not installed — run: pip install pyfftw)") - self._btn_pyfftw.setEnabled(PYFFTW_AVAILABLE) - - self._backend_group = QButtonGroup(self) - self._backend_group.addButton(self._btn_scipy, id=0) - self._backend_group.addButton(self._btn_pyfftw, id=1) - - if current_backend == "pyfftw" and PYFFTW_AVAILABLE: - self._btn_pyfftw.setChecked(True) - else: - self._btn_scipy.setChecked(True) - - bl.addWidget(self._btn_scipy) - bl.addWidget(self._btn_pyfftw) - layout.addWidget(grp_backend) - - # ---- Zero-padding ---------------------------------------------- - grp_zp = QGroupBox("Zero-Padding") - zl = QVBoxLayout(grp_zp) - - pad_row = QHBoxLayout() - pad_row.addWidget(QLabel("Pad factor:")) - self._spin_pad = QSpinBox() - self._spin_pad.setRange(1, 256) - self._spin_pad.setValue(max(1, current_pad_factor)) - self._spin_pad.setToolTip( - "Multiply the waveform length by this factor via zero-padding\n" - "before computing the FFT.\n" - "1 = no padding (natural length).\n" - "Powers of 2 (2, 4, 8 …) give the best performance." - ) - self._spin_pad.valueChanged.connect(self._update_info) - pad_row.addWidget(self._spin_pad) - zl.addLayout(pad_row) - - self._lbl_nfft = QLabel() - self._lbl_freq_res = QLabel() - self._lbl_vel_res = QLabel() - for lbl in (self._lbl_nfft, self._lbl_freq_res, self._lbl_vel_res): - lbl.setStyleSheet(_CSS_HINT) - zl.addWidget(lbl) - - layout.addWidget(grp_zp) - - # ---- Buttons --------------------------------------------------- - buttons = QDialogButtonBox() - buttons.addButton("Apply", QDialogButtonBox.ButtonRole.AcceptRole - ).clicked.connect(self.accept) - buttons.addButton("Cancel", QDialogButtonBox.ButtonRole.RejectRole - ).clicked.connect(self.reject) - layout.addWidget(buttons) - - self._update_info() - - def _update_info(self): - spf = self._samples_per_frame - sr = self._sample_rate_hz - pad = self._spin_pad.value() - - if spf is None or sr is None: - self._lbl_nfft.setText("Load a file to preview FFT parameters.") - self._lbl_freq_res.setText("") - self._lbl_vel_res.setText("") - return - - n_fft = spf * pad - freq_res_hz = sr / n_fft - freq_res_mhz = freq_res_hz / 1e6 - # v (m/s) = freq (MHz) × grating (µm) - vel_res_ms = freq_res_mhz * self._grating_um - - self._lbl_nfft.setText(f"FFT points: {spf} × {pad} = {n_fft:,}") - self._lbl_freq_res.setText( - f"Frequency bin: {freq_res_mhz:.4f} MHz ({freq_res_hz / 1e3:.2f} kHz)") - self._lbl_vel_res.setText( - f"Velocity bin: {vel_res_ms:.3f} m/s " - f"(at grating = {self._grating_um:.2f} µm)") - - def get_backend(self) -> str: - return "pyfftw" if self._btn_pyfftw.isChecked() and PYFFTW_AVAILABLE else "scipy" - - def get_pad_factor(self) -> int: - return max(1, self._spin_pad.value()) - - -# --------------------------------------------------------------------------- -# Manual alignment dialog (Fusion -> Manual Alignment...) -# --------------------------------------------------------------------------- - -class ManualAlignOverlayCanvas(FigureCanvasQTAgg): - """Renders ManualAlignmentDialog's multi-angle mask overlay and turns - keyboard input into translate/rotate nudge requests for whichever angle - the dialog currently has active. - - A pure input+render widget — it holds no alignment state and never - touches SrasFile itself; ManualAlignmentDialog owns all of that and - decides, from these signals, whether a cheap single-layer refresh or a - full preview-canvas rebuild is needed. - - FigureCanvasQTAgg is a real QWidget, so keyPressEvent works like on any - other widget, but Qt only ever delivers key events to whichever widget - currently has focus — StrongFocus, plus grabbing focus on click and once - right after the dialog is shown, are both required or arrow keys - silently do nothing. - - Rotate keys are letters (Q/E), not punctuation (comma/period or - brackets): Shift+letter still reports the same Qt.Key on every platform, - whereas Shift+comma/bracket can report a different virtual key - (Key_Less / Key_BraceLeft) depending on platform and keyboard layout — - which would silently break the "Shift = coarse step" modifier for - rotation specifically. Arrow keys have no such hazard. - """ - nudge_translate = pyqtSignal(int, int, bool) # dir_x, dir_y in {-1,0,1}; coarse - nudge_rotate = pyqtSignal(int, bool) # dir in {-1,1} (CCW/CW); coarse - - _TRANSLATE_KEYS = { - Qt.Key.Key_Left: (-1, 0), - Qt.Key.Key_Right: (1, 0), - Qt.Key.Key_Up: (0, -1), - Qt.Key.Key_Down: (0, 1), - } - _ROTATE_KEYS = {Qt.Key.Key_Q: 1, Qt.Key.Key_E: -1} # CCW, CW - - def __init__(self, parent=None): - fig = Figure(figsize=(6, 6), tight_layout=True) - self.ax = fig.add_subplot(111) - super().__init__(fig) - self.setParent(parent) - self.setFocusPolicy(Qt.FocusPolicy.StrongFocus) - self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding) - self.mpl_connect("button_press_event", lambda _e: self.setFocus()) - - def show_overlay(self, rgba: np.ndarray, extent: list[float], title: str): - self.figure.clf() - self.ax = self.figure.add_subplot(111) - self.ax.imshow(rgba, extent=extent, origin="upper", aspect="auto") - self.ax.set_xlabel("X (mm)") - self.ax.set_ylabel("Y (mm)") - self.ax.set_title(title) - self.draw_idle() # coalesces rapid redraws — matters for key-repeat. - - def keyPressEvent(self, event: QKeyEvent): - key = event.key() - coarse = bool(event.modifiers() & Qt.KeyboardModifier.ShiftModifier) - if key in self._TRANSLATE_KEYS: - dx, dy = self._TRANSLATE_KEYS[key] - self.nudge_translate.emit(dx, dy, coarse) - event.accept() - elif key in self._ROTATE_KEYS: - self.nudge_rotate.emit(self._ROTATE_KEYS[key], coarse) - event.accept() - else: - super().keyPressEvent(event) - - -class ManualAlignmentDialog(QDialog): - """Non-modal manual angle-alignment editor (Fusion -> Manual Alignment...). - - 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. 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 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 - SrasViewerWindow two ways: it reuses parent._run_worker/_jobs directly - 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) - alignment_cleared = pyqtSignal() - - _PREVIEW_MARGIN_FRAC = 0.15 - _BASE_ALPHA = 0.42 - _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, - cached_dc4_mv: dict[int, np.ndarray]): - super().__init__(parent) - self._parent = parent - self._sras = sras - self._ref_angle_idx = ref_angle_idx - self._downsample = (1, 1) # (rows, cols) block-mean factors - self._dc4_mv: dict[int, np.ndarray] = {} - self._masks_small: dict[int, np.ndarray] = {} - self._preview_layers: dict[int, np.ndarray] = {} - self._preview_origin_mm = (0.0, 0.0) - self._preview_shape = (1, 1) - 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) - - self._seed_initial_params(seed_per_angle) - 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._active_angle = 1 if ref_angle_idx == 0 and n > 1 else 0 - - self._build_ui(dc_threshold_mv) - self._set_controls_enabled(False) # re-enabled once masks are ready - self._start_mask_prep(cached_dc4_mv) - - def showEvent(self, event): - super().showEvent(event) - self.canvas.setFocus() - - # ------------------------------------------------------------------ - # Construction - # ------------------------------------------------------------------ - - def _seed_initial_params(self, seed_per_angle: dict[int, ManualAngleParams] | None): - seed = seed_per_angle or {} - self._angle_params: dict[int, ManualAngleParams] = { - a: (ManualAngleParams(seed[a].rotation_deg, seed[a].shift_mm) - if a in seed else ManualAngleParams()) - for a in range(self._sras.n_angles) - } - self._angle_params[self._ref_angle_idx] = ManualAngleParams() - - def _build_ui(self, dc_threshold_mv: float): - root = QHBoxLayout(self) - - self.canvas = ManualAlignOverlayCanvas() - left = QWidget() - left_l = QVBoxLayout(left) - left_l.setContentsMargins(0, 0, 0, 0) - left_l.setSpacing(4) - left_l.addWidget(NavigationToolbar2QT(self.canvas, left)) - left_l.addWidget(self.canvas) - root.addWidget(left, stretch=1) - - panel = QWidget() - panel_l = QVBoxLayout(panel) - panel_l.setContentsMargins(0, 0, 0, 0) - panel_l.setSpacing(8) - - # ---- Active Angle ------------------------------------------------- - grp_angle, al = _group("Active Angle") - self.combo_active_angle = QComboBox() - for a in range(self._sras.n_angles): - label = f"Angle {a} ({self._sras.angles_deg[a]:.1f}°)" - if a == self._ref_angle_idx: - label += " [reference]" - self.combo_active_angle.addItem(label) - al.addWidget(self.combo_active_angle) - self.lbl_active_note = _wrap_label("", _CSS_WARN) - al.addWidget(self.lbl_active_note) - panel_l.addWidget(grp_angle) - - # ---- Manual Adjustment --------------------------------------------- - self.grp_manual_adjust, mform_box = _group("Manual Adjustment") - mform = _form() - self.spin_active_rotation_deg = QDoubleSpinBox() - self.spin_active_rotation_deg.setRange(-3600.0, 3600.0) - self.spin_active_rotation_deg.setDecimals(3) - self.spin_active_rotation_deg.setSuffix(" °") - self.spin_active_rotation_deg.setMinimumWidth(_SPIN_MIN_W) - mform.addRow("Rotation:", self.spin_active_rotation_deg) - - self.spin_active_shift_x_mm = QDoubleSpinBox() - self.spin_active_shift_x_mm.setRange(-1e5, 1e5) - self.spin_active_shift_x_mm.setDecimals(4) - self.spin_active_shift_x_mm.setSuffix(" mm") - self.spin_active_shift_x_mm.setMinimumWidth(_SPIN_MIN_W) - mform.addRow("Shift X:", self.spin_active_shift_x_mm) - - self.spin_active_shift_y_mm = QDoubleSpinBox() - self.spin_active_shift_y_mm.setRange(-1e5, 1e5) - self.spin_active_shift_y_mm.setDecimals(4) - self.spin_active_shift_y_mm.setSuffix(" mm") - self.spin_active_shift_y_mm.setMinimumWidth(_SPIN_MIN_W) - mform.addRow("Shift Y:", self.spin_active_shift_y_mm) - mform_box.addLayout(mform) - panel_l.addWidget(self.grp_manual_adjust) - - # ---- Nudge Step Sizes ------------------------------------------------ - self.grp_step_sizes, sl = _group("Nudge Step Sizes") - sform = _form() - self.spin_step_translate_mm = QDoubleSpinBox() - self.spin_step_translate_mm.setRange(0.0001, 1000.0) - self.spin_step_translate_mm.setDecimals(4) - self.spin_step_translate_mm.setSuffix(" mm") - self.spin_step_translate_mm.setValue(0.01) - self.spin_step_translate_mm.setMinimumWidth(_SPIN_MIN_W) - sform.addRow("Translate step:", self.spin_step_translate_mm) - - self.spin_step_rotate_deg = QDoubleSpinBox() - self.spin_step_rotate_deg.setRange(0.001, 90.0) - self.spin_step_rotate_deg.setDecimals(3) - self.spin_step_rotate_deg.setSuffix(" °") - self.spin_step_rotate_deg.setValue(0.1) - self.spin_step_rotate_deg.setMinimumWidth(_SPIN_MIN_W) - sform.addRow("Rotate step:", self.spin_step_rotate_deg) - - self.spin_step_multiplier = QDoubleSpinBox() - self.spin_step_multiplier.setRange(1.0, 1000.0) - self.spin_step_multiplier.setDecimals(1) - self.spin_step_multiplier.setValue(10.0) - self.spin_step_multiplier.setMinimumWidth(_SPIN_MIN_W) - sform.addRow("Coarse × (Shift):", self.spin_step_multiplier) - sl.addLayout(sform) - sl.addWidget(_wrap_label( - "Arrow keys nudge X/Y translation; Q/E nudge rotation (CCW/CW). " - "Hold Shift for the coarse step. Click the image once so it has " - "keyboard focus.", _CSS_HINT)) - panel_l.addWidget(self.grp_step_sizes) - - # ---- Mask Threshold --------------------------------------------------- - self.grp_mask_threshold, tl = _group("Mask Threshold") - tform = _form() - self.spin_mask_threshold_mv = QDoubleSpinBox() - self.spin_mask_threshold_mv.setRange(-500.0, 500.0) - self.spin_mask_threshold_mv.setDecimals(3) - self.spin_mask_threshold_mv.setSuffix(" mV") - self.spin_mask_threshold_mv.setValue(dc_threshold_mv) - self.spin_mask_threshold_mv.setMinimumWidth(_SPIN_MIN_W) - tform.addRow("DC threshold:", self.spin_mask_threshold_mv) - 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() - 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_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( - "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 ------------------------------------------------------ - grp_actions, acl = _group("Actions") - self.btn_auto_derotate = QPushButton("Auto De-rotate (use known angles)") - self.btn_save = QPushButton("Save Alignment") - self.btn_clear = QPushButton("Clear Alignment…") - self.btn_close = QPushButton("Close") - for btn in (self.btn_auto_derotate, self.btn_save, self.btn_clear, self.btn_close): - acl.addWidget(btn) - panel_l.addWidget(grp_actions) - - self.lbl_status = _wrap_label("", _CSS_MUTED) - panel_l.addWidget(self.lbl_status) - panel_l.addStretch() - - 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) - self.spin_active_shift_x_mm.editingFinished.connect(self._on_shift_spin_edited) - 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) - self.canvas.nudge_translate.connect(self._on_nudge_translate) - self.canvas.nudge_rotate.connect(self._on_nudge_rotate) - - self.combo_active_angle.blockSignals(True) - self.combo_active_angle.setCurrentIndex(self._active_angle) - self.combo_active_angle.blockSignals(False) - self._on_active_angle_changed(self._active_angle) - - # ------------------------------------------------------------------ - # Mask preparation (initial CH4 fetch + threshold + downsample) - # ------------------------------------------------------------------ - - def _start_mask_prep(self, cached_dc4_mv: dict[int, np.ndarray]): - self._dc4_mv = dict(cached_dc4_mv) - missing = [a for a in range(self._sras.n_angles) if a not in self._dc4_mv] - if not missing: - self._finish_mask_prep() - return - self.lbl_status.setText(f"Preparing masks: 0/{len(missing)} angle(s) needed…") - started = self._parent._run_worker( - "manual_align_masks", Ch4MaskWorker(self._sras, missing), - connect=( - ("angle_done", self._on_mask_angle_done), - ("error", lambda msg: self.lbl_status.setText(f"Mask prep error: {msg}")), - ), - on_done=self._finish_mask_prep) - if not started: - self.lbl_status.setText( - "Could not start mask preparation (busy) — close and reopen.") - - def _on_mask_angle_done(self, angle_idx: int, dc4_mv: np.ndarray): - self._dc4_mv[angle_idx] = dc4_mv - self.lbl_status.setText( - f"Preparing masks: {len(self._dc4_mv)}/{self._sras.n_angles} ready…") - - def _finish_mask_prep(self): - if len(self._dc4_mv) < self._sras.n_angles: - return # a mask-worker error left some angles unfetched - # 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() - self._rebuild_preview_canvas() - self._set_controls_enabled(True) - self.lbl_status.setText("Ready.") - - def _recompute_masks_small(self): - """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: no alignment geometry depends on this - threshold, only which pixels the overlay paints.""" - threshold = self.spin_mask_threshold_mv.value() - fy, fx = self._downsample - self._masks_small = { - a: compute.block_mean_2d((img >= threshold).astype(np.float32), fy, fx) - for a, img in self._dc4_mv.items() - } - - # ------------------------------------------------------------------ - # Preview canvas: full rebuild vs. incremental single-layer refresh - # ------------------------------------------------------------------ - - def _rebuild_preview_canvas(self): - """Full geometry rebuild: recomputes the shared preview canvas's - origin/shape (rotation can grow the union bbox — translation alone - cannot, per the padding baked in via _PREVIEW_MARGIN_FRAC) and every - angle's reprojected mask layer. Triggered by: dialog open, - mask-threshold change, Auto De-rotate, a rotation nudge/edit of the - 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) - 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_pitch_mm = pitch - self._preview_layers = { - 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.""" - self._preview_layers[self._active_angle] = self._reproject(self._active_angle) - self._redraw_overlay() - - def _redraw_overlay(self): - """Alpha-composite every angle's colored mask layer into one RGBA - image ("all thresholds overlaid with varying opacity"). Each angle - keeps a fixed, distinct color regardless of which is active; the - active angle is drawn last (on top) at a visibly higher alpha so - it's easy to track while nudging.""" - if not self._preview_layers: - return # mask prep hasn't finished yet — nothing to draw - n_rows, n_cols = self._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 = self._preview_layers.get(a) - if layer is None: - continue - alpha = self._ACTIVE_ALPHA if a == self._active_angle else self._BASE_ALPHA - color = self._angle_colors[a] - fg_a = layer * alpha - for c in range(3): - rgba[..., c] = color[c] * fg_a + rgba[..., c] * rgba[..., 3] * (1 - fg_a) - rgba[..., 3] = fg_a + rgba[..., 3] * (1 - fg_a) - - x0, y0 = self._preview_origin_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, - y_axis[-1] + dy / 2, y_axis[0] - dy / 2] - title = (f"Angle {self._active_angle} active " - f"({self._sras.angles_deg[self._active_angle]:.1f}°)") - self.canvas.show_overlay(rgba, extent, title) - - # ------------------------------------------------------------------ - # Angle selection / nudge / edit handlers - # ------------------------------------------------------------------ - - def _on_active_angle_changed(self, angle_idx: int): - self._active_angle = angle_idx - is_ref = angle_idx == self._ref_angle_idx - self.grp_manual_adjust.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_active_spinboxes() - self._redraw_overlay() - - def _sync_active_spinboxes(self): - p = self._angle_params[self._active_angle] - for spin, val in ((self.spin_active_rotation_deg, p.rotation_deg), - (self.spin_active_shift_x_mm, p.shift_mm[0]), - (self.spin_active_shift_y_mm, p.shift_mm[1])): - spin.blockSignals(True) - spin.setValue(val) - spin.blockSignals(False) - - def _on_nudge_translate(self, dir_x: int, dir_y: int, coarse: bool): - if not self._masks_ready or self._active_angle == self._ref_angle_idx: - return - step = self.spin_step_translate_mm.value() - if coarse: - step *= self.spin_step_multiplier.value() - p = self._angle_params[self._active_angle] - p.shift_mm = (p.shift_mm[0] + dir_x * step, p.shift_mm[1] + dir_y * step) - self._sync_active_spinboxes() - self._refresh_active_preview_layer() - - def _on_nudge_rotate(self, direction: int, coarse: bool): - if not self._masks_ready or self._active_angle == self._ref_angle_idx: - return - step = self.spin_step_rotate_deg.value() - if coarse: - step *= self.spin_step_multiplier.value() - self._angle_params[self._active_angle].rotation_deg += direction * step - self._sync_active_spinboxes() - self._rebuild_preview_canvas() - - def _on_rotation_spin_edited(self): - if self._active_angle == self._ref_angle_idx: - return - self._angle_params[self._active_angle].rotation_deg = self.spin_active_rotation_deg.value() - self._rebuild_preview_canvas() - - def _on_shift_spin_edited(self): - if self._active_angle == self._ref_angle_idx: - return - p = self._angle_params[self._active_angle] - p.shift_mm = (self.spin_active_shift_x_mm.value(), self.spin_active_shift_y_mm.value()) - self._refresh_active_preview_layer() - - def _on_mask_threshold_edited(self): - if not self._masks_ready: - return - self._recompute_masks_small() - self._rebuild_preview_canvas() - - # ------------------------------------------------------------------ - # Actions - # ------------------------------------------------------------------ - - 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 = 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 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: - return - angles = [a for a in range(self._sras.n_angles) if a != self._ref_angle_idx] - if not angles: - return - worker = CrossCorrelateWorker( - 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( - "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, - 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)…") - - 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) - self.lbl_status.setText( - f"Cross-correlated {self._correlate_done_count} angle(s) against " - 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) - except OSError as exc: - QMessageBox.warning(self, "Save Alignment Failed", str(exc)) - return - self.lbl_status.setText(f"Saved to {path.name}.") - self.alignment_saved.emit(result, str(path)) - - def _on_clear(self): - reply = QMessageBox.question( - self, "Clear Alignment", - "This resets every angle back to raw/unaligned (0° rotation, no " - "shift) and deletes the saved alignment file for this scan, if " - "any. This cannot be undone. Continue?", - QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No, - QMessageBox.StandardButton.No) - if reply != QMessageBox.StandardButton.Yes: - return - try: - existed = delete_manual_alignment(self._sras) - except OSError as exc: - QMessageBox.warning(self, "Clear Alignment Failed", - 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( - "Alignment cleared; saved file removed." if existed - else "Alignment cleared (there was no saved file).") - self.alignment_cleared.emit() - - def _set_controls_enabled(self, enabled: bool): - self._masks_ready = enabled - self.combo_active_angle.setEnabled(enabled) - 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) - - -# --------------------------------------------------------------------------- -# Main window -# --------------------------------------------------------------------------- - -class SrasViewerWindow(QMainWindow): - def __init__(self, initial_path: str | None = None): - super().__init__() - self.setWindowTitle("SRAS Scan Viewer") - self.resize(1560, 840) - self.setMinimumSize(960, 560) - self.setAcceptDrops(True) - - self._sras: SrasFile | None = None - self._current_image: np.ndarray | None = None - self._current_angle: int = 0 - self._current_ch: int = 0 - self._pending_angle: int = 0 - self._pending_ch: int = 0 - self._pending_bg_sub: bool = True - self._pending_threshold: float = 50.0 # mV - self._pending_fft_pad_factor: int = 1 - - # Live background jobs, keyed by role — see _run_worker. - self._jobs: dict[str, tuple] = {} - self._progress_dlgs: dict[str, QProgressDialog] = {} - - # FFT settings (configured via FFT Options dialog, persisted across - # sessions). IniFormat: predictable cross-platform and redirectable - # in tests. - self._settings = QSettings(QSettings.Format.IniFormat, - QSettings.Scope.UserScope, - "sras-viewer", "sras-viewer") - compute.set_fft_backend(str(self._settings.value("fft/backend", "scipy"))) - try: - pad = int(self._settings.value("fft/pad_factor", 1)) - except (TypeError, ValueError): - pad = 1 - self._fft_pad_factor: int = max(1, min(256, pad)) # 1 = no padding - - # Convert menu: batch DC/FFT compute-and-store (v6 -> v7) - self._batch_errors: list[str] = [] - - # Display-only settings (colormap, grating) never trigger a - # recompute — they're applied to cached data on redraw. DC images - # (CH3/CH4) are cheap and precomputed for every angle in the - # background right after load. CH1/Velocity FFT images are - # computed lazily (with a progress popup) the first time an - # angle/threshold combination is viewed — using the cached DC4 - # image to skip the FFT entirely for masked-out pixels — and - # cached per (angle, bg_sub, n_fft, threshold) so revisiting the - # same combination is free. - self._dc_cache: dict[tuple[int, int], np.ndarray] = {} - self._fft_cache: dict[tuple[int, bool, int | None, float], np.ndarray] = {} - self._dc_generation: int = 0 - - # Angle alignment ("Fusion" menu) - self._alignment_result = None - self._alignment_generation: int = 0 - self._aligned_cache: dict[tuple, np.ndarray] = {} - self._manual_align_dialog: ManualAlignmentDialog | None = None - - self._build_ui() - - if initial_path: - self._load_file(initial_path) - - # ------------------------------------------------------------------ - # Background job plumbing - # ------------------------------------------------------------------ - - def _run_worker(self, key: str, worker: QObject, *, - connect: tuple = (), quit_on: tuple = ("finished",), - on_done=None) -> bool: - """Move *worker* onto its own QThread and start it. Returns False if - a job under *key* is already running. - - Centralises two lifetime hazards that each cost a process abort: - - 1. The job is claimed in self._jobs *before* start() and before - anything below that can pump the Qt event loop (a - QProgressDialog.show() does on first display). If it weren't, a - re-entrant editingFinished could slip past the busy check, start a - second thread, and then have the first call's own assignment - clobber — and destroy while still running — that second QThread. - - 2. thread.finished fires as the thread winds down but does not - guarantee the OS thread has joined. Dropping the last reference to - a QThread whose thread is still running logs "QThread: Destroyed - while thread is still running" and aborts, so wait() first. - """ - if key in self._jobs: - return False - - thread = QThread() - self._jobs[key] = (thread, worker, on_done) # claim before anything pumps - worker.moveToThread(thread) - thread.started.connect(worker.run) - for signal_name, slot in connect: - getattr(worker, signal_name).connect(slot) - for signal_name in quit_on: - getattr(worker, signal_name).connect(thread.quit) - thread.finished.connect(lambda k=key: self._on_job_finished(k)) - thread.start() - return True - - def _on_job_finished(self, key: str): - job = self._jobs.pop(key, None) - if job is None: - return - thread, _worker, on_done = job - thread.wait() # join before releasing our last reference - if on_done is not None: - on_done() - - def _job_running(self, key: str) -> bool: - return key in self._jobs - - # ------------------------------------------------------------------ - # UI construction - # ------------------------------------------------------------------ - - def _build_ui(self): - central = QWidget() - self.setCentralWidget(central) - root = QHBoxLayout(central) - root.setContentsMargins(8, 8, 8, 8) - root.setSpacing(8) - - root.addWidget(self._build_left_panel()) - root.addWidget(self._build_canvases(), stretch=1) - root.addWidget(self._build_right_panel()) - - self.statusBar().showMessage("Open an .sras file to begin.") - self._build_menus() - - def _build_left_panel(self) -> QWidget: - panel = QWidget() - panel_layout = QVBoxLayout(panel) - panel_layout.setContentsMargins(0, 0, 0, 0) - panel_layout.setSpacing(8) - - # ---- File ------------------------------------------------------- - grp_file, fl = _group("File") - self.btn_open = QPushButton("Open .sras…") - self.btn_open.clicked.connect(self._on_open) - self.lbl_filename = _wrap_label("No file loaded", _CSS_MUTED) - fl.addWidget(self.btn_open) - fl.addWidget(self.lbl_filename) - panel_layout.addWidget(grp_file) - - # ---- Scan info -------------------------------------------------- - grp_info, il = _group("Scan Info") - il.setSpacing(3) - self._info = {} - for key in ("Angles", "Rows", "Frames / row", "Samples / frame", - "Sample rate", "X start", "Pixel Δx", "Laser freq"): - lbl = _wrap_label(f"{key}: —", _CSS_INFO) - il.addWidget(lbl) - self._info[key] = lbl - - # frame-count / format notes - self.lbl_frame_warn = _wrap_label("", _CSS_WARN) - il.addWidget(self.lbl_frame_warn) - - # background DC-precompute progress - self.lbl_dc_precompute = _wrap_label("", _CSS_BUSY) - il.addWidget(self.lbl_dc_precompute) - panel_layout.addWidget(grp_info) - - # ---- View settings ---------------------------------------------- - grp_view, vl = _group("View Settings") - - view_form = _form() - - self.spin_angle = QSpinBox() - self.spin_angle.setRange(0, 0) - self.spin_angle.setEnabled(False) - self.spin_angle.setMinimumWidth(64) - self.spin_angle.editingFinished.connect(self._on_view_changed) - self.lbl_angle_deg = QLabel("—") - angle_field = QWidget() - ar = QHBoxLayout(angle_field) - ar.setContentsMargins(0, 0, 0, 0) - ar.setSpacing(6) - ar.addWidget(self.spin_angle) - ar.addWidget(self.lbl_angle_deg) - ar.addStretch() - view_form.addRow("Angle:", angle_field) - - self.combo_channel = QComboBox() - self.combo_channel.addItems(CH_LABELS) - self.combo_channel.setEnabled(False) - self.combo_channel.setSizePolicy(QSizePolicy.Policy.Expanding, - QSizePolicy.Policy.Fixed) - self.combo_channel.setSizeAdjustPolicy( - QComboBox.SizeAdjustPolicy.AdjustToMinimumContentsLengthWithIcon) - self.combo_channel.setMinimumContentsLength(12) - self.combo_channel.currentIndexChanged.connect(self._on_channel_changed) - view_form.addRow("Channel:", self.combo_channel) - vl.addLayout(view_form) - - sep = QFrame() - sep.setFrameShape(QFrame.Shape.HLine) - sep.setStyleSheet("color: #555;") - vl.addWidget(sep) - - # DC threshold (for RF / CH1 masking) - self.grp_threshold, tl = _group("RF Mask Threshold (CH1 only)") - thr_form = _form() - self.spin_threshold_mv = QDoubleSpinBox() - self.spin_threshold_mv.setRange(-500.0, 500.0) - self.spin_threshold_mv.setDecimals(3) - self.spin_threshold_mv.setSingleStep(0.025) - self.spin_threshold_mv.setSuffix(" mV") - self.spin_threshold_mv.setValue(50.0) - self.spin_threshold_mv.setEnabled(False) - self.spin_threshold_mv.setMinimumWidth(_SPIN_MIN_W) - self.spin_threshold_mv.editingFinished.connect(self._on_threshold_changed) - thr_form.addRow("DC threshold:", self.spin_threshold_mv) - tl.addLayout(thr_form) - self.lbl_threshold_adc = _wrap_label( - f"≈ {mv_to_adc(50.0):.1f} ADC counts", _CSS_MUTED) - tl.addWidget(self.lbl_threshold_adc) - vl.addWidget(self.grp_threshold) - - # Background subtraction (v4+ files only) - self.chk_bg_sub = QCheckBox("Background subtraction (CH1 only)") - self.chk_bg_sub.setChecked(True) - self.chk_bg_sub.setEnabled(False) - self.chk_bg_sub.setToolTip( - "Subtract the stored background waveform from each CH1 frame\n" - "before computing the FFT (v4+ files only)." - ) - self.chk_bg_sub.toggled.connect(self._on_bg_sub_toggled) - vl.addWidget(self.chk_bg_sub) - - # Aligned View (Fusion → Angle Alignment result) - self.chk_aligned_view = QCheckBox("Aligned View (Fusion)") - self.chk_aligned_view.setChecked(False) - self.chk_aligned_view.setEnabled(False) - self.chk_aligned_view.setToolTip( - "Show the current angle/channel resampled onto the shared,\n" - "rotation+translation-aligned canvas from Fusion → Angle\n" - "Alignment. Uncheck to see the raw per-angle scan grid." - ) - self.chk_aligned_view.toggled.connect(self._on_aligned_view_toggled) - vl.addWidget(self.chk_aligned_view) - - self.btn_export_csv = QPushButton("Export Image as CSV…") - self.btn_export_csv.setEnabled(False) - self.btn_export_csv.setToolTip( - "Save the current CH1 image (one scan row per CSV line).") - self.btn_export_csv.clicked.connect(self._on_export_csv) - vl.addWidget(self.btn_export_csv) - - panel_layout.addWidget(grp_view) - - # ---- ROI --------------------------------------------------------- - grp_roi, rl = _group("ROI (Region of Interest)") - - self.btn_draw_roi = QPushButton("Draw ROI") - self.btn_draw_roi.setCheckable(True) - self.btn_draw_roi.setEnabled(False) - self.btn_draw_roi.setToolTip( - "Arm next click+drag on the image to draw a new ROI\n" - "(replaces any existing one). Click again to cancel.\n" - "After drawing, drag inside to move, or grab corners to reshape.\n" - "The ROI is persistent across channels / modes / angles." - ) - self.btn_draw_roi.toggled.connect(self._on_draw_roi_toggled) - rl.addWidget(self.btn_draw_roi) - - self.btn_clear_roi = QPushButton("Clear ROI") - self.btn_clear_roi.setEnabled(False) - self.btn_clear_roi.clicked.connect(self._on_clear_roi) - rl.addWidget(self.btn_clear_roi) - - self.btn_export_roi = QPushButton("Export ROI as CSV…") - self.btn_export_roi.setEnabled(False) - self.btn_export_roi.setToolTip( - "Save every pixel whose centre lies inside the ROI as CSV.\n" - "Columns: row, frame, x_mm, y_mm, value.\n" - "Corner coordinates of the quad are written in the file header." - ) - self.btn_export_roi.clicked.connect(self._on_export_roi_csv) - rl.addWidget(self.btn_export_roi) - - self.lbl_roi_center = _wrap_label("centroid: —", _CSS_HINT) - self.lbl_roi_size = _wrap_label("bbox: —", _CSS_HINT) - self.lbl_roi_npix = _wrap_label("pixels inside: —", _CSS_HINT) - for lbl in (self.lbl_roi_center, self.lbl_roi_size, self.lbl_roi_npix): - rl.addWidget(lbl) - - panel_layout.addWidget(grp_roi) - panel_layout.addStretch() - return _scroll_panel(panel, _LEFT_PANEL_W) - - def _build_canvases(self) -> QWidget: - splitter = QSplitter(Qt.Orientation.Vertical) - splitter.setChildrenCollapsible(False) - - img_widget = QWidget() - img_vl = QVBoxLayout(img_widget) - img_vl.setContentsMargins(0, 0, 0, 0) - img_vl.setSpacing(4) - self.image_canvas = ImageCanvas() - self.image_canvas.setMinimumHeight(220) - self.image_canvas.pixel_clicked.connect(self._on_pixel_clicked) - self.image_canvas.roi_changed.connect(self._update_roi_ui) - self.image_canvas.draw_mode_changed.connect(self._on_draw_mode_changed) - img_vl.addWidget(NavigationToolbar2QT(self.image_canvas, img_widget)) - img_vl.addWidget(self.image_canvas) - splitter.addWidget(img_widget) - - wave_widget = QWidget() - wave_vl = QVBoxLayout(wave_widget) - wave_vl.setContentsMargins(0, 0, 0, 0) - wave_vl.setSpacing(4) - self.lbl_wave_hint = QLabel( - "Click a pixel in the image above to inspect its waveform.") - self.lbl_wave_hint.setAlignment(Qt.AlignmentFlag.AlignCenter) - self.lbl_wave_hint.setStyleSheet(_CSS_MUTED) - self.wave_canvas = WaveformCanvas() - self.wave_canvas.setMinimumHeight(150) - wave_vl.addWidget(self.lbl_wave_hint) - wave_vl.addWidget(self.wave_canvas) - splitter.addWidget(wave_widget) - - splitter.setStretchFactor(0, 3) - splitter.setStretchFactor(1, 1) - splitter.setSizes([580, 250]) - return splitter - - def _build_right_panel(self) -> QWidget: - # Velocity settings (visible only in velocity mode) - self.grp_velocity, vel_l = _group("Velocity Settings (CH1 only)") - vel_form = _form() - self.spin_grating_um = QDoubleSpinBox() - self.spin_grating_um.setRange(0.1, 1000.0) - self.spin_grating_um.setDecimals(2) - self.spin_grating_um.setSingleStep(0.5) - self.spin_grating_um.setSuffix(" µm") - self.spin_grating_um.setValue(25) - self.spin_grating_um.setEnabled(False) - self.spin_grating_um.setMinimumWidth(_SPIN_MIN_W) - self.spin_grating_um.editingFinished.connect(self._on_grating_changed) - vel_form.addRow("Grating size:", self.spin_grating_um) - vel_l.addLayout(vel_form) - vel_l.addWidget(_wrap_label("v (m/s) = freq (MHz) × grating (µm)", - "font-size: 10px; color: #888;")) - self.grp_velocity.setVisible(False) - - grp_display, dl = _group("Display Options") - - cmap_form = _form() - self.combo_cmap = QComboBox() - self.combo_cmap.addItems(CMAPS) - self.combo_cmap.setCurrentText("gray") - self.combo_cmap.setEnabled(False) - self.combo_cmap.setSizePolicy(QSizePolicy.Policy.Expanding, - QSizePolicy.Policy.Fixed) - self.combo_cmap.currentIndexChanged.connect(self._on_cmap_changed) - cmap_form.addRow("Colormap:", self.combo_cmap) - dl.addLayout(cmap_form) - - self.chk_auto = QCheckBox("Auto-scale colormap") - self.chk_auto.setChecked(True) - self.chk_auto.toggled.connect(self._on_autoscale_toggled) - dl.addWidget(self.chk_auto) - - range_form = _form() - for label, attr in (("min:", "spin_vmin"), ("max:", "spin_vmax")): - spin = QDoubleSpinBox() - spin.setRange(-1e9, 1e9) - spin.setDecimals(4) - spin.setEnabled(False) - spin.setMinimumWidth(_SPIN_MIN_W) - spin.editingFinished.connect(self._on_manual_range_changed) - setattr(self, attr, spin) - range_form.addRow(label, spin) - dl.addLayout(range_form) - - right_panel = QWidget() - layout = QVBoxLayout(right_panel) - layout.setContentsMargins(0, 0, 0, 0) - layout.setSpacing(8) - layout.addWidget(self.grp_velocity) - layout.addWidget(grp_display) - layout.addStretch() - return _scroll_panel(right_panel, _RIGHT_PANEL_W) - - def _build_menus(self): - menubar = self.menuBar() - - fft_menu = menubar.addMenu("&FFT") - fft_act = QAction("FFT &Options…", self) - fft_act.setStatusTip("Configure FFT backend and zero-padding") - fft_act.triggered.connect(self._on_fft_options) - fft_menu.addAction(fft_act) - - fusion_menu = menubar.addMenu("&Fusion") - self._alignment_act = QAction("Angle &Alignment", self) - self._alignment_act.setStatusTip( - "Compute a rotation+translation alignment across all angles " - "(from CH4 masks) and enable Aligned View. Requires >1 angle.") - self._alignment_act.setEnabled(False) - self._alignment_act.triggered.connect(self._on_angle_alignment) - fusion_menu.addAction(self._alignment_act) - - self._manual_align_act = QAction("&Manual Alignment…", self) - self._manual_align_act.setStatusTip( - "Open an interactive dialog to align angles by eye: overlaid CH4 " - "threshold masks, keyboard nudge (translate + rotate), auto " - "de-rotate to the known scan angles, and save/clear a persistent " - "alignment.") - self._manual_align_act.setEnabled(False) - self._manual_align_act.triggered.connect(self._on_manual_alignment) - fusion_menu.addAction(self._manual_align_act) - - convert_menu = menubar.addMenu("&Convert") - self._batch_dc_act = QAction("Batch Compute DC and &Store…", self) - self._batch_dc_act.setStatusTip( - "Select .sras files and compute+store DC images (CH3/CH4 mean) " - "for every angle, converting v6 files to v7 in place.") - self._batch_dc_act.triggered.connect(lambda: self._on_batch_compute("dc")) - convert_menu.addAction(self._batch_dc_act) - - self._batch_fft_act = QAction("Batch Compute FFT and Sto&re…", self) - self._batch_fft_act.setStatusTip( - "Select .sras files and compute+store FFT peak-frequency images " - "for every angle, converting v6 files to v7 in place. Stored " - "images are natural-resolution (pad 1); padded views compute live.") - self._batch_fft_act.triggered.connect(lambda: self._on_batch_compute("fft")) - convert_menu.addAction(self._batch_fft_act) - - # ------------------------------------------------------------------ - # Drag-and-drop - # ------------------------------------------------------------------ - - def dragEnterEvent(self, event): - urls = event.mimeData().urls() - if urls and urls[0].toLocalFile().lower().endswith(".sras"): - event.acceptProposedAction() - - def dropEvent(self, event): - self._load_file(event.mimeData().urls()[0].toLocalFile()) - - # ------------------------------------------------------------------ - # File loading - # ------------------------------------------------------------------ - - def _on_open(self): - path, _ = QFileDialog.getOpenFileName( - self, "Open SRAS File", "", "SRAS Files (*.sras);;All Files (*)") - if path: - self._load_file(path) - - def _load_file(self, path: str): - started = self._run_worker( - "load", LoadWorker(path), - connect=( - ("finished", self._on_load_done), - ("error", lambda msg: self.statusBar().showMessage(f"Error: {msg}")), - ), - ) - if not started: - return - self.btn_open.setEnabled(False) - self.statusBar().showMessage(f"Loading {Path(path).name}…") - self._show_progress("main", f"Loading {Path(path).name}…") - - def _on_load_done(self, sras): - self._close_progress("main") - self.btn_open.setEnabled(True) - if sras is None: - return - self._sras = sras - self._current_image = None - - # A manual-alignment dialog bound to the previous file must not - # survive a reload — its per-angle state (and the sras it was - # constructed against) no longer matches the new file's geometry. - if self._manual_align_dialog is not None: - self._manual_align_dialog.close() - self._manual_align_dialog = None - - # Caches (and any in-flight DC precompute) belong to the previous - # file's geometry — discard and start fresh. Bumping the generation - # counters makes any still-running worker's result get dropped when - # it lands. - self._dc_cache = {} - self._fft_cache = {} - self._dc_generation += 1 - self.lbl_dc_precompute.setText("") - - self._alignment_result = None - self._aligned_cache = {} - self._alignment_generation += 1 - self.chk_aligned_view.blockSignals(True) - self.chk_aligned_view.setChecked(False) - self.chk_aligned_view.setEnabled(False) - self.chk_aligned_view.blockSignals(False) - - # Silently restore a previously-saved manual alignment, if any, so - # the work survives closing and reopening the file. - sidecar = load_manual_alignment(sras) - if sidecar is not None: - try: - self._alignment_result = build_manual_alignment( - sras, sidecar.ref_angle_idx, sidecar.dc_threshold_mv, - sidecar.per_angle) - self.chk_aligned_view.blockSignals(True) - self.chk_aligned_view.setChecked(True) - self.chk_aligned_view.blockSignals(False) - self.statusBar().showMessage( - f"Restored saved manual alignment from " - f"{sidecar_path(sras.path).name}") - except Exception as exc: - # A corrupt/foreign sidecar or a rescan that shrank n_angles - # below ref_angle_idx must not block opening the .sras file. - self.statusBar().showMessage( - f"Could not restore saved alignment: {exc}") - - # A ROI from the previous file no longer matches the new scan's - # geometry, so discard it on every load. - self.image_canvas.clear_roi() - - self.lbl_filename.setText(sras.path.name) - - self.spin_angle.blockSignals(True) - self.spin_angle.setRange(0, max(0, sras.n_angles - 1)) - self.spin_angle.setValue(0) - self.spin_angle.blockSignals(False) - - # DC channels are cheap and give an instant, fluid overview of a - # scan; CH1/Velocity require an FFT per pixel that can take minutes - # on a large scan, so don't default to it. - self.combo_channel.blockSignals(True) - self.combo_channel.setCurrentIndex(CH4_IDX) - self.combo_channel.blockSignals(False) - - self._update_controls_enabled(True) - self._on_threshold_changed() # refresh ADC label with file calibration - self._on_view_changed() - self._start_dc_precompute() - - # ------------------------------------------------------------------ - # Scan info panel - # ------------------------------------------------------------------ - - def _update_scan_info_labels(self): - s = self._sras - if s is None: - return - a = self.spin_angle.value() - for key, text in ( - ("Angles", f"{s.n_angles}"), - ("Rows", f"{s.n_rows[a]}"), - ("Frames / row", f"{s.n_frames[a]}"), - ("Samples / frame", f"{s.samples_per_frame}"), - ("Sample rate", f"{s.sample_rate_hz / 1e9:.4g} GS/s"), - ("X start", f"{s.x_start_mm[a]:.4g} mm"), - ("Pixel Δx", f"{s.pixel_x_mm * 1e3:.3g} µm"), - ("Laser freq", f"{s.laser_freq_hz / 1e3:.4g} kHz"), - ): - self._info[key].setText(f"{key}: {text}") - - notes = [] - if s.frame_count_mismatch: - notes.append(f"! Header n_frames={s.n_frames_header}, " - f"actual={s.n_frames[a]} (scanner bug — corrected)") - if s.scan_aborted: - notes.append(f"! Scan aborted: {s.n_angles}/{s.n_angles_declared} " - "angles complete") - if s.background is not None: - notes.append(f"Background waveform: {len(s.background)} samples") - if s.version in (6, 7): - notes.append("v6/v7 format: rows / frames / x_start are per-angle") - - n_dc = sum(1 for x in s.precomputed_dc4_mv if x is not None) - n_fft = sum(1 for x in s.precomputed_freq_mhz if x is not None) - if n_dc or n_fft: - bg_note = " (bg-sub)" if s.precomputed_bg_sub else " (no bg-sub)" - notes.append( - f"Cached images: DC {n_dc}/{s.n_angles} angles, " - f"FFT {n_fft}/{s.n_angles} angles{bg_note if n_fft else ''} " - "— display is instant for cached angles") - elif s.version == 7: - notes.append("v7 format: no cache blocks stored yet") - self.lbl_frame_warn.setText("\n".join(notes)) - - # ------------------------------------------------------------------ - # Controls - # ------------------------------------------------------------------ - - def _update_controls_enabled(self, enabled: bool): - s = self._sras - has_file = enabled and s is not None - ch_idx = self.combo_channel.currentIndex() - is_ch1 = enabled and ch_idx in CH1_DERIVED_MODES - is_vel = enabled and ch_idx == VELOCITY_MODE_IDX - - self.spin_angle.setEnabled(has_file and s.n_angles > 1) - self.combo_channel.setEnabled(enabled) - self.combo_cmap.setEnabled(enabled) - self.chk_auto.setEnabled(enabled) - manual = enabled and not self.chk_auto.isChecked() - self.spin_vmin.setEnabled(manual) - self.spin_vmax.setEnabled(manual) - - # Threshold and bg-sub apply to all CH1 modes - self.spin_threshold_mv.setEnabled(is_ch1) - self.chk_bg_sub.setEnabled(has_file and s.background is not None and is_ch1) - self.spin_grating_um.setEnabled(is_vel) - self.grp_velocity.setVisible(is_vel) - - self.btn_export_csv.setEnabled(is_ch1 and self._current_image is not None) - # ROI: always usable once a file is loaded (independent of channel) - self.btn_draw_roi.setEnabled(has_file) - - # Batch Convert actions pick their own files, independent of - # whatever's currently open — only gated on no batch already running. - can_batch = not self._job_running("batch") - self._batch_dc_act.setEnabled(can_batch) - self._batch_fft_act.setEnabled(can_batch) - - self._alignment_act.setEnabled( - has_file and s.n_angles > 1 and not self._job_running("align")) - self._manual_align_act.setEnabled( - has_file and s.n_angles > 1 and not self._job_running("align")) - self.chk_aligned_view.setEnabled(enabled and self._alignment_result is not None) - self._update_roi_ui() - - def _on_channel_changed(self): - self._update_controls_enabled(self._sras is not None) - self._on_view_changed() - - def _on_bg_sub_toggled(self): - # Background subtraction changes the FFT input, so it genuinely - # invalidates the cached raw FFT (the cache key includes it) — - # _refresh_display() recomputes only on a miss for the new state. - if self._sras is not None and self.combo_channel.currentIndex() in CH1_DERIVED_MODES: - self._refresh_display() - - def _on_grating_changed(self): - # Grating is a pure post-multiply on the cached frequency image — - # never needs a recompute. - if self._sras is not None and self.combo_channel.currentIndex() == VELOCITY_MODE_IDX: - self._refresh_display() - - def _on_threshold_changed(self): - mv = self.spin_threshold_mv.value() - cal = (self._sras.cal(CH4_IDX) if self._sras is not None - else (_FALLBACK_YMULT_MV, _FALLBACK_YOFF_ADC, 0.0)) - self.lbl_threshold_adc.setText(f"≈ {mv_to_adc(mv, *cal):.1f} ADC counts") - # Threshold decides which pixels get an FFT at all, so changing it is - # a genuine cache-key change — but the recompute reuses the cached DC4 - # image to skip masked-out pixels. - if self._sras is not None and self.combo_channel.currentIndex() in CH1_DERIVED_MODES: - self._refresh_display() - - def _on_autoscale_toggled(self, checked: bool): - manual = not checked - self.spin_vmin.setEnabled(manual and self._sras is not None) - self.spin_vmax.setEnabled(manual and self._sras is not None) - if self._sras is not None and self._current_image is not None: - self._redraw_image(self._current_image) - - def _on_manual_range_changed(self): - if not self.chk_auto.isChecked() and self._current_image is not None: - self._redraw_image(self._current_image) - - def _on_cmap_changed(self): - # Colormap is purely how the existing image is rendered. - if self._current_image is not None: - self._redraw_image(self._current_image) - - def _on_view_changed(self): - if self._sras is None: - return - idx = self.spin_angle.value() - self.lbl_angle_deg.setText(f"({self._sras.angles_deg[idx]:.1f}°)") - self._update_scan_info_labels() - self._refresh_display() - - def _on_aligned_view_toggled(self, checked: bool): - if self._current_image is not None: - self._redraw_image(self._current_image) - - # ------------------------------------------------------------------ - # CSV export - # ------------------------------------------------------------------ - - def _on_export_csv(self): - if self._current_image is None or self._sras is None: - return - default_name = (f"{self._sras.path.stem}_angle{self._current_angle}" - f"_{CH_NAMES[self._current_ch]}.csv") - path, _ = QFileDialog.getSaveFileName( - self, "Export Image as CSV", - str(self._sras.path.parent / default_name), - "CSV files (*.csv);;All files (*)") - if not path: - return - np.savetxt(path, self._current_image, delimiter=",", fmt="%.6g") - self.statusBar().showMessage(f"Exported {Path(path).name}") - - def _on_export_roi_csv(self): - if self._current_image is None or self._sras is None: - return - roi = self.image_canvas.get_roi() - if roi is None: - self.statusBar().showMessage("No ROI — draw one first") - return - s = self._sras - x_axis = s.x_axis_mm(self._current_angle) - y_axis = s.y_positions_mm(self._current_angle) - mask = roi.mask_for_grid(x_axis, y_axis) - if not mask.any(): - self.statusBar().showMessage("ROI does not overlap any pixel") - return - img = self._current_image - if img.shape != mask.shape: - self.statusBar().showMessage( - f"ROI shape {mask.shape} does not match image {img.shape}") - return - - X, Y = np.meshgrid(np.asarray(x_axis, dtype=np.float64), - np.asarray(y_axis, dtype=np.float64)) - rows_idx, frames_idx = np.where(mask) - n_pix = int(mask.sum()) - - ch_name = CH_NAMES[self._current_ch] - angle = self._current_angle - default_name = f"{s.path.stem}_angle{angle}_{ch_name}_ROI.csv" - path, _ = QFileDialog.getSaveFileName( - self, "Export ROI as CSV", - str(s.path.parent / default_name), - "CSV files (*.csv);;All files (*)") - if not path: - return - - corners_str = " ".join(f"({p[0]:.6g},{p[1]:.6g})" for p in roi.corners()) - header = ( - f"# ROI quad corners (BL BR TR TL) mm: {corners_str}\n" - f"# source: {s.path.name}, channel={ch_name}, " - f"angle_idx={angle}, angle_deg={s.angles_deg[angle]:.4g}\n" - f"# n_pixels={n_pix}\n" - "row,frame,x_mm,y_mm,value" - ) - data = np.column_stack([ - rows_idx.astype(np.int64), frames_idx.astype(np.int64), - X[mask], Y[mask], img[mask].astype(np.float64), - ]) - # integer columns first, floats after — use a per-column format list - np.savetxt(path, data, delimiter=",", - fmt=["%d", "%d", "%.6g", "%.6g", "%.6g"], - header=header, comments="") - self.statusBar().showMessage( - f"Exported ROI ({n_pix} pixels) to {Path(path).name}") - - # ------------------------------------------------------------------ - # ROI - # ------------------------------------------------------------------ - - def _on_draw_roi_toggled(self, checked: bool): - if checked: - self.image_canvas.start_drawing() - self.statusBar().showMessage( - "Click and drag on the image to draw a new rectangle.") - else: - self.image_canvas.cancel_drawing() - - def _on_draw_mode_changed(self, active: bool): - # Keep the toggle button's visual state in sync with the canvas. - self.btn_draw_roi.blockSignals(True) - self.btn_draw_roi.setChecked(active) - self.btn_draw_roi.blockSignals(False) - - def _on_clear_roi(self): - self.image_canvas.clear_roi() - self.statusBar().showMessage("ROI cleared") - - def _update_roi_ui(self): - roi = self.image_canvas.get_roi() - if roi is None: - self.lbl_roi_center.setText("centroid: —") - self.lbl_roi_size.setText("bbox: —") - self.lbl_roi_npix.setText("pixels inside: —") - self.btn_clear_roi.setEnabled(False) - self.btn_export_roi.setEnabled(False) - return - - cen = roi.centroid() - bbox = roi.bbox_size() - self.lbl_roi_center.setText(f"centroid: ({cen[0]:.3f}, {cen[1]:.3f}) mm") - self.lbl_roi_size.setText(f"bbox: {bbox[0]:.3f} × {bbox[1]:.3f} mm") - - npix = 0 - if self._sras is not None: - try: - # Deliberately always the raw per-angle grid, even when - # Aligned View is on: _on_export_roi_csv also exports on the - # raw grid (never synthetically-resampled pixels), so this - # readout must match what Export ROI actually writes. - mask = roi.mask_for_grid( - self._sras.x_axis_mm(self._current_angle), - self._sras.y_positions_mm(self._current_angle)) - npix = int(mask.sum()) - except Exception: - npix = 0 - self.lbl_roi_npix.setText(f"pixels inside: {npix}") - self.btn_clear_roi.setEnabled(True) - self.btn_export_roi.setEnabled(self._current_image is not None and npix > 0) - - # ------------------------------------------------------------------ - # Display - # ------------------------------------------------------------------ - - def _current_n_fft(self) -> int | None: - if self._fft_pad_factor <= 1 or self._sras is None: - return None - return self._sras.samples_per_frame * self._fft_pad_factor - - def _scale_for_display(self, freq_mhz: np.ndarray, ch_idx: int) -> np.ndarray: - """Velocity is a pure post-multiply of the (already DC-masked) - cached frequency image — never worth a recompute on its own.""" - if ch_idx == VELOCITY_MODE_IDX: - return freq_mhz * self.spin_grating_um.value() - return freq_mhz - - def _fft_cache_key(self, angle_idx: int) -> tuple: - return (angle_idx, self.chk_bg_sub.isChecked(), self._current_n_fft(), - self.spin_threshold_mv.value()) - - def _aligned_cache_key(self, angle_idx: int, ch_idx: int) -> tuple: - """Mirrors _fft_cache's key granularity so a stale aligned image is - never shown after bg_sub/threshold/pad/grating changes.""" - if ch_idx in CH1_DERIVED_MODES: - return (*self._fft_cache_key(angle_idx), ch_idx, - self.spin_grating_um.value() if ch_idx == VELOCITY_MODE_IDX else None) - return (angle_idx, ch_idx) - - def _aligned_canvas_axes(self) -> tuple[np.ndarray, np.ndarray]: - r = self._alignment_result - n_rows, n_cols = r.canvas_shape - return (r.canvas_origin_mm[0] + np.arange(n_cols) * r.canvas_dx_mm, - r.canvas_origin_mm[1] + np.arange(n_rows) * r.canvas_dy_mm) - - def _get_aligned_display_image(self, raw_img: np.ndarray, angle_idx: int, - ch_idx: int) -> np.ndarray: - key = self._aligned_cache_key(angle_idx, ch_idx) - cached = self._aligned_cache.get(key) - if cached is None: - cached = apply_alignment(self._alignment_result, angle_idx, raw_img) - self._aligned_cache[key] = cached - return cached - - def _refresh_display(self): - """Show the image for the current angle/channel/threshold, using - cached data whenever possible and only falling back to a background - compute (with progress popup) when genuinely nothing is cached yet.""" - if self._sras is None: - return - angle_idx = self.spin_angle.value() - ch_idx = self.combo_channel.currentIndex() - - if ch_idx in CH1_DERIVED_MODES: - raw = self._fft_cache.get(self._fft_cache_key(angle_idx)) - if raw is not None: - self._show_image_now(self._scale_for_display(raw, ch_idx), - angle_idx, ch_idx) - return - else: - cached = self._dc_cache.get((angle_idx, ch_idx)) - if cached is not None: - self._show_image_now(cached, angle_idx, ch_idx) - return - - # Nothing cached for these settings — need a real compute. Changing - # the DC threshold changes *which* pixels get an FFT at all, so it - # can't be satisfied from the cache — but with the DC map already - # known, the recompute skips the FFT for masked-out pixels. - self._start_compute() - - def _show_image_now(self, img: np.ndarray, angle_idx: int, ch_idx: int): - """Display an already-available image with no compute involved.""" - self._current_image = img - self._current_angle = angle_idx - self._current_ch = ch_idx - self.btn_export_csv.setEnabled(ch_idx in CH1_DERIVED_MODES) - self._redraw_image(img) - self._update_roi_ui() - - def _redraw_image(self, img: np.ndarray): - s = self._sras - angle_idx = self._current_angle - ch_idx = self._current_ch - - aligned = (self.chk_aligned_view.isChecked() - and self._alignment_result is not None - and angle_idx in self._alignment_result.per_angle) - if aligned: - display_img = self._get_aligned_display_image(img, angle_idx, ch_idx) - x_axis, y_axis = self._aligned_canvas_axes() - else: - display_img = img - x_axis = s.x_axis_mm(angle_idx) - y_axis = s.y_positions_mm(angle_idx) - - dx = x_axis[1] - x_axis[0] if len(x_axis) > 1 else s.pixel_x_mm - dy = float(y_axis[1] - y_axis[0]) if len(y_axis) > 1 else 1.0 - extent = [x_axis[0] - dx / 2, x_axis[-1] + dx / 2, - y_axis[-1] + dy / 2, y_axis[0] - dy / 2] - - if self.chk_auto.isChecked(): - vmin, vmax = float(display_img.min()), float(display_img.max()) - for spin, val in ((self.spin_vmin, vmin), (self.spin_vmax, vmax)): - spin.blockSignals(True) - spin.setValue(val) - spin.blockSignals(False) - else: - vmin, vmax = self.spin_vmin.value(), self.spin_vmax.value() - - angle_deg = s.angles_deg[angle_idx] - mode_str, unit, colorbar_label = _CHANNEL_DISPLAY[ch_idx] - if ch_idx == VELOCITY_MODE_IDX: - ch_label = f"Velocity [grating={self.spin_grating_um.value():.2f} µm]" - else: - ch_label = CH_LABELS[ch_idx] - - title = f"{CH_NAMES[ch_idx]} | {mode_str} | {angle_deg:.1f}°" - if aligned: - title += " [Aligned]" - - self.image_canvas.show_image( - display_img, extent, - cmap=self.combo_cmap.currentText(), - vmin=vmin, vmax=vmax, - xlabel="X (mm)", ylabel="Y (mm)", - title=title, colorbar_label=colorbar_label, - ) - self.statusBar().showMessage( - f"{s.path.name} | {ch_label} @ {angle_deg:.1f}° " - f"| {display_img.shape[1]} × {display_img.shape[0]} px | {unit}" - f"{' | Aligned' if aligned else ''}" - ) - - # ------------------------------------------------------------------ - # Background compute (only reached on a genuine cache miss) - # ------------------------------------------------------------------ - - def _start_compute(self): - if self._sras is None or self._job_running("compute"): - return # re-checked when the running compute finishes - - angle_idx = self.spin_angle.value() - ch_idx = self.combo_channel.currentIndex() - is_fft = ch_idx in CH1_DERIVED_MODES - - self._pending_angle = angle_idx - self._pending_ch = ch_idx - self._pending_bg_sub = self.chk_bg_sub.isChecked() - self._pending_threshold = self.spin_threshold_mv.value() - self._pending_fft_pad_factor = self._fft_pad_factor - - worker = ComputeWorker( - self._sras, angle_idx, ch_idx, - apply_bg_sub=self._pending_bg_sub, - n_fft=self._current_n_fft(), - dc_threshold_mv=self._pending_threshold, - # Reuse the cached DC4 image (if the precompute has reached this - # angle) so the FFT skips masked-out pixels entirely and doesn't - # need to re-read the CH4 channel from disk. - dc4_mv=self._dc_cache.get((angle_idx, CH4_IDX)), - is_fft_mode=is_fft, - ) - if not self._run_worker( - "compute", worker, - connect=( - ("finished", self._on_compute_done), - ("error", lambda msg: self.statusBar().showMessage( - f"Compute error: {msg}")), - ), - on_done=self._after_compute): - return - - if is_fft: - self.statusBar().showMessage("Computing FFT…") - self._show_progress( - "main", - f"Computing FFT for angle {angle_idx}…\n" - "This can take a while on a large scan — result is cached " - "so revisiting this angle/mode/threshold will be instant.") - else: - self.statusBar().showMessage("Computing DC image…") - self._show_progress("main", f"Computing DC image for angle {angle_idx}…") - - def _after_compute(self): - """If settings changed while the compute was running, re-dispatch - through the cache-aware path — the now-current combination may - already be cached.""" - if (self.spin_angle.value(), self.combo_channel.currentIndex(), - self.chk_bg_sub.isChecked(), self.spin_threshold_mv.value(), - self._fft_pad_factor) != ( - self._pending_angle, self._pending_ch, self._pending_bg_sub, - self._pending_threshold, self._pending_fft_pad_factor): - self._refresh_display() - - def _on_compute_done(self, result): - self._close_progress("main") - if result is None: - return # cancelled mid-compute; the partial image must not cache - angle_idx = self._pending_angle - ch_idx = self._pending_ch - - if ch_idx in CH1_DERIVED_MODES: - self._fft_cache[(angle_idx, self._pending_bg_sub, - self._current_n_fft(), self._pending_threshold)] = result - img = self._scale_for_display(result, ch_idx) - else: - img = result - self._dc_cache[(angle_idx, ch_idx)] = img - - self._show_image_now(img, angle_idx, ch_idx) - - # ------------------------------------------------------------------ - # Background DC precompute (all angles, so switching is fluid) - # ------------------------------------------------------------------ - - def _start_dc_precompute(self): - if self._sras is None: - return - generation = self._dc_generation - n_angles = self._sras.n_angles - - worker = DcPrecomputeWorker(self._sras) - self._run_worker( - "dc_precompute", worker, - connect=( - ("angle_done", lambda a, dc3, dc4, g=generation: - self._on_dc_precompute_angle_done(g, a, dc3, dc4, n_angles)), - ("error", lambda msg: self.statusBar().showMessage( - f"DC precompute error: {msg}", 5000)), - ), - quit_on=("finished", "error"), - ) - - def _on_dc_precompute_angle_done(self, generation: int, angle_idx: int, - dc3_mv: np.ndarray, dc4_mv: np.ndarray, - n_angles: int): - if generation != self._dc_generation: - return # stale result from a previously-loaded file — discard - self._dc_cache[(angle_idx, CH3_IDX)] = dc3_mv - self._dc_cache[(angle_idx, CH4_IDX)] = dc4_mv - - done = sum(1 for a in range(n_angles) if (a, CH4_IDX) in self._dc_cache) - self.lbl_dc_precompute.setText( - f"Precomputing DC images: {done}/{n_angles} angles ready…" - if done < n_angles else "DC images ready for all angles.") - - # If we just finished the angle/channel the user is currently looking - # at and it wasn't shown yet (they switched here before the precompute - # caught up and are still waiting), show it now. - current_ch = self.combo_channel.currentIndex() - if (angle_idx == self.spin_angle.value() - and not self._job_running("compute") - and current_ch in (CH3_IDX, CH4_IDX) - and (self._current_angle != angle_idx or self._current_ch != current_ch)): - self._refresh_display() - - # ------------------------------------------------------------------ - # Pixel inspector - # ------------------------------------------------------------------ - - def _on_pixel_clicked(self, row_idx: int, frame_idx: int): - if self._sras is None or self._current_image is None: - return - angle_idx = self._current_angle - if (self.chk_aligned_view.isChecked() and self._alignment_result is not None - and angle_idx in self._alignment_result.per_angle): - # The click landed on the shared aligned canvas — invert the same - # canvas->raw affine used to display it back to a raw (row, frame) - # index before looking up the waveform. - t = self._alignment_result.per_angle[angle_idx] - raw = t.matrix @ np.array([row_idx, frame_idx], dtype=np.float64) + t.offset - row_idx, frame_idx = int(round(raw[0])), int(round(raw[1])) - n_rows_a, n_frames_a = self._sras.image_shape(angle_idx) - if not (0 <= row_idx < n_rows_a and 0 <= frame_idx < n_frames_a): - self.statusBar().showMessage( - "No source waveform here (padding region of the aligned canvas).") - return - - self.lbl_wave_hint.hide() - if self._current_ch in CH1_DERIVED_MODES: - self.wave_canvas.show_rf_waveform( - self._sras, angle_idx, row_idx, frame_idx, - apply_bg_sub=self.chk_bg_sub.isChecked()) - else: - self.wave_canvas.show_dc_waveform( - self._sras, angle_idx, self._current_ch, row_idx, frame_idx) - - # ------------------------------------------------------------------ - # Progress dialogs - # ------------------------------------------------------------------ - - def _show_progress(self, key: str, message: str, maximum: int = 0): - """Show (or relabel) the progress dialog under *key*. maximum=0 gives - an indeterminate busy indicator.""" - dlg = self._progress_dlgs.get(key) - if dlg is not None: - dlg.setLabelText(message) - return - dlg = QProgressDialog(message, "", 0, maximum, self) - dlg.setWindowTitle("Please wait…") - dlg.setCancelButton(None) - dlg.setWindowModality(Qt.WindowModality.WindowModal) - dlg.setMinimumDuration(300) # only appears if it takes > 300 ms - dlg.show() - self._progress_dlgs[key] = dlg - - def _set_progress(self, key: str, pct: int): - dlg = self._progress_dlgs.get(key) - if dlg is not None: - dlg.setValue(pct) - - def _close_progress(self, key: str): - dlg = self._progress_dlgs.pop(key, None) - if dlg is not None: - dlg.close() - - # ------------------------------------------------------------------ - # Convert menu: batch DC/FFT compute-and-store (v6 -> v7) - # ------------------------------------------------------------------ - - def _on_batch_compute(self, mode: str): - if self._job_running("batch"): - return - label = "DC" if mode == "dc" else "FFT" - paths, _ = QFileDialog.getOpenFileNames( - self, f"Select .sras files to batch-compute {label}", "", - "SRAS files (*.sras);;All files (*)") - if not paths: - return - - self._batch_errors = [] - worker = BatchCacheWorker(paths, mode, self.chk_bg_sub.isChecked()) - started = self._run_worker( - "batch", worker, - connect=( - ("progress", lambda pct: self._set_progress("batch", pct)), - ("file_done", self._on_batch_file_done), - ("finished", lambda p=paths: self._on_batch_finished(p)), - ), - on_done=self._after_batch, - ) - if not started: - return # a second trigger snuck in while the file dialog was open - - self._batch_dc_act.setEnabled(False) - self._batch_fft_act.setEnabled(False) - self._show_progress( - "batch", f"Batch computing {label} for {len(paths)} file(s)…", - maximum=100) - - def _on_batch_file_done(self, path: str, err: str): - if err: - self._batch_errors.append(f"{Path(path).name} — {err}") - self._show_progress("batch", f"Processed {Path(path).name}…") - - def _on_batch_finished(self, paths: list[str]): - self._close_progress("batch") - - n_total = len(paths) - n_failed = len(self._batch_errors) - n_ok = n_total - n_failed - if n_failed: - summary = (f"Batch store: {n_ok}/{n_total} file(s) updated, " - f"{n_failed} failed: {'; '.join(self._batch_errors)}") - else: - summary = f"Batch store: {n_ok}/{n_total} file(s) updated." - self.statusBar().showMessage(summary) - self._batch_errors = [] - - # If the currently-open file was in this batch, reload it so the GUI - # picks up the newly-written v7 cache instead of stale state. - if self._sras is not None and str(self._sras.path) in paths: - self._load_file(str(self._sras.path)) - - def _after_batch(self): - self._batch_dc_act.setEnabled(True) - self._batch_fft_act.setEnabled(True) - - # ------------------------------------------------------------------ - # Fusion: angle alignment - # ------------------------------------------------------------------ - - def _on_angle_alignment(self): - if self._sras is None or self._sras.n_angles <= 1: - return - ref_idx = 0 - threshold_mv = self.spin_threshold_mv.value() - generation = self._alignment_generation - - started = self._run_worker( - "align", AngleAlignmentWorker(self._sras, ref_idx, threshold_mv), - connect=( - ("progress", lambda pct: self._set_progress("main", pct)), - ("finished", lambda result, err, g=generation: - self._on_alignment_done(g, result, err)), - ), - on_done=lambda: self._update_controls_enabled(self._sras is not None), - ) - if not started: - return - - self._alignment_act.setEnabled(False) - self._show_progress( - "main", - f"Computing angle alignment ({self._sras.n_angles} angles, " - f"ref=angle 0, CH4 mask ≥ {threshold_mv:.3f} mV)…", - maximum=100) - - def _on_alignment_done(self, generation: int, result, error_msg: str): - self._close_progress("main") - if generation != self._alignment_generation: - return # a new file was loaded while this was computing — discard - if error_msg: - self.statusBar().showMessage(f"Angle alignment failed: {error_msg}") - return - self._alignment_result = result - self._aligned_cache = {} - self.chk_aligned_view.setEnabled(True) - self.chk_aligned_view.blockSignals(True) - self.chk_aligned_view.setChecked(True) - self.chk_aligned_view.blockSignals(False) - nr, nc = result.canvas_shape - self.statusBar().showMessage( - f"Angle alignment computed ({self._sras.n_angles} angles, " - f"canvas {nc}×{nr} px).") - self._refresh_display() - - # ------------------------------------------------------------------ - # Fusion: manual alignment - # ------------------------------------------------------------------ - - def _on_manual_alignment(self): - if self._sras is None or self._sras.n_angles <= 1: - return - if self._manual_align_dialog is not None: - self._manual_align_dialog.raise_() - self._manual_align_dialog.activateWindow() - return - - ref_idx = 0 - threshold_mv = self.spin_threshold_mv.value() - seed: dict[int, ManualAngleParams] = {} - # Seed only from a previously *saved manual* alignment (this dialog's - # own Save also writes this sidecar) -- never from self._alignment_result - # when it holds the automatic Fusion -> Angle Alignment's output. That - # path's translation comes from FFT phase correlation, which is the - # very thing manual mode exists to work around; inheriting it here - # would silently reintroduce the same bad translations under a - # "manual" label, on top of the (correct) analytic rotation, which is - # exactly what makes manual mode look like it "still does the same - # thing" the automatic one does. - sidecar = load_manual_alignment(self._sras) - if sidecar is not None and sidecar.ref_angle_idx == ref_idx: - seed = dict(sidecar.per_angle) - threshold_mv = sidecar.dc_threshold_mv - - cached_dc4 = {a: img for (a, ch), img in self._dc_cache.items() if ch == CH4_IDX} - dlg = ManualAlignmentDialog( - self, self._sras, ref_angle_idx=ref_idx, dc_threshold_mv=threshold_mv, - seed_per_angle=seed, cached_dc4_mv=cached_dc4) - dlg.alignment_saved.connect(self._on_manual_alignment_saved) - dlg.alignment_cleared.connect(self._on_manual_alignment_cleared) - dlg.finished.connect(self._on_manual_align_dialog_closed) - dlg.setAttribute(Qt.WidgetAttribute.WA_DeleteOnClose) - self._manual_align_dialog = dlg - dlg.show() - - def _on_manual_align_dialog_closed(self, _result_code: int): - self._manual_align_dialog = None - - def _on_manual_alignment_saved(self, result, sidecar_path_str: str): - self._alignment_result = result - self._aligned_cache = {} - self._alignment_generation += 1 - self.chk_aligned_view.setEnabled(True) - self.chk_aligned_view.blockSignals(True) - self.chk_aligned_view.setChecked(True) - self.chk_aligned_view.blockSignals(False) - self._update_controls_enabled(self._sras is not None) - self.statusBar().showMessage( - f"Manual alignment saved to {Path(sidecar_path_str).name}") - if self._current_image is not None: - self._refresh_display() - - def _on_manual_alignment_cleared(self): - self._alignment_result = None - self._aligned_cache = {} - self._alignment_generation += 1 - self.chk_aligned_view.blockSignals(True) - self.chk_aligned_view.setChecked(False) - self.chk_aligned_view.setEnabled(False) - self.chk_aligned_view.blockSignals(False) - self._update_controls_enabled(self._sras is not None) - self.statusBar().showMessage("Manual alignment cleared.") - if self._current_image is not None: - self._refresh_display() - - # ------------------------------------------------------------------ - # FFT Options - # ------------------------------------------------------------------ - - def _on_fft_options(self): - dlg = FftOptionsDialog( - self, - current_backend=compute.get_fft_backend(), - current_pad_factor=self._fft_pad_factor, - samples_per_frame=self._sras.samples_per_frame if self._sras else None, - sample_rate_hz=self._sras.sample_rate_hz if self._sras else None, - grating_um=self.spin_grating_um.value(), - ) - if dlg.exec() != QDialog.DialogCode.Accepted: - return - compute.set_fft_backend(dlg.get_backend()) - self._fft_pad_factor = dlg.get_pad_factor() - self._settings.setValue("fft/backend", compute.get_fft_backend()) - self._settings.setValue("fft/pad_factor", self._fft_pad_factor) - # Pad factor changes the FFT bin count, so it genuinely invalidates - # the cached raw FFT (part of the cache key) — _refresh_display() - # recomputes only on a cache miss. - if self._sras is not None and self.combo_channel.currentIndex() in CH1_DERIVED_MODES: - self._refresh_display() - - # ------------------------------------------------------------------ - - def closeEvent(self, event): - if self._manual_align_dialog is not None: - self._manual_align_dialog.close() - - # Signal every cancellable worker first, then wait. Waiting without - # signalling means sitting out whatever is in flight — on a large - # scan a single angle is ~40 s. - jobs = list(self._jobs.values()) - for _thread, worker, _on_done in jobs: - stop = getattr(worker, "stop", None) - if callable(stop): - stop() - for thread, _worker, _on_done in jobs: - thread.quit() - thread.wait(5000) - super().closeEvent(event) - - -# --------------------------------------------------------------------------- - -def main(): - app = QApplication(sys.argv) - window = SrasViewerWindow( - initial_path=sys.argv[1] if len(sys.argv) > 1 else None) - window.show() - sys.exit(app.exec()) - - -if __name__ == "__main__": - main() diff --git a/sras_viewer/__init__.py b/sras_viewer/__init__.py new file mode 100644 index 0000000..b218d72 --- /dev/null +++ b/sras_viewer/__init__.py @@ -0,0 +1,24 @@ +""" +SRAS Scan File Viewer +PyQt6 application for visualizing channel data from .sras binary scan files. + +Channel semantics (fixed by sc3_aui_app.py acquisition settings): + CH1 — RF Acoustic Packet (AC-coupled, 100 mV/div): FFT → peak frequency + CH3 — Bias A (DC-coupled, 50 mV/div): waveform mean + CH4 — Bias B (DC-coupled, 50 mV/div): waveform mean + +RF images are masked: pixels where CH4_dc < dc_threshold show 0. + +File parsing lives in sras_format, image/alignment math in sras_compute, and +background workers in sras_workers — none of which import Qt or matplotlib, +so multiprocessing children can load them cheaply. +""" + +import faulthandler + +faulthandler.enable() # print a native stack trace on SIGSEGV/SIGABRT/etc. + +from .canvases import ImageCanvas, RoiQuad, WaveformCanvas # noqa: E402,F401 +from .common import CH_LABELS, CMAPS, VELOCITY_MODE_IDX # noqa: E402,F401 +from .dialogs import FftOptionsDialog, ManualAlignmentDialog # noqa: E402,F401 +from .main_window import SrasViewerWindow, main # noqa: E402,F401 diff --git a/sras_viewer/__main__.py b/sras_viewer/__main__.py new file mode 100644 index 0000000..cbd143d --- /dev/null +++ b/sras_viewer/__main__.py @@ -0,0 +1,4 @@ +from .main_window import main + +if __name__ == "__main__": + main() diff --git a/sras_viewer/canvases.py b/sras_viewer/canvases.py new file mode 100644 index 0000000..f7b5e83 --- /dev/null +++ b/sras_viewer/canvases.py @@ -0,0 +1,542 @@ +"""Matplotlib canvases and the ROI primitive.""" + +import numpy as np +from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg +from matplotlib.figure import Figure +from matplotlib.patches import Polygon +from matplotlib.path import Path as MplPath +from PyQt6.QtCore import Qt, pyqtSignal +from PyQt6.QtGui import QKeyEvent +from PyQt6.QtWidgets import QSizePolicy + +from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, CH_NAMES, SrasFile, adc_to_mv + +# --------------------------------------------------------------------------- +# ROI (free quadrilateral in data coordinates) +# --------------------------------------------------------------------------- + +class RoiQuad: + """Free quadrilateral defined in data coordinates (mm). + + Stored as 4 corner points (shape (4, 2)) in CCW order: BL, BR, TR, TL. + Each corner can be positioned independently, allowing skewed / + non-orthogonal regions of interest. Because it lives in scan/data + coords it persists unchanged when the displayed channel/mode switches. + """ + + def __init__(self, pts: np.ndarray): + """pts : array-like, shape (4, 2).""" + self._pts = np.asarray(pts, dtype=np.float64).reshape(4, 2).copy() + + @classmethod + def from_bbox(cls, x0: float, y0: float, x1: float, y1: float) -> "RoiQuad": + """Create an axis-aligned rectangle from two opposite corners.""" + lx, rx = min(x0, x1), max(x0, x1) + by, ty = min(y0, y1), max(y0, y1) + return cls(np.array([[lx, by], [rx, by], [rx, ty], [lx, ty]])) + + def copy(self) -> "RoiQuad": + return RoiQuad(self._pts.copy()) + + def corners(self) -> np.ndarray: + """World-coord corners, shape (4, 2), CCW: BL, BR, TR, TL.""" + return self._pts.copy() + + def centroid(self) -> np.ndarray: + return self._pts.mean(axis=0) + + def bbox_size(self) -> np.ndarray: + """Width and height of the axis-aligned bounding box, shape (2,).""" + return self._pts.max(axis=0) - self._pts.min(axis=0) + + def contains(self, x: float, y: float) -> bool: + return bool(MplPath(self._pts).contains_point((x, y))) + + def mask_for_grid(self, x_axis: np.ndarray, + y_axis: np.ndarray) -> np.ndarray: + """Boolean mask (n_rows, n_frames) of pixels whose centres lie + inside the quadrilateral. + + Only the quad's axis-aligned bounding box is tested — meshgrid and + contains_points over the *whole* grid would be tens of millions of + point-in-polygon tests (and hundreds of MB of float64 temporaries) + on a large scan, on every ROI edit. + """ + x = np.asarray(x_axis, dtype=np.float64) + y = np.asarray(y_axis, dtype=np.float64) + mask = np.zeros((y.size, x.size), dtype=bool) + + (x0, y0), (x1, y1) = self._pts.min(axis=0), self._pts.max(axis=0) + cols = np.nonzero((x >= x0) & (x <= x1))[0] + rows = np.nonzero((y >= y0) & (y <= y1))[0] + if cols.size == 0 or rows.size == 0: + return mask + + c0, c1 = int(cols[0]), int(cols[-1]) + 1 + r0, r1 = int(rows[0]), int(rows[-1]) + 1 + X, Y = np.meshgrid(x[c0:c1], y[r0:r1]) + inside = MplPath(self._pts).contains_points( + np.column_stack([X.ravel(), Y.ravel()])) + mask[r0:r1, c0:c1] = inside.reshape(X.shape) + return mask + + +# --------------------------------------------------------------------------- +# Matplotlib canvases +# --------------------------------------------------------------------------- + +class ImageCanvas(FigureCanvasQTAgg): + pixel_clicked = pyqtSignal(int, int) # row_idx, frame_idx + roi_changed = pyqtSignal() # ROI created / edited / cleared + draw_mode_changed = pyqtSignal(bool) # "draw new ROI" arm toggled + + # Interaction state values + _IDLE = "idle" + _DRAW_NEW = "draw_new" + _MOVE = "move" + _DRAG_CORNER = "drag_corner" + + # Hit tolerance (display pixels) for handles. + _HANDLE_PX = 12 + _CLICK_THRESH_PX = 4 # releases within this of press count as a click + + def __init__(self, parent=None): + fig = Figure(figsize=(7, 5), tight_layout=True) + self.ax = fig.add_subplot(111) + super().__init__(fig) + self.setParent(parent) + self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding) + self._extent = None + self._img_shape = None + + # ROI state + self._roi: RoiQuad | None = None + self._roi_artists: list = [] + self._state = self._IDLE + self._draw_mode = False + + # Per-interaction snapshots / anchors + self._press_xy: tuple[float, float] | None = None + self._press_pixel: tuple[float, float] | None = None + self._press_button = None + self._snapshot: RoiQuad | None = None + self._drag_corner_idx: int = -1 + self._move_anchor = None # press-point in world coords + self._draw_previous: RoiQuad | None = None + + self.mpl_connect("button_press_event", self._on_press) + self.mpl_connect("motion_notify_event", self._on_motion) + self.mpl_connect("button_release_event", self._on_release) + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def show_image(self, img: np.ndarray, extent: list[float], cmap: str, + vmin: float, vmax: float, xlabel: str, ylabel: str, title: str, + colorbar_label: str = ""): + self.figure.clf() + self.ax = self.figure.add_subplot(111) + # Patches and lines are destroyed by figure.clf(); drop stale refs. + self._roi_artists = [] + + self._extent = extent + self._img_shape = img.shape + + im = self.ax.imshow( + img, aspect="auto", origin="upper", + extent=extent, cmap=cmap, vmin=vmin, vmax=vmax, + interpolation="nearest", + ) + cb = self.figure.colorbar(im, ax=self.ax, fraction=0.046, pad=0.04) + if colorbar_label: + cb.set_label(colorbar_label) + + self.ax.set_xlabel(xlabel) + self.ax.set_ylabel(ylabel) + self.ax.set_title(title) + + # Re-draw the ROI (if any) on top of the fresh image so it persists + # unchanged across mode / angle / channel switches. + self._draw_roi() + self.draw() + + def get_roi(self) -> RoiQuad | None: + return self._roi + + def set_roi(self, roi: RoiQuad | None): + self._roi = roi.copy() if roi is not None else None + self._draw_roi() + self.draw_idle() + self.roi_changed.emit() + + def clear_roi(self): + self._roi = None + self._remove_roi_artists() + self.draw_idle() + self.roi_changed.emit() + + def start_drawing(self): + """Arm the next click+drag on the image to create a new ROI, + replacing any existing one.""" + self._draw_mode = True + self.setCursor(Qt.CursorShape.CrossCursor) + self.draw_mode_changed.emit(True) + + def cancel_drawing(self): + if self._draw_mode: + self._draw_mode = False + self.setCursor(Qt.CursorShape.ArrowCursor) + self.draw_mode_changed.emit(False) + + # ------------------------------------------------------------------ + # Rendering + # ------------------------------------------------------------------ + + def _remove_roi_artists(self): + for a in self._roi_artists: + try: + a.remove() + except (ValueError, AttributeError, NotImplementedError): + pass + self._roi_artists = [] + + def _draw_roi(self): + self._remove_roi_artists() + if self._roi is None or self.ax is None: + return + corners = self._roi.corners() + + # Filled quad, then a sharp unfilled edge for visibility over bright + # images, then draggable corner handles. + for kwargs in ( + dict(fill=True, facecolor="#ffd93a", edgecolor="#e53935", + alpha=0.22, linewidth=2.0, zorder=10), + dict(fill=False, edgecolor="#e53935", linewidth=1.8, zorder=11), + ): + patch = Polygon(corners, closed=True, **kwargs) + self.ax.add_patch(patch) + self._roi_artists.append(patch) + + self._roi_artists.append(self.ax.scatter( + corners[:, 0], corners[:, 1], s=60, c="white", + edgecolors="#e53935", linewidths=1.6, zorder=13)) + + # ------------------------------------------------------------------ + # Hit testing (display pixels for handles, data coords for "inside") + # ------------------------------------------------------------------ + + def _hit_test(self, event) -> tuple[str, int | None] | None: + if self._roi is None or self.ax is None: + return None + if event.x is None or event.y is None: + return None + corners_disp = self.ax.transData.transform(self._roi.corners()) + click = np.array([event.x, event.y]) + + for i in range(4): + if np.hypot(*(corners_disp[i] - click)) <= self._HANDLE_PX: + return ("corner", i) + + if event.xdata is not None and event.ydata is not None: + if self._roi.contains(event.xdata, event.ydata): + return ("inside", None) + return None + + # ------------------------------------------------------------------ + # Mouse event handlers + # ------------------------------------------------------------------ + + def _on_press(self, event): + if event.inaxes is not self.ax or self._extent is None: + return + if event.button != 1: # only left mouse button + return + # If the matplotlib toolbar is in pan / zoom mode, let it handle + # the interaction instead of starting a ROI manipulation. + tb = getattr(self, "toolbar", None) + if tb is not None and getattr(tb, "mode", ""): + return + + self._press_xy = (event.xdata, event.ydata) + self._press_pixel = (event.x, event.y) + self._press_button = event.button + + if self._draw_mode: + self._draw_previous = self._roi.copy() if self._roi else None + self._roi = RoiQuad.from_bbox(event.xdata, event.ydata, + event.xdata, event.ydata) + self._state = self._DRAW_NEW + self._draw_roi() + self.draw_idle() + return + + hit = self._hit_test(event) + if hit is None: + self._state = self._IDLE + return + + kind, idx = hit + self._snapshot = self._roi.copy() + if kind == "corner": + self._state = self._DRAG_CORNER + self._drag_corner_idx = idx + else: + self._state = self._MOVE + self._move_anchor = (event.xdata, event.ydata) + + def _on_motion(self, event): + if self._state == self._IDLE: + return + if event.xdata is None or event.ydata is None: + return + if event.inaxes is not self.ax: + return + + if self._state == self._DRAW_NEW: + x0, y0 = self._press_xy + self._roi = RoiQuad.from_bbox(x0, y0, event.xdata, event.ydata) + elif self._state == self._MOVE: + delta = np.array([event.xdata - self._move_anchor[0], + event.ydata - self._move_anchor[1]]) + self._roi._pts = self._snapshot.corners() + delta + elif self._state == self._DRAG_CORNER: + self._roi._pts[self._drag_corner_idx] = [event.xdata, event.ydata] + + self._draw_roi() + self.draw_idle() + + def _on_release(self, event): + if event.button != 1 and self._press_button != 1: + return + prev_state = self._state + self._state = self._IDLE + try: + if prev_state == self._DRAW_NEW: + self._finish_draw() + elif prev_state in (self._MOVE, self._DRAG_CORNER): + self._draw_roi() + self.draw_idle() + self.roi_changed.emit() + else: + self._maybe_emit_pixel_click(event) + finally: + self._press_xy = self._press_pixel = None + self._press_button = None + + def _finish_draw(self): + """Commit (or reject) a freshly-dragged quad.""" + if self._extent is not None: + x0, x1, y_bot, y_top = self._extent + min_w = abs(x1 - x0) * 0.01 # minimum: 1% of each axis range + min_h = abs(y_bot - y_top) * 0.01 + else: + min_w = min_h = 1e-6 + + if self._roi is None: + too_small = True + else: + bbox = self._roi.bbox_size() + too_small = bbox[0] < min_w or bbox[1] < min_h + if too_small: + self._roi = self._draw_previous + + self._draw_previous = None + self.cancel_drawing() + self._draw_roi() + self.draw_idle() + self.roi_changed.emit() + + def _maybe_emit_pixel_click(self, event): + """A release close enough to its press counts as a pixel click.""" + if (self._press_pixel is None or event.x is None or event.y is None + or self._extent is None or event.inaxes is not self.ax + or event.xdata is None): + return + dx_px = event.x - self._press_pixel[0] + dy_px = event.y - self._press_pixel[1] + if dx_px * dx_px + dy_px * dy_px > self._CLICK_THRESH_PX ** 2: + return + + x0, x1, y_bot, y_top = self._extent + n_rows, n_frames = self._img_shape + col = int((event.xdata - x0) / (x1 - x0) * n_frames) + row = int((event.ydata - y_top) / (y_bot - y_top) * n_rows) + self.pixel_clicked.emit(max(0, min(row, n_rows - 1)), + max(0, min(col, n_frames - 1))) + + +class WaveformCanvas(FigureCanvasQTAgg): + def __init__(self, parent=None): + fig = Figure(figsize=(8, 3), tight_layout=True) + self.ax_wave = fig.add_subplot(121) + self.ax_right = fig.add_subplot(122) + super().__init__(fig) + self.setParent(parent) + self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding) + + def show_rf_waveform(self, sras: SrasFile, angle_idx: int, + row_idx: int, frame_idx: int, + apply_bg_sub: bool = True): + """CH1 RF: time-domain + FFT spectrum. + + If apply_bg_sub is True and sras.background is not None, the background + waveform is overlaid on the time-domain plot and the FFT is computed + on the subtracted signal. The unsubtracted FFT is also shown faintly + for comparison. + """ + data = sras.data[angle_idx] + waveform = data[row_idx, CH1_IDX, frame_idx, :].astype(np.float32) + t_ns = sras.time_axis_ns() + f_mhz = sras.freq_axis_mhz() + dc3_val = data[row_idx, CH3_IDX, frame_idx, :].astype(np.float32).mean() + dc4_val = data[row_idx, CH4_IDX, frame_idx, :].astype(np.float32).mean() + + bg = sras.background if (apply_bg_sub and sras.background is not None) else None + waveform_plot = waveform - bg if bg is not None else waveform + + self.ax_wave.cla() + self.ax_right.cla() + + if bg is not None: + self.ax_wave.plot(t_ns, waveform, linewidth=0.5, color="#aaaaaa", + label="raw", zorder=1) + self.ax_wave.plot(t_ns, bg, linewidth=0.5, color="#e07030", + linestyle="--", label="background", zorder=2) + self.ax_wave.plot(t_ns, waveform_plot, linewidth=0.7, color="#4488cc", + label="subtracted", zorder=3) + self.ax_wave.legend(fontsize=7, loc="upper right") + else: + self.ax_wave.plot(t_ns, waveform, linewidth=0.7, color="#4488cc") + + self.ax_wave.set_xlabel("Time (ns)") + self.ax_wave.set_ylabel("ADC counts") + bg_tag = " [bg sub]" if bg is not None else "" + dc3_mv = adc_to_mv(dc3_val, *sras.cal(CH3_IDX)) + dc4_mv = adc_to_mv(dc4_val, *sras.cal(CH4_IDX)) + self.ax_wave.set_title( + f"CH1 RF row={row_idx} frame={frame_idx}{bg_tag}\n" + f"CH3={dc3_val:.1f} CH4={dc4_val:.1f} " + f"({dc3_mv:.2f} / {dc4_mv:.2f} mV)", + fontsize=8, + ) + + # FFT of the (possibly subtracted) waveform + power_sub = np.abs(np.fft.rfft(waveform_plot)) ** 2 + power_sub[0] = 0.0 + peak_mhz = f_mhz[int(np.argmax(power_sub))] + + if bg is not None: + # Also show the unsubtracted FFT for reference + power_raw = np.abs(np.fft.rfft(waveform)) ** 2 + power_raw[0] = 0.0 + self.ax_right.plot(f_mhz, power_raw, linewidth=0.5, color="#aaaaaa", + label="raw FFT", zorder=1) + + self.ax_right.plot(f_mhz, power_sub, linewidth=0.7, color="#4488cc", + label="subtracted FFT" if bg is not None else None, zorder=2) + self.ax_right.axvline(peak_mhz, color="tomato", linestyle="--", + linewidth=1.2, label=f"peak = {peak_mhz:.1f} MHz") + self.ax_right.set_xlabel("Frequency (MHz)") + self.ax_right.set_ylabel("Power (arb.)") + self.ax_right.set_title("FFT Power Spectrum") + self.ax_right.set_xlim(0, 500) + self.ax_right.legend(fontsize=8) + + self.draw() + + def show_dc_waveform(self, sras: SrasFile, angle_idx: int, ch_idx: int, + row_idx: int, frame_idx: int): + """CH3 or CH4 DC: time-domain + mean annotation.""" + waveform = sras.data[angle_idx][row_idx, ch_idx, frame_idx, :].astype(np.float32) + mean_val = float(waveform.mean()) + mean_mv = adc_to_mv(mean_val, *sras.cal(ch_idx)) + + self.ax_wave.cla() + self.ax_right.cla() + + self.ax_wave.plot(sras.time_axis_ns(), waveform, linewidth=0.7, color="#4488cc") + self.ax_wave.axhline(mean_val, color="tomato", linestyle="--", + linewidth=1.2, label=f"mean = {mean_val:.2f} ADC") + self.ax_wave.set_xlabel("Time (ns)") + self.ax_wave.set_ylabel("ADC counts") + self.ax_wave.set_title( + f"{CH_NAMES[ch_idx]} DC row={row_idx} frame={frame_idx}") + self.ax_wave.legend(fontsize=8) + + self.ax_right.text( + 0.5, 0.5, + f"DC mode\n\nmean = {mean_val:.3f} ADC\n = {mean_mv:.3f} mV", + ha="center", va="center", + transform=self.ax_right.transAxes, fontsize=11, + ) + self.ax_right.set_axis_off() + + self.draw() + + +class ManualAlignOverlayCanvas(FigureCanvasQTAgg): + """Renders ManualAlignmentDialog's multi-angle mask overlay and turns + keyboard input into translate/rotate nudge requests for whichever angle + the dialog currently has active. + + A pure input+render widget — it holds no alignment state and never + touches SrasFile itself; ManualAlignmentDialog owns all of that and + decides, from these signals, whether a cheap single-layer refresh or a + full preview-canvas rebuild is needed. + + FigureCanvasQTAgg is a real QWidget, so keyPressEvent works like on any + other widget, but Qt only ever delivers key events to whichever widget + currently has focus — StrongFocus, plus grabbing focus on click and once + right after the dialog is shown, are both required or arrow keys + silently do nothing. + + Rotate keys are letters (Q/E), not punctuation (comma/period or + brackets): Shift+letter still reports the same Qt.Key on every platform, + whereas Shift+comma/bracket can report a different virtual key + (Key_Less / Key_BraceLeft) depending on platform and keyboard layout — + which would silently break the "Shift = coarse step" modifier for + rotation specifically. Arrow keys have no such hazard. + """ + nudge_translate = pyqtSignal(int, int, bool) # dir_x, dir_y in {-1,0,1}; coarse + nudge_rotate = pyqtSignal(int, bool) # dir in {-1,1} (CCW/CW); coarse + + _TRANSLATE_KEYS = { + Qt.Key.Key_Left: (-1, 0), + Qt.Key.Key_Right: (1, 0), + Qt.Key.Key_Up: (0, -1), + Qt.Key.Key_Down: (0, 1), + } + _ROTATE_KEYS = {Qt.Key.Key_Q: 1, Qt.Key.Key_E: -1} # CCW, CW + + def __init__(self, parent=None): + fig = Figure(figsize=(6, 6), tight_layout=True) + self.ax = fig.add_subplot(111) + super().__init__(fig) + self.setParent(parent) + self.setFocusPolicy(Qt.FocusPolicy.StrongFocus) + self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding) + self.mpl_connect("button_press_event", lambda _e: self.setFocus()) + + def show_overlay(self, rgba: np.ndarray, extent: list[float], title: str): + self.figure.clf() + self.ax = self.figure.add_subplot(111) + self.ax.imshow(rgba, extent=extent, origin="upper", aspect="auto") + self.ax.set_xlabel("X (mm)") + self.ax.set_ylabel("Y (mm)") + self.ax.set_title(title) + self.draw_idle() # coalesces rapid redraws — matters for key-repeat. + + def keyPressEvent(self, event: QKeyEvent): + key = event.key() + coarse = bool(event.modifiers() & Qt.KeyboardModifier.ShiftModifier) + if key in self._TRANSLATE_KEYS: + dx, dy = self._TRANSLATE_KEYS[key] + self.nudge_translate.emit(dx, dy, coarse) + event.accept() + elif key in self._ROTATE_KEYS: + self.nudge_rotate.emit(self._ROTATE_KEYS[key], coarse) + event.accept() + else: + super().keyPressEvent(event) + diff --git a/sras_viewer/common.py b/sras_viewer/common.py new file mode 100644 index 0000000..7d07a50 --- /dev/null +++ b/sras_viewer/common.py @@ -0,0 +1,115 @@ +"""Shared constants and small layout helpers for the viewer widgets.""" + +from PyQt6.QtCore import Qt +from PyQt6.QtWidgets import ( + QFormLayout, QFrame, QGroupBox, QLabel, QScrollArea, QSizePolicy, + QVBoxLayout, QWidget, +) + +from sras_format import CH1_IDX, CH3_IDX, CH4_IDX + +# --------------------------------------------------------------------------- +# Display constants +# --------------------------------------------------------------------------- + +CH_LABELS = [ + "CH1 — RF (FFT peak freq)", + "CH3 — Bias A (DC mean)", + "CH4 — Bias B (DC mean)", + "CH1 — Velocity (SRAS)", +] + +# Combo index for the derived velocity mode (uses CH1_IDX data) +VELOCITY_MODE_IDX = 3 +# All modes that operate on CH1 waveforms +CH1_DERIVED_MODES = (CH1_IDX, VELOCITY_MODE_IDX) + +CMAPS = ["gray", "viridis", "plasma", "inferno", "hot", "jet", "RdBu_r", "seismic"] + +# (mode_str, status-bar unit, colorbar label) per channel index +_CHANNEL_DISPLAY = { + CH1_IDX: ("RF", "Peak frequency (MHz)", "MHz"), + CH3_IDX: ("DC", "DC mean (mV)", "mV"), + CH4_IDX: ("DC", "DC mean (mV)", "mV"), + VELOCITY_MODE_IDX: ("Velocity", "Velocity (m/s)", "m/s"), +} + +_CSS_HINT = "font-size: 11px; color: #aaa;" +_CSS_INFO = "font-size: 11px;" +_CSS_MUTED = "color: #888; font-size: 11px;" +_CSS_WARN = "color: #e07000; font-size: 11px;" +_CSS_BUSY = "color: #4a90d9; font-size: 11px;" + +# Side-panel column widths (the scroll areas that hold the controls). +_LEFT_PANEL_W = 288 +_RIGHT_PANEL_W = 272 + +# Minimum width for a spin box so its value + suffix are never clipped. +_SPIN_MIN_W = 96 + + +# --------------------------------------------------------------------------- +# Small layout helpers +# --------------------------------------------------------------------------- + +def _wrap_label(text: str = "", css: str | None = None) -> QLabel: + """A word-wrapped QLabel that reports its *wrapped* height to the layout. + + A plain word-wrapped QLabel advertises a single-line minimum height, so in a + fixed-width column the layout happily shrinks it and the extra lines get + clipped. Enabling height-for-width makes the box layout ask for the real + height at the column's width instead. + """ + lbl = QLabel(text) + lbl.setWordWrap(True) + sp = lbl.sizePolicy() + sp.setVerticalPolicy(QSizePolicy.Policy.Minimum) + sp.setHeightForWidth(True) + lbl.setSizePolicy(sp) + if css: + lbl.setStyleSheet(css) + return lbl + + +def _group(title: str) -> tuple[QGroupBox, QVBoxLayout]: + """A group box with consistent, non-cramped internal margins.""" + grp = QGroupBox(title) + lay = QVBoxLayout(grp) + lay.setContentsMargins(10, 8, 10, 10) + lay.setSpacing(6) + return grp, lay + + +def _form() -> QFormLayout: + """A label/field form layout for a narrow side panel.""" + form = QFormLayout() + form.setContentsMargins(0, 0, 0, 0) + form.setHorizontalSpacing(8) + form.setVerticalSpacing(6) + form.setLabelAlignment(Qt.AlignmentFlag.AlignRight + | Qt.AlignmentFlag.AlignVCenter) + form.setFormAlignment(Qt.AlignmentFlag.AlignLeft | Qt.AlignmentFlag.AlignTop) + form.setFieldGrowthPolicy( + QFormLayout.FieldGrowthPolicy.AllNonFixedFieldsGrow) + form.setRowWrapPolicy(QFormLayout.RowWrapPolicy.DontWrapRows) + return form + + +def _scroll_panel(inner: QWidget, width: int) -> QScrollArea: + """Put a side panel in a fixed-width scroll area. + + Without this the panels are sized by the window: a short window squeezes the + controls past their minimum heights, which is what makes text overlap the + widget below it. Scrolling keeps every control at its natural size. + """ + area = QScrollArea() + area.setWidget(inner) + area.setWidgetResizable(True) + area.setFrameShape(QFrame.Shape.NoFrame) + area.setHorizontalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAlwaysOff) + area.setVerticalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAsNeeded) + area.setFixedWidth(width) + area.viewport().setAutoFillBackground(False) + inner.setAutoFillBackground(False) + return area + diff --git a/sras_viewer/dialogs.py b/sras_viewer/dialogs.py new file mode 100644 index 0000000..d1efed7 --- /dev/null +++ b/sras_viewer/dialogs.py @@ -0,0 +1,767 @@ +"""FFT Options and Manual Alignment dialogs.""" + +from typing import TYPE_CHECKING + +import matplotlib as mpl +import numpy as np +from matplotlib.backends.backend_qtagg import NavigationToolbar2QT +from PyQt6.QtCore import pyqtSignal +from PyQt6.QtWidgets import ( + QButtonGroup, QComboBox, QDialog, QDialogButtonBox, QDoubleSpinBox, + QGroupBox, QHBoxLayout, QLabel, QMessageBox, QPushButton, QRadioButton, + QSpinBox, QVBoxLayout, QWidget, +) + +import sras_compute as compute +from sras_compute import ( + PYFFTW_AVAILABLE, ManualAngleParams, build_manual_alignment, + delete_manual_alignment, save_manual_alignment, +) +from sras_format import SrasFile +from sras_workers import Ch4MaskWorker, CrossCorrelateWorker + +from .canvases import ManualAlignOverlayCanvas +from .common import ( + _CSS_HINT, _CSS_MUTED, _CSS_WARN, _SPIN_MIN_W, _form, _group, + _scroll_panel, _wrap_label, +) + +if TYPE_CHECKING: + from .main_window import SrasViewerWindow + + +# --------------------------------------------------------------------------- +# FFT Options dialog +# --------------------------------------------------------------------------- + +class FftOptionsDialog(QDialog): + """Configure FFT backend and zero-padding. + + Changes take effect only when the user clicks Apply. Cancel discards + all pending edits. The live 'frequency resolution' label updates as + the user adjusts the pad factor so they can see the trade-off before + committing. + """ + + def __init__(self, parent=None, *, + current_backend: str, + current_pad_factor: int, + samples_per_frame: int | None, + sample_rate_hz: float | None, + grating_um: float): + super().__init__(parent) + self.setWindowTitle("FFT Options") + self.setModal(True) + self.setMinimumWidth(380) + + self._samples_per_frame = samples_per_frame + self._sample_rate_hz = sample_rate_hz + self._grating_um = grating_um + + layout = QVBoxLayout(self) + + # ---- Backend --------------------------------------------------- + grp_backend = QGroupBox("FFT Backend") + bl = QVBoxLayout(grp_backend) + + self._btn_scipy = QRadioButton("SciPy FFT (pocketfft) (always available)") + self._btn_pyfftw = QRadioButton( + "pyFFTW (faster for large arrays)" if PYFFTW_AVAILABLE + else "pyFFTW (not installed — run: pip install pyfftw)") + self._btn_pyfftw.setEnabled(PYFFTW_AVAILABLE) + + self._backend_group = QButtonGroup(self) + self._backend_group.addButton(self._btn_scipy, id=0) + self._backend_group.addButton(self._btn_pyfftw, id=1) + + if current_backend == "pyfftw" and PYFFTW_AVAILABLE: + self._btn_pyfftw.setChecked(True) + else: + self._btn_scipy.setChecked(True) + + bl.addWidget(self._btn_scipy) + bl.addWidget(self._btn_pyfftw) + layout.addWidget(grp_backend) + + # ---- Zero-padding ---------------------------------------------- + grp_zp = QGroupBox("Zero-Padding") + zl = QVBoxLayout(grp_zp) + + pad_row = QHBoxLayout() + pad_row.addWidget(QLabel("Pad factor:")) + self._spin_pad = QSpinBox() + self._spin_pad.setRange(1, 256) + self._spin_pad.setValue(max(1, current_pad_factor)) + self._spin_pad.setToolTip( + "Multiply the waveform length by this factor via zero-padding\n" + "before computing the FFT.\n" + "1 = no padding (natural length).\n" + "Powers of 2 (2, 4, 8 …) give the best performance." + ) + self._spin_pad.valueChanged.connect(self._update_info) + pad_row.addWidget(self._spin_pad) + zl.addLayout(pad_row) + + self._lbl_nfft = QLabel() + self._lbl_freq_res = QLabel() + self._lbl_vel_res = QLabel() + for lbl in (self._lbl_nfft, self._lbl_freq_res, self._lbl_vel_res): + lbl.setStyleSheet(_CSS_HINT) + zl.addWidget(lbl) + + layout.addWidget(grp_zp) + + # ---- Buttons --------------------------------------------------- + buttons = QDialogButtonBox() + buttons.addButton("Apply", QDialogButtonBox.ButtonRole.AcceptRole + ).clicked.connect(self.accept) + buttons.addButton("Cancel", QDialogButtonBox.ButtonRole.RejectRole + ).clicked.connect(self.reject) + layout.addWidget(buttons) + + self._update_info() + + def _update_info(self): + spf = self._samples_per_frame + sr = self._sample_rate_hz + pad = self._spin_pad.value() + + if spf is None or sr is None: + self._lbl_nfft.setText("Load a file to preview FFT parameters.") + self._lbl_freq_res.setText("") + self._lbl_vel_res.setText("") + return + + n_fft = spf * pad + freq_res_hz = sr / n_fft + freq_res_mhz = freq_res_hz / 1e6 + # v (m/s) = freq (MHz) × grating (µm) + vel_res_ms = freq_res_mhz * self._grating_um + + self._lbl_nfft.setText(f"FFT points: {spf} × {pad} = {n_fft:,}") + self._lbl_freq_res.setText( + f"Frequency bin: {freq_res_mhz:.4f} MHz ({freq_res_hz / 1e3:.2f} kHz)") + self._lbl_vel_res.setText( + f"Velocity bin: {vel_res_ms:.3f} m/s " + f"(at grating = {self._grating_um:.2f} µm)") + + def get_backend(self) -> str: + return "pyfftw" if self._btn_pyfftw.isChecked() and PYFFTW_AVAILABLE else "scipy" + + def get_pad_factor(self) -> int: + return max(1, self._spin_pad.value()) + + + +class ManualAlignmentDialog(QDialog): + """Non-modal manual angle-alignment editor (Fusion -> Manual Alignment...). + + 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. 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 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 + SrasViewerWindow two ways: it reuses parent._run_worker/_jobs directly + 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) + alignment_cleared = pyqtSignal() + + _PREVIEW_MARGIN_FRAC = 0.15 + _BASE_ALPHA = 0.42 + _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, + cached_dc4_mv: dict[int, np.ndarray]): + super().__init__(parent) + self._parent = parent + self._sras = sras + self._ref_angle_idx = ref_angle_idx + self._downsample = (1, 1) # (rows, cols) block-mean factors + self._dc4_mv: dict[int, np.ndarray] = {} + self._masks_small: dict[int, np.ndarray] = {} + self._preview_layers: dict[int, np.ndarray] = {} + self._preview_origin_mm = (0.0, 0.0) + self._preview_shape = (1, 1) + 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) + + self._seed_initial_params(seed_per_angle) + 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._active_angle = 1 if ref_angle_idx == 0 and n > 1 else 0 + + self._build_ui(dc_threshold_mv) + self._set_controls_enabled(False) # re-enabled once masks are ready + self._start_mask_prep(cached_dc4_mv) + + def showEvent(self, event): + super().showEvent(event) + self.canvas.setFocus() + + # ------------------------------------------------------------------ + # Construction + # ------------------------------------------------------------------ + + def _seed_initial_params(self, seed_per_angle: dict[int, ManualAngleParams] | None): + seed = seed_per_angle or {} + self._angle_params: dict[int, ManualAngleParams] = { + a: (ManualAngleParams(seed[a].rotation_deg, seed[a].shift_mm) + if a in seed else ManualAngleParams()) + for a in range(self._sras.n_angles) + } + self._angle_params[self._ref_angle_idx] = ManualAngleParams() + + def _build_ui(self, dc_threshold_mv: float): + root = QHBoxLayout(self) + + self.canvas = ManualAlignOverlayCanvas() + left = QWidget() + left_l = QVBoxLayout(left) + left_l.setContentsMargins(0, 0, 0, 0) + left_l.setSpacing(4) + left_l.addWidget(NavigationToolbar2QT(self.canvas, left)) + left_l.addWidget(self.canvas) + root.addWidget(left, stretch=1) + + panel = QWidget() + panel_l = QVBoxLayout(panel) + panel_l.setContentsMargins(0, 0, 0, 0) + panel_l.setSpacing(8) + + # ---- Active Angle ------------------------------------------------- + grp_angle, al = _group("Active Angle") + self.combo_active_angle = QComboBox() + for a in range(self._sras.n_angles): + label = f"Angle {a} ({self._sras.angles_deg[a]:.1f}°)" + if a == self._ref_angle_idx: + label += " [reference]" + self.combo_active_angle.addItem(label) + al.addWidget(self.combo_active_angle) + self.lbl_active_note = _wrap_label("", _CSS_WARN) + al.addWidget(self.lbl_active_note) + panel_l.addWidget(grp_angle) + + # ---- Manual Adjustment --------------------------------------------- + self.grp_manual_adjust, mform_box = _group("Manual Adjustment") + mform = _form() + self.spin_active_rotation_deg = QDoubleSpinBox() + self.spin_active_rotation_deg.setRange(-3600.0, 3600.0) + self.spin_active_rotation_deg.setDecimals(3) + self.spin_active_rotation_deg.setSuffix(" °") + self.spin_active_rotation_deg.setMinimumWidth(_SPIN_MIN_W) + mform.addRow("Rotation:", self.spin_active_rotation_deg) + + self.spin_active_shift_x_mm = QDoubleSpinBox() + self.spin_active_shift_x_mm.setRange(-1e5, 1e5) + self.spin_active_shift_x_mm.setDecimals(4) + self.spin_active_shift_x_mm.setSuffix(" mm") + self.spin_active_shift_x_mm.setMinimumWidth(_SPIN_MIN_W) + mform.addRow("Shift X:", self.spin_active_shift_x_mm) + + self.spin_active_shift_y_mm = QDoubleSpinBox() + self.spin_active_shift_y_mm.setRange(-1e5, 1e5) + self.spin_active_shift_y_mm.setDecimals(4) + self.spin_active_shift_y_mm.setSuffix(" mm") + self.spin_active_shift_y_mm.setMinimumWidth(_SPIN_MIN_W) + mform.addRow("Shift Y:", self.spin_active_shift_y_mm) + mform_box.addLayout(mform) + panel_l.addWidget(self.grp_manual_adjust) + + # ---- Nudge Step Sizes ------------------------------------------------ + self.grp_step_sizes, sl = _group("Nudge Step Sizes") + sform = _form() + self.spin_step_translate_mm = QDoubleSpinBox() + self.spin_step_translate_mm.setRange(0.0001, 1000.0) + self.spin_step_translate_mm.setDecimals(4) + self.spin_step_translate_mm.setSuffix(" mm") + self.spin_step_translate_mm.setValue(0.01) + self.spin_step_translate_mm.setMinimumWidth(_SPIN_MIN_W) + sform.addRow("Translate step:", self.spin_step_translate_mm) + + self.spin_step_rotate_deg = QDoubleSpinBox() + self.spin_step_rotate_deg.setRange(0.001, 90.0) + self.spin_step_rotate_deg.setDecimals(3) + self.spin_step_rotate_deg.setSuffix(" °") + self.spin_step_rotate_deg.setValue(0.1) + self.spin_step_rotate_deg.setMinimumWidth(_SPIN_MIN_W) + sform.addRow("Rotate step:", self.spin_step_rotate_deg) + + self.spin_step_multiplier = QDoubleSpinBox() + self.spin_step_multiplier.setRange(1.0, 1000.0) + self.spin_step_multiplier.setDecimals(1) + self.spin_step_multiplier.setValue(10.0) + self.spin_step_multiplier.setMinimumWidth(_SPIN_MIN_W) + sform.addRow("Coarse × (Shift):", self.spin_step_multiplier) + sl.addLayout(sform) + sl.addWidget(_wrap_label( + "Arrow keys nudge X/Y translation; Q/E nudge rotation (CCW/CW). " + "Hold Shift for the coarse step. Click the image once so it has " + "keyboard focus.", _CSS_HINT)) + panel_l.addWidget(self.grp_step_sizes) + + # ---- Mask Threshold --------------------------------------------------- + self.grp_mask_threshold, tl = _group("Mask Threshold") + tform = _form() + self.spin_mask_threshold_mv = QDoubleSpinBox() + self.spin_mask_threshold_mv.setRange(-500.0, 500.0) + self.spin_mask_threshold_mv.setDecimals(3) + self.spin_mask_threshold_mv.setSuffix(" mV") + self.spin_mask_threshold_mv.setValue(dc_threshold_mv) + self.spin_mask_threshold_mv.setMinimumWidth(_SPIN_MIN_W) + tform.addRow("DC threshold:", self.spin_mask_threshold_mv) + 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() + 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_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( + "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 ------------------------------------------------------ + grp_actions, acl = _group("Actions") + self.btn_auto_derotate = QPushButton("Auto De-rotate (use known angles)") + self.btn_save = QPushButton("Save Alignment") + self.btn_clear = QPushButton("Clear Alignment…") + self.btn_close = QPushButton("Close") + for btn in (self.btn_auto_derotate, self.btn_save, self.btn_clear, self.btn_close): + acl.addWidget(btn) + panel_l.addWidget(grp_actions) + + self.lbl_status = _wrap_label("", _CSS_MUTED) + panel_l.addWidget(self.lbl_status) + panel_l.addStretch() + + 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) + self.spin_active_shift_x_mm.editingFinished.connect(self._on_shift_spin_edited) + 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) + self.canvas.nudge_translate.connect(self._on_nudge_translate) + self.canvas.nudge_rotate.connect(self._on_nudge_rotate) + + self.combo_active_angle.blockSignals(True) + self.combo_active_angle.setCurrentIndex(self._active_angle) + self.combo_active_angle.blockSignals(False) + self._on_active_angle_changed(self._active_angle) + + # ------------------------------------------------------------------ + # Mask preparation (initial CH4 fetch + threshold + downsample) + # ------------------------------------------------------------------ + + def _start_mask_prep(self, cached_dc4_mv: dict[int, np.ndarray]): + self._dc4_mv = dict(cached_dc4_mv) + missing = [a for a in range(self._sras.n_angles) if a not in self._dc4_mv] + if not missing: + self._finish_mask_prep() + return + self.lbl_status.setText(f"Preparing masks: 0/{len(missing)} angle(s) needed…") + started = self._parent._run_worker( + "manual_align_masks", Ch4MaskWorker(self._sras, missing), + connect=( + ("angle_done", self._on_mask_angle_done), + ("error", lambda msg: self.lbl_status.setText(f"Mask prep error: {msg}")), + ), + on_done=self._finish_mask_prep) + if not started: + self.lbl_status.setText( + "Could not start mask preparation (busy) — close and reopen.") + + def _on_mask_angle_done(self, angle_idx: int, dc4_mv: np.ndarray): + self._dc4_mv[angle_idx] = dc4_mv + self.lbl_status.setText( + f"Preparing masks: {len(self._dc4_mv)}/{self._sras.n_angles} ready…") + + def _finish_mask_prep(self): + if len(self._dc4_mv) < self._sras.n_angles: + return # a mask-worker error left some angles unfetched + # 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() + self._rebuild_preview_canvas() + self._set_controls_enabled(True) + self.lbl_status.setText("Ready.") + + def _recompute_masks_small(self): + """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: no alignment geometry depends on this + threshold, only which pixels the overlay paints.""" + threshold = self.spin_mask_threshold_mv.value() + fy, fx = self._downsample + self._masks_small = { + a: compute.block_mean_2d((img >= threshold).astype(np.float32), fy, fx) + for a, img in self._dc4_mv.items() + } + + # ------------------------------------------------------------------ + # Preview canvas: full rebuild vs. incremental single-layer refresh + # ------------------------------------------------------------------ + + def _rebuild_preview_canvas(self): + """Full geometry rebuild: recomputes the shared preview canvas's + origin/shape (rotation can grow the union bbox — translation alone + cannot, per the padding baked in via _PREVIEW_MARGIN_FRAC) and every + angle's reprojected mask layer. Triggered by: dialog open, + mask-threshold change, Auto De-rotate, a rotation nudge/edit of the + 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) + 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_pitch_mm = pitch + self._preview_layers = { + 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.""" + self._preview_layers[self._active_angle] = self._reproject(self._active_angle) + self._redraw_overlay() + + def _redraw_overlay(self): + """Alpha-composite every angle's colored mask layer into one RGBA + image ("all thresholds overlaid with varying opacity"). Each angle + keeps a fixed, distinct color regardless of which is active; the + active angle is drawn last (on top) at a visibly higher alpha so + it's easy to track while nudging.""" + if not self._preview_layers: + return # mask prep hasn't finished yet — nothing to draw + n_rows, n_cols = self._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 = self._preview_layers.get(a) + if layer is None: + continue + alpha = self._ACTIVE_ALPHA if a == self._active_angle else self._BASE_ALPHA + color = self._angle_colors[a] + fg_a = layer * alpha + for c in range(3): + rgba[..., c] = color[c] * fg_a + rgba[..., c] * rgba[..., 3] * (1 - fg_a) + rgba[..., 3] = fg_a + rgba[..., 3] * (1 - fg_a) + + x0, y0 = self._preview_origin_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, + y_axis[-1] + dy / 2, y_axis[0] - dy / 2] + title = (f"Angle {self._active_angle} active " + f"({self._sras.angles_deg[self._active_angle]:.1f}°)") + self.canvas.show_overlay(rgba, extent, title) + + # ------------------------------------------------------------------ + # Angle selection / nudge / edit handlers + # ------------------------------------------------------------------ + + def _on_active_angle_changed(self, angle_idx: int): + self._active_angle = angle_idx + is_ref = angle_idx == self._ref_angle_idx + self.grp_manual_adjust.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_active_spinboxes() + self._redraw_overlay() + + def _sync_active_spinboxes(self): + p = self._angle_params[self._active_angle] + for spin, val in ((self.spin_active_rotation_deg, p.rotation_deg), + (self.spin_active_shift_x_mm, p.shift_mm[0]), + (self.spin_active_shift_y_mm, p.shift_mm[1])): + spin.blockSignals(True) + spin.setValue(val) + spin.blockSignals(False) + + def _on_nudge_translate(self, dir_x: int, dir_y: int, coarse: bool): + if not self._masks_ready or self._active_angle == self._ref_angle_idx: + return + step = self.spin_step_translate_mm.value() + if coarse: + step *= self.spin_step_multiplier.value() + p = self._angle_params[self._active_angle] + p.shift_mm = (p.shift_mm[0] + dir_x * step, p.shift_mm[1] + dir_y * step) + self._sync_active_spinboxes() + self._refresh_active_preview_layer() + + def _on_nudge_rotate(self, direction: int, coarse: bool): + if not self._masks_ready or self._active_angle == self._ref_angle_idx: + return + step = self.spin_step_rotate_deg.value() + if coarse: + step *= self.spin_step_multiplier.value() + self._angle_params[self._active_angle].rotation_deg += direction * step + self._sync_active_spinboxes() + self._rebuild_preview_canvas() + + def _on_rotation_spin_edited(self): + if self._active_angle == self._ref_angle_idx: + return + self._angle_params[self._active_angle].rotation_deg = self.spin_active_rotation_deg.value() + self._rebuild_preview_canvas() + + def _on_shift_spin_edited(self): + if self._active_angle == self._ref_angle_idx: + return + p = self._angle_params[self._active_angle] + p.shift_mm = (self.spin_active_shift_x_mm.value(), self.spin_active_shift_y_mm.value()) + self._refresh_active_preview_layer() + + def _on_mask_threshold_edited(self): + if not self._masks_ready: + return + self._recompute_masks_small() + self._rebuild_preview_canvas() + + # ------------------------------------------------------------------ + # Actions + # ------------------------------------------------------------------ + + 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 = 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 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: + return + angles = [a for a in range(self._sras.n_angles) if a != self._ref_angle_idx] + if not angles: + return + worker = CrossCorrelateWorker( + 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( + "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, + 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)…") + + 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) + self.lbl_status.setText( + f"Cross-correlated {self._correlate_done_count} angle(s) against " + 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) + except OSError as exc: + QMessageBox.warning(self, "Save Alignment Failed", str(exc)) + return + self.lbl_status.setText(f"Saved to {path.name}.") + self.alignment_saved.emit(result, str(path)) + + def _on_clear(self): + reply = QMessageBox.question( + self, "Clear Alignment", + "This resets every angle back to raw/unaligned (0° rotation, no " + "shift) and deletes the saved alignment file for this scan, if " + "any. This cannot be undone. Continue?", + QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No, + QMessageBox.StandardButton.No) + if reply != QMessageBox.StandardButton.Yes: + return + try: + existed = delete_manual_alignment(self._sras) + except OSError as exc: + QMessageBox.warning(self, "Clear Alignment Failed", + 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( + "Alignment cleared; saved file removed." if existed + else "Alignment cleared (there was no saved file).") + self.alignment_cleared.emit() + + def _set_controls_enabled(self, enabled: bool): + self._masks_ready = enabled + self.combo_active_angle.setEnabled(enabled) + 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_viewer/main_window.py b/sras_viewer/main_window.py new file mode 100644 index 0000000..b0fa6a5 --- /dev/null +++ b/sras_viewer/main_window.py @@ -0,0 +1,1403 @@ +"""The SrasViewerWindow main window and application entry point.""" + +import sys +from pathlib import Path + +import numpy as np +from matplotlib.backends.backend_qtagg import NavigationToolbar2QT +from PyQt6.QtCore import QObject, QSettings, Qt, QThread +from PyQt6.QtGui import QAction +from PyQt6.QtWidgets import ( + QApplication, QCheckBox, QComboBox, QDialog, QDoubleSpinBox, QFileDialog, + QFrame, QHBoxLayout, QLabel, QMainWindow, QProgressDialog, QPushButton, + QSizePolicy, QSpinBox, QSplitter, QVBoxLayout, QWidget, +) + +import sras_compute as compute +from sras_compute import ( + ManualAngleParams, apply_alignment, build_manual_alignment, + load_manual_alignment, sidecar_path, +) +from sras_format import ( + CH3_IDX, CH4_IDX, CH_NAMES, SrasFile, mv_to_adc, + _FALLBACK_YMULT_MV, _FALLBACK_YOFF_ADC, +) +from sras_workers import ( + AngleAlignmentWorker, BatchCacheWorker, ComputeWorker, DcPrecomputeWorker, + LoadWorker, +) + +from .canvases import ImageCanvas, WaveformCanvas +from .common import ( + CH1_DERIVED_MODES, CH_LABELS, CMAPS, VELOCITY_MODE_IDX, _CHANNEL_DISPLAY, + _CSS_BUSY, _CSS_HINT, _CSS_INFO, _CSS_MUTED, _CSS_WARN, _LEFT_PANEL_W, + _RIGHT_PANEL_W, _SPIN_MIN_W, _form, _group, _scroll_panel, _wrap_label, +) +from .dialogs import FftOptionsDialog, ManualAlignmentDialog + +# --------------------------------------------------------------------------- +# Main window +# --------------------------------------------------------------------------- + +class SrasViewerWindow(QMainWindow): + def __init__(self, initial_path: str | None = None): + super().__init__() + self.setWindowTitle("SRAS Scan Viewer") + self.resize(1560, 840) + self.setMinimumSize(960, 560) + self.setAcceptDrops(True) + + self._sras: SrasFile | None = None + self._current_image: np.ndarray | None = None + self._current_angle: int = 0 + self._current_ch: int = 0 + self._pending_angle: int = 0 + self._pending_ch: int = 0 + self._pending_bg_sub: bool = True + self._pending_threshold: float = 50.0 # mV + self._pending_fft_pad_factor: int = 1 + + # Live background jobs, keyed by role — see _run_worker. + self._jobs: dict[str, tuple] = {} + self._progress_dlgs: dict[str, QProgressDialog] = {} + + # FFT settings (configured via FFT Options dialog, persisted across + # sessions). IniFormat: predictable cross-platform and redirectable + # in tests. + self._settings = QSettings(QSettings.Format.IniFormat, + QSettings.Scope.UserScope, + "sras-viewer", "sras-viewer") + compute.set_fft_backend(str(self._settings.value("fft/backend", "scipy"))) + try: + pad = int(self._settings.value("fft/pad_factor", 1)) + except (TypeError, ValueError): + pad = 1 + self._fft_pad_factor: int = max(1, min(256, pad)) # 1 = no padding + + # Convert menu: batch DC/FFT compute-and-store (v6 -> v7) + self._batch_errors: list[str] = [] + + # Display-only settings (colormap, grating) never trigger a + # recompute — they're applied to cached data on redraw. DC images + # (CH3/CH4) are cheap and precomputed for every angle in the + # background right after load. CH1/Velocity FFT images are + # computed lazily (with a progress popup) the first time an + # angle/threshold combination is viewed — using the cached DC4 + # image to skip the FFT entirely for masked-out pixels — and + # cached per (angle, bg_sub, n_fft, threshold) so revisiting the + # same combination is free. + self._dc_cache: dict[tuple[int, int], np.ndarray] = {} + self._fft_cache: dict[tuple[int, bool, int | None, float], np.ndarray] = {} + self._dc_generation: int = 0 + + # Angle alignment ("Fusion" menu) + self._alignment_result = None + self._alignment_generation: int = 0 + self._aligned_cache: dict[tuple, np.ndarray] = {} + self._manual_align_dialog: ManualAlignmentDialog | None = None + + self._build_ui() + + if initial_path: + self._load_file(initial_path) + + # ------------------------------------------------------------------ + # Background job plumbing + # ------------------------------------------------------------------ + + def _run_worker(self, key: str, worker: QObject, *, + connect: tuple = (), quit_on: tuple = ("finished",), + on_done=None) -> bool: + """Move *worker* onto its own QThread and start it. Returns False if + a job under *key* is already running. + + Centralises two lifetime hazards that each cost a process abort: + + 1. The job is claimed in self._jobs *before* start() and before + anything below that can pump the Qt event loop (a + QProgressDialog.show() does on first display). If it weren't, a + re-entrant editingFinished could slip past the busy check, start a + second thread, and then have the first call's own assignment + clobber — and destroy while still running — that second QThread. + + 2. thread.finished fires as the thread winds down but does not + guarantee the OS thread has joined. Dropping the last reference to + a QThread whose thread is still running logs "QThread: Destroyed + while thread is still running" and aborts, so wait() first. + """ + if key in self._jobs: + return False + + thread = QThread() + self._jobs[key] = (thread, worker, on_done) # claim before anything pumps + worker.moveToThread(thread) + thread.started.connect(worker.run) + for signal_name, slot in connect: + getattr(worker, signal_name).connect(slot) + for signal_name in quit_on: + getattr(worker, signal_name).connect(thread.quit) + thread.finished.connect(lambda k=key: self._on_job_finished(k)) + thread.start() + return True + + def _on_job_finished(self, key: str): + job = self._jobs.pop(key, None) + if job is None: + return + thread, _worker, on_done = job + thread.wait() # join before releasing our last reference + if on_done is not None: + on_done() + + def _job_running(self, key: str) -> bool: + return key in self._jobs + + # ------------------------------------------------------------------ + # UI construction + # ------------------------------------------------------------------ + + def _build_ui(self): + central = QWidget() + self.setCentralWidget(central) + root = QHBoxLayout(central) + root.setContentsMargins(8, 8, 8, 8) + root.setSpacing(8) + + root.addWidget(self._build_left_panel()) + root.addWidget(self._build_canvases(), stretch=1) + root.addWidget(self._build_right_panel()) + + self.statusBar().showMessage("Open an .sras file to begin.") + self._build_menus() + + def _build_left_panel(self) -> QWidget: + panel = QWidget() + panel_layout = QVBoxLayout(panel) + panel_layout.setContentsMargins(0, 0, 0, 0) + panel_layout.setSpacing(8) + + # ---- File ------------------------------------------------------- + grp_file, fl = _group("File") + self.btn_open = QPushButton("Open .sras…") + self.btn_open.clicked.connect(self._on_open) + self.lbl_filename = _wrap_label("No file loaded", _CSS_MUTED) + fl.addWidget(self.btn_open) + fl.addWidget(self.lbl_filename) + panel_layout.addWidget(grp_file) + + # ---- Scan info -------------------------------------------------- + grp_info, il = _group("Scan Info") + il.setSpacing(3) + self._info = {} + for key in ("Angles", "Rows", "Frames / row", "Samples / frame", + "Sample rate", "X start", "Pixel Δx", "Laser freq"): + lbl = _wrap_label(f"{key}: —", _CSS_INFO) + il.addWidget(lbl) + self._info[key] = lbl + + # frame-count / format notes + self.lbl_frame_warn = _wrap_label("", _CSS_WARN) + il.addWidget(self.lbl_frame_warn) + + # background DC-precompute progress + self.lbl_dc_precompute = _wrap_label("", _CSS_BUSY) + il.addWidget(self.lbl_dc_precompute) + panel_layout.addWidget(grp_info) + + # ---- View settings ---------------------------------------------- + grp_view, vl = _group("View Settings") + + view_form = _form() + + self.spin_angle = QSpinBox() + self.spin_angle.setRange(0, 0) + self.spin_angle.setEnabled(False) + self.spin_angle.setMinimumWidth(64) + self.spin_angle.editingFinished.connect(self._on_view_changed) + self.lbl_angle_deg = QLabel("—") + angle_field = QWidget() + ar = QHBoxLayout(angle_field) + ar.setContentsMargins(0, 0, 0, 0) + ar.setSpacing(6) + ar.addWidget(self.spin_angle) + ar.addWidget(self.lbl_angle_deg) + ar.addStretch() + view_form.addRow("Angle:", angle_field) + + self.combo_channel = QComboBox() + self.combo_channel.addItems(CH_LABELS) + self.combo_channel.setEnabled(False) + self.combo_channel.setSizePolicy(QSizePolicy.Policy.Expanding, + QSizePolicy.Policy.Fixed) + self.combo_channel.setSizeAdjustPolicy( + QComboBox.SizeAdjustPolicy.AdjustToMinimumContentsLengthWithIcon) + self.combo_channel.setMinimumContentsLength(12) + self.combo_channel.currentIndexChanged.connect(self._on_channel_changed) + view_form.addRow("Channel:", self.combo_channel) + vl.addLayout(view_form) + + sep = QFrame() + sep.setFrameShape(QFrame.Shape.HLine) + sep.setStyleSheet("color: #555;") + vl.addWidget(sep) + + # DC threshold (for RF / CH1 masking) + self.grp_threshold, tl = _group("RF Mask Threshold (CH1 only)") + thr_form = _form() + self.spin_threshold_mv = QDoubleSpinBox() + self.spin_threshold_mv.setRange(-500.0, 500.0) + self.spin_threshold_mv.setDecimals(3) + self.spin_threshold_mv.setSingleStep(0.025) + self.spin_threshold_mv.setSuffix(" mV") + self.spin_threshold_mv.setValue(50.0) + self.spin_threshold_mv.setEnabled(False) + self.spin_threshold_mv.setMinimumWidth(_SPIN_MIN_W) + self.spin_threshold_mv.editingFinished.connect(self._on_threshold_changed) + thr_form.addRow("DC threshold:", self.spin_threshold_mv) + tl.addLayout(thr_form) + self.lbl_threshold_adc = _wrap_label( + f"≈ {mv_to_adc(50.0):.1f} ADC counts", _CSS_MUTED) + tl.addWidget(self.lbl_threshold_adc) + vl.addWidget(self.grp_threshold) + + # Background subtraction (v4+ files only) + self.chk_bg_sub = QCheckBox("Background subtraction (CH1 only)") + self.chk_bg_sub.setChecked(True) + self.chk_bg_sub.setEnabled(False) + self.chk_bg_sub.setToolTip( + "Subtract the stored background waveform from each CH1 frame\n" + "before computing the FFT (v4+ files only)." + ) + self.chk_bg_sub.toggled.connect(self._on_bg_sub_toggled) + vl.addWidget(self.chk_bg_sub) + + # Aligned View (Fusion → Angle Alignment result) + self.chk_aligned_view = QCheckBox("Aligned View (Fusion)") + self.chk_aligned_view.setChecked(False) + self.chk_aligned_view.setEnabled(False) + self.chk_aligned_view.setToolTip( + "Show the current angle/channel resampled onto the shared,\n" + "rotation+translation-aligned canvas from Fusion → Angle\n" + "Alignment. Uncheck to see the raw per-angle scan grid." + ) + self.chk_aligned_view.toggled.connect(self._on_aligned_view_toggled) + vl.addWidget(self.chk_aligned_view) + + self.btn_export_csv = QPushButton("Export Image as CSV…") + self.btn_export_csv.setEnabled(False) + self.btn_export_csv.setToolTip( + "Save the current CH1 image (one scan row per CSV line).") + self.btn_export_csv.clicked.connect(self._on_export_csv) + vl.addWidget(self.btn_export_csv) + + panel_layout.addWidget(grp_view) + + # ---- ROI --------------------------------------------------------- + grp_roi, rl = _group("ROI (Region of Interest)") + + self.btn_draw_roi = QPushButton("Draw ROI") + self.btn_draw_roi.setCheckable(True) + self.btn_draw_roi.setEnabled(False) + self.btn_draw_roi.setToolTip( + "Arm next click+drag on the image to draw a new ROI\n" + "(replaces any existing one). Click again to cancel.\n" + "After drawing, drag inside to move, or grab corners to reshape.\n" + "The ROI is persistent across channels / modes / angles." + ) + self.btn_draw_roi.toggled.connect(self._on_draw_roi_toggled) + rl.addWidget(self.btn_draw_roi) + + self.btn_clear_roi = QPushButton("Clear ROI") + self.btn_clear_roi.setEnabled(False) + self.btn_clear_roi.clicked.connect(self._on_clear_roi) + rl.addWidget(self.btn_clear_roi) + + self.btn_export_roi = QPushButton("Export ROI as CSV…") + self.btn_export_roi.setEnabled(False) + self.btn_export_roi.setToolTip( + "Save every pixel whose centre lies inside the ROI as CSV.\n" + "Columns: row, frame, x_mm, y_mm, value.\n" + "Corner coordinates of the quad are written in the file header." + ) + self.btn_export_roi.clicked.connect(self._on_export_roi_csv) + rl.addWidget(self.btn_export_roi) + + self.lbl_roi_center = _wrap_label("centroid: —", _CSS_HINT) + self.lbl_roi_size = _wrap_label("bbox: —", _CSS_HINT) + self.lbl_roi_npix = _wrap_label("pixels inside: —", _CSS_HINT) + for lbl in (self.lbl_roi_center, self.lbl_roi_size, self.lbl_roi_npix): + rl.addWidget(lbl) + + panel_layout.addWidget(grp_roi) + panel_layout.addStretch() + return _scroll_panel(panel, _LEFT_PANEL_W) + + def _build_canvases(self) -> QWidget: + splitter = QSplitter(Qt.Orientation.Vertical) + splitter.setChildrenCollapsible(False) + + img_widget = QWidget() + img_vl = QVBoxLayout(img_widget) + img_vl.setContentsMargins(0, 0, 0, 0) + img_vl.setSpacing(4) + self.image_canvas = ImageCanvas() + self.image_canvas.setMinimumHeight(220) + self.image_canvas.pixel_clicked.connect(self._on_pixel_clicked) + self.image_canvas.roi_changed.connect(self._update_roi_ui) + self.image_canvas.draw_mode_changed.connect(self._on_draw_mode_changed) + img_vl.addWidget(NavigationToolbar2QT(self.image_canvas, img_widget)) + img_vl.addWidget(self.image_canvas) + splitter.addWidget(img_widget) + + wave_widget = QWidget() + wave_vl = QVBoxLayout(wave_widget) + wave_vl.setContentsMargins(0, 0, 0, 0) + wave_vl.setSpacing(4) + self.lbl_wave_hint = QLabel( + "Click a pixel in the image above to inspect its waveform.") + self.lbl_wave_hint.setAlignment(Qt.AlignmentFlag.AlignCenter) + self.lbl_wave_hint.setStyleSheet(_CSS_MUTED) + self.wave_canvas = WaveformCanvas() + self.wave_canvas.setMinimumHeight(150) + wave_vl.addWidget(self.lbl_wave_hint) + wave_vl.addWidget(self.wave_canvas) + splitter.addWidget(wave_widget) + + splitter.setStretchFactor(0, 3) + splitter.setStretchFactor(1, 1) + splitter.setSizes([580, 250]) + return splitter + + def _build_right_panel(self) -> QWidget: + # Velocity settings (visible only in velocity mode) + self.grp_velocity, vel_l = _group("Velocity Settings (CH1 only)") + vel_form = _form() + self.spin_grating_um = QDoubleSpinBox() + self.spin_grating_um.setRange(0.1, 1000.0) + self.spin_grating_um.setDecimals(2) + self.spin_grating_um.setSingleStep(0.5) + self.spin_grating_um.setSuffix(" µm") + self.spin_grating_um.setValue(25) + self.spin_grating_um.setEnabled(False) + self.spin_grating_um.setMinimumWidth(_SPIN_MIN_W) + self.spin_grating_um.editingFinished.connect(self._on_grating_changed) + vel_form.addRow("Grating size:", self.spin_grating_um) + vel_l.addLayout(vel_form) + vel_l.addWidget(_wrap_label("v (m/s) = freq (MHz) × grating (µm)", + "font-size: 10px; color: #888;")) + self.grp_velocity.setVisible(False) + + grp_display, dl = _group("Display Options") + + cmap_form = _form() + self.combo_cmap = QComboBox() + self.combo_cmap.addItems(CMAPS) + self.combo_cmap.setCurrentText("gray") + self.combo_cmap.setEnabled(False) + self.combo_cmap.setSizePolicy(QSizePolicy.Policy.Expanding, + QSizePolicy.Policy.Fixed) + self.combo_cmap.currentIndexChanged.connect(self._on_cmap_changed) + cmap_form.addRow("Colormap:", self.combo_cmap) + dl.addLayout(cmap_form) + + self.chk_auto = QCheckBox("Auto-scale colormap") + self.chk_auto.setChecked(True) + self.chk_auto.toggled.connect(self._on_autoscale_toggled) + dl.addWidget(self.chk_auto) + + range_form = _form() + for label, attr in (("min:", "spin_vmin"), ("max:", "spin_vmax")): + spin = QDoubleSpinBox() + spin.setRange(-1e9, 1e9) + spin.setDecimals(4) + spin.setEnabled(False) + spin.setMinimumWidth(_SPIN_MIN_W) + spin.editingFinished.connect(self._on_manual_range_changed) + setattr(self, attr, spin) + range_form.addRow(label, spin) + dl.addLayout(range_form) + + right_panel = QWidget() + layout = QVBoxLayout(right_panel) + layout.setContentsMargins(0, 0, 0, 0) + layout.setSpacing(8) + layout.addWidget(self.grp_velocity) + layout.addWidget(grp_display) + layout.addStretch() + return _scroll_panel(right_panel, _RIGHT_PANEL_W) + + def _build_menus(self): + menubar = self.menuBar() + + fft_menu = menubar.addMenu("&FFT") + fft_act = QAction("FFT &Options…", self) + fft_act.setStatusTip("Configure FFT backend and zero-padding") + fft_act.triggered.connect(self._on_fft_options) + fft_menu.addAction(fft_act) + + fusion_menu = menubar.addMenu("&Fusion") + self._alignment_act = QAction("Angle &Alignment", self) + self._alignment_act.setStatusTip( + "Compute a rotation+translation alignment across all angles " + "(from CH4 masks) and enable Aligned View. Requires >1 angle.") + self._alignment_act.setEnabled(False) + self._alignment_act.triggered.connect(self._on_angle_alignment) + fusion_menu.addAction(self._alignment_act) + + self._manual_align_act = QAction("&Manual Alignment…", self) + self._manual_align_act.setStatusTip( + "Open an interactive dialog to align angles by eye: overlaid CH4 " + "threshold masks, keyboard nudge (translate + rotate), auto " + "de-rotate to the known scan angles, and save/clear a persistent " + "alignment.") + self._manual_align_act.setEnabled(False) + self._manual_align_act.triggered.connect(self._on_manual_alignment) + fusion_menu.addAction(self._manual_align_act) + + convert_menu = menubar.addMenu("&Convert") + self._batch_dc_act = QAction("Batch Compute DC and &Store…", self) + self._batch_dc_act.setStatusTip( + "Select .sras files and compute+store DC images (CH3/CH4 mean) " + "for every angle, converting v6 files to v7 in place.") + self._batch_dc_act.triggered.connect(lambda: self._on_batch_compute("dc")) + convert_menu.addAction(self._batch_dc_act) + + self._batch_fft_act = QAction("Batch Compute FFT and Sto&re…", self) + self._batch_fft_act.setStatusTip( + "Select .sras files and compute+store FFT peak-frequency images " + "for every angle, converting v6 files to v7 in place. Stored " + "images are natural-resolution (pad 1); padded views compute live.") + self._batch_fft_act.triggered.connect(lambda: self._on_batch_compute("fft")) + convert_menu.addAction(self._batch_fft_act) + + # ------------------------------------------------------------------ + # Drag-and-drop + # ------------------------------------------------------------------ + + def dragEnterEvent(self, event): + urls = event.mimeData().urls() + if urls and urls[0].toLocalFile().lower().endswith(".sras"): + event.acceptProposedAction() + + def dropEvent(self, event): + self._load_file(event.mimeData().urls()[0].toLocalFile()) + + # ------------------------------------------------------------------ + # File loading + # ------------------------------------------------------------------ + + def _on_open(self): + path, _ = QFileDialog.getOpenFileName( + self, "Open SRAS File", "", "SRAS Files (*.sras);;All Files (*)") + if path: + self._load_file(path) + + def _load_file(self, path: str): + started = self._run_worker( + "load", LoadWorker(path), + connect=( + ("finished", self._on_load_done), + ("error", lambda msg: self.statusBar().showMessage(f"Error: {msg}")), + ), + ) + if not started: + return + self.btn_open.setEnabled(False) + self.statusBar().showMessage(f"Loading {Path(path).name}…") + self._show_progress("main", f"Loading {Path(path).name}…") + + def _on_load_done(self, sras): + self._close_progress("main") + self.btn_open.setEnabled(True) + if sras is None: + return + self._sras = sras + self._current_image = None + + # A manual-alignment dialog bound to the previous file must not + # survive a reload — its per-angle state (and the sras it was + # constructed against) no longer matches the new file's geometry. + if self._manual_align_dialog is not None: + self._manual_align_dialog.close() + self._manual_align_dialog = None + + # Caches (and any in-flight DC precompute) belong to the previous + # file's geometry — discard and start fresh. Bumping the generation + # counters makes any still-running worker's result get dropped when + # it lands. + self._dc_cache = {} + self._fft_cache = {} + self._dc_generation += 1 + self.lbl_dc_precompute.setText("") + + self._alignment_result = None + self._aligned_cache = {} + self._alignment_generation += 1 + self.chk_aligned_view.blockSignals(True) + self.chk_aligned_view.setChecked(False) + self.chk_aligned_view.setEnabled(False) + self.chk_aligned_view.blockSignals(False) + + # Silently restore a previously-saved manual alignment, if any, so + # the work survives closing and reopening the file. + sidecar = load_manual_alignment(sras) + if sidecar is not None: + try: + self._alignment_result = build_manual_alignment( + sras, sidecar.ref_angle_idx, sidecar.dc_threshold_mv, + sidecar.per_angle) + self.chk_aligned_view.blockSignals(True) + self.chk_aligned_view.setChecked(True) + self.chk_aligned_view.blockSignals(False) + self.statusBar().showMessage( + f"Restored saved manual alignment from " + f"{sidecar_path(sras.path).name}") + except Exception as exc: + # A corrupt/foreign sidecar or a rescan that shrank n_angles + # below ref_angle_idx must not block opening the .sras file. + self.statusBar().showMessage( + f"Could not restore saved alignment: {exc}") + + # A ROI from the previous file no longer matches the new scan's + # geometry, so discard it on every load. + self.image_canvas.clear_roi() + + self.lbl_filename.setText(sras.path.name) + + self.spin_angle.blockSignals(True) + self.spin_angle.setRange(0, max(0, sras.n_angles - 1)) + self.spin_angle.setValue(0) + self.spin_angle.blockSignals(False) + + # DC channels are cheap and give an instant, fluid overview of a + # scan; CH1/Velocity require an FFT per pixel that can take minutes + # on a large scan, so don't default to it. + self.combo_channel.blockSignals(True) + self.combo_channel.setCurrentIndex(CH4_IDX) + self.combo_channel.blockSignals(False) + + self._update_controls_enabled(True) + self._on_threshold_changed() # refresh ADC label with file calibration + self._on_view_changed() + self._start_dc_precompute() + + # ------------------------------------------------------------------ + # Scan info panel + # ------------------------------------------------------------------ + + def _update_scan_info_labels(self): + s = self._sras + if s is None: + return + a = self.spin_angle.value() + for key, text in ( + ("Angles", f"{s.n_angles}"), + ("Rows", f"{s.n_rows[a]}"), + ("Frames / row", f"{s.n_frames[a]}"), + ("Samples / frame", f"{s.samples_per_frame}"), + ("Sample rate", f"{s.sample_rate_hz / 1e9:.4g} GS/s"), + ("X start", f"{s.x_start_mm[a]:.4g} mm"), + ("Pixel Δx", f"{s.pixel_x_mm * 1e3:.3g} µm"), + ("Laser freq", f"{s.laser_freq_hz / 1e3:.4g} kHz"), + ): + self._info[key].setText(f"{key}: {text}") + + notes = [] + if s.frame_count_mismatch: + notes.append(f"! Header n_frames={s.n_frames_header}, " + f"actual={s.n_frames[a]} (scanner bug — corrected)") + if s.scan_aborted: + notes.append(f"! Scan aborted: {s.n_angles}/{s.n_angles_declared} " + "angles complete") + if s.background is not None: + notes.append(f"Background waveform: {len(s.background)} samples") + if s.version in (6, 7): + notes.append("v6/v7 format: rows / frames / x_start are per-angle") + + n_dc = sum(1 for x in s.precomputed_dc4_mv if x is not None) + n_fft = sum(1 for x in s.precomputed_freq_mhz if x is not None) + if n_dc or n_fft: + bg_note = " (bg-sub)" if s.precomputed_bg_sub else " (no bg-sub)" + notes.append( + f"Cached images: DC {n_dc}/{s.n_angles} angles, " + f"FFT {n_fft}/{s.n_angles} angles{bg_note if n_fft else ''} " + "— display is instant for cached angles") + elif s.version == 7: + notes.append("v7 format: no cache blocks stored yet") + self.lbl_frame_warn.setText("\n".join(notes)) + + # ------------------------------------------------------------------ + # Controls + # ------------------------------------------------------------------ + + def _update_controls_enabled(self, enabled: bool): + s = self._sras + has_file = enabled and s is not None + ch_idx = self.combo_channel.currentIndex() + is_ch1 = enabled and ch_idx in CH1_DERIVED_MODES + is_vel = enabled and ch_idx == VELOCITY_MODE_IDX + + self.spin_angle.setEnabled(has_file and s.n_angles > 1) + self.combo_channel.setEnabled(enabled) + self.combo_cmap.setEnabled(enabled) + self.chk_auto.setEnabled(enabled) + manual = enabled and not self.chk_auto.isChecked() + self.spin_vmin.setEnabled(manual) + self.spin_vmax.setEnabled(manual) + + # Threshold and bg-sub apply to all CH1 modes + self.spin_threshold_mv.setEnabled(is_ch1) + self.chk_bg_sub.setEnabled(has_file and s.background is not None and is_ch1) + self.spin_grating_um.setEnabled(is_vel) + self.grp_velocity.setVisible(is_vel) + + self.btn_export_csv.setEnabled(is_ch1 and self._current_image is not None) + # ROI: always usable once a file is loaded (independent of channel) + self.btn_draw_roi.setEnabled(has_file) + + # Batch Convert actions pick their own files, independent of + # whatever's currently open — only gated on no batch already running. + can_batch = not self._job_running("batch") + self._batch_dc_act.setEnabled(can_batch) + self._batch_fft_act.setEnabled(can_batch) + + self._alignment_act.setEnabled( + has_file and s.n_angles > 1 and not self._job_running("align")) + self._manual_align_act.setEnabled( + has_file and s.n_angles > 1 and not self._job_running("align")) + self.chk_aligned_view.setEnabled(enabled and self._alignment_result is not None) + self._update_roi_ui() + + def _on_channel_changed(self): + self._update_controls_enabled(self._sras is not None) + self._on_view_changed() + + def _on_bg_sub_toggled(self): + # Background subtraction changes the FFT input, so it genuinely + # invalidates the cached raw FFT (the cache key includes it) — + # _refresh_display() recomputes only on a miss for the new state. + if self._sras is not None and self.combo_channel.currentIndex() in CH1_DERIVED_MODES: + self._refresh_display() + + def _on_grating_changed(self): + # Grating is a pure post-multiply on the cached frequency image — + # never needs a recompute. + if self._sras is not None and self.combo_channel.currentIndex() == VELOCITY_MODE_IDX: + self._refresh_display() + + def _on_threshold_changed(self): + mv = self.spin_threshold_mv.value() + cal = (self._sras.cal(CH4_IDX) if self._sras is not None + else (_FALLBACK_YMULT_MV, _FALLBACK_YOFF_ADC, 0.0)) + self.lbl_threshold_adc.setText(f"≈ {mv_to_adc(mv, *cal):.1f} ADC counts") + # Threshold decides which pixels get an FFT at all, so changing it is + # a genuine cache-key change — but the recompute reuses the cached DC4 + # image to skip masked-out pixels. + if self._sras is not None and self.combo_channel.currentIndex() in CH1_DERIVED_MODES: + self._refresh_display() + + def _on_autoscale_toggled(self, checked: bool): + manual = not checked + self.spin_vmin.setEnabled(manual and self._sras is not None) + self.spin_vmax.setEnabled(manual and self._sras is not None) + if self._sras is not None and self._current_image is not None: + self._redraw_image(self._current_image) + + def _on_manual_range_changed(self): + if not self.chk_auto.isChecked() and self._current_image is not None: + self._redraw_image(self._current_image) + + def _on_cmap_changed(self): + # Colormap is purely how the existing image is rendered. + if self._current_image is not None: + self._redraw_image(self._current_image) + + def _on_view_changed(self): + if self._sras is None: + return + idx = self.spin_angle.value() + self.lbl_angle_deg.setText(f"({self._sras.angles_deg[idx]:.1f}°)") + self._update_scan_info_labels() + self._refresh_display() + + def _on_aligned_view_toggled(self, checked: bool): + if self._current_image is not None: + self._redraw_image(self._current_image) + + # ------------------------------------------------------------------ + # CSV export + # ------------------------------------------------------------------ + + def _on_export_csv(self): + if self._current_image is None or self._sras is None: + return + default_name = (f"{self._sras.path.stem}_angle{self._current_angle}" + f"_{CH_NAMES[self._current_ch]}.csv") + path, _ = QFileDialog.getSaveFileName( + self, "Export Image as CSV", + str(self._sras.path.parent / default_name), + "CSV files (*.csv);;All files (*)") + if not path: + return + np.savetxt(path, self._current_image, delimiter=",", fmt="%.6g") + self.statusBar().showMessage(f"Exported {Path(path).name}") + + def _on_export_roi_csv(self): + if self._current_image is None or self._sras is None: + return + roi = self.image_canvas.get_roi() + if roi is None: + self.statusBar().showMessage("No ROI — draw one first") + return + s = self._sras + x_axis = s.x_axis_mm(self._current_angle) + y_axis = s.y_positions_mm(self._current_angle) + mask = roi.mask_for_grid(x_axis, y_axis) + if not mask.any(): + self.statusBar().showMessage("ROI does not overlap any pixel") + return + img = self._current_image + if img.shape != mask.shape: + self.statusBar().showMessage( + f"ROI shape {mask.shape} does not match image {img.shape}") + return + + X, Y = np.meshgrid(np.asarray(x_axis, dtype=np.float64), + np.asarray(y_axis, dtype=np.float64)) + rows_idx, frames_idx = np.where(mask) + n_pix = int(mask.sum()) + + ch_name = CH_NAMES[self._current_ch] + angle = self._current_angle + default_name = f"{s.path.stem}_angle{angle}_{ch_name}_ROI.csv" + path, _ = QFileDialog.getSaveFileName( + self, "Export ROI as CSV", + str(s.path.parent / default_name), + "CSV files (*.csv);;All files (*)") + if not path: + return + + corners_str = " ".join(f"({p[0]:.6g},{p[1]:.6g})" for p in roi.corners()) + header = ( + f"# ROI quad corners (BL BR TR TL) mm: {corners_str}\n" + f"# source: {s.path.name}, channel={ch_name}, " + f"angle_idx={angle}, angle_deg={s.angles_deg[angle]:.4g}\n" + f"# n_pixels={n_pix}\n" + "row,frame,x_mm,y_mm,value" + ) + data = np.column_stack([ + rows_idx.astype(np.int64), frames_idx.astype(np.int64), + X[mask], Y[mask], img[mask].astype(np.float64), + ]) + # integer columns first, floats after — use a per-column format list + np.savetxt(path, data, delimiter=",", + fmt=["%d", "%d", "%.6g", "%.6g", "%.6g"], + header=header, comments="") + self.statusBar().showMessage( + f"Exported ROI ({n_pix} pixels) to {Path(path).name}") + + # ------------------------------------------------------------------ + # ROI + # ------------------------------------------------------------------ + + def _on_draw_roi_toggled(self, checked: bool): + if checked: + self.image_canvas.start_drawing() + self.statusBar().showMessage( + "Click and drag on the image to draw a new rectangle.") + else: + self.image_canvas.cancel_drawing() + + def _on_draw_mode_changed(self, active: bool): + # Keep the toggle button's visual state in sync with the canvas. + self.btn_draw_roi.blockSignals(True) + self.btn_draw_roi.setChecked(active) + self.btn_draw_roi.blockSignals(False) + + def _on_clear_roi(self): + self.image_canvas.clear_roi() + self.statusBar().showMessage("ROI cleared") + + def _update_roi_ui(self): + roi = self.image_canvas.get_roi() + if roi is None: + self.lbl_roi_center.setText("centroid: —") + self.lbl_roi_size.setText("bbox: —") + self.lbl_roi_npix.setText("pixels inside: —") + self.btn_clear_roi.setEnabled(False) + self.btn_export_roi.setEnabled(False) + return + + cen = roi.centroid() + bbox = roi.bbox_size() + self.lbl_roi_center.setText(f"centroid: ({cen[0]:.3f}, {cen[1]:.3f}) mm") + self.lbl_roi_size.setText(f"bbox: {bbox[0]:.3f} × {bbox[1]:.3f} mm") + + npix = 0 + if self._sras is not None: + try: + # Deliberately always the raw per-angle grid, even when + # Aligned View is on: _on_export_roi_csv also exports on the + # raw grid (never synthetically-resampled pixels), so this + # readout must match what Export ROI actually writes. + mask = roi.mask_for_grid( + self._sras.x_axis_mm(self._current_angle), + self._sras.y_positions_mm(self._current_angle)) + npix = int(mask.sum()) + except Exception: + npix = 0 + self.lbl_roi_npix.setText(f"pixels inside: {npix}") + self.btn_clear_roi.setEnabled(True) + self.btn_export_roi.setEnabled(self._current_image is not None and npix > 0) + + # ------------------------------------------------------------------ + # Display + # ------------------------------------------------------------------ + + def _current_n_fft(self) -> int | None: + if self._fft_pad_factor <= 1 or self._sras is None: + return None + return self._sras.samples_per_frame * self._fft_pad_factor + + def _scale_for_display(self, freq_mhz: np.ndarray, ch_idx: int) -> np.ndarray: + """Velocity is a pure post-multiply of the (already DC-masked) + cached frequency image — never worth a recompute on its own.""" + if ch_idx == VELOCITY_MODE_IDX: + return freq_mhz * self.spin_grating_um.value() + return freq_mhz + + def _fft_cache_key(self, angle_idx: int) -> tuple: + return (angle_idx, self.chk_bg_sub.isChecked(), self._current_n_fft(), + self.spin_threshold_mv.value()) + + def _aligned_cache_key(self, angle_idx: int, ch_idx: int) -> tuple: + """Mirrors _fft_cache's key granularity so a stale aligned image is + never shown after bg_sub/threshold/pad/grating changes.""" + if ch_idx in CH1_DERIVED_MODES: + return (*self._fft_cache_key(angle_idx), ch_idx, + self.spin_grating_um.value() if ch_idx == VELOCITY_MODE_IDX else None) + return (angle_idx, ch_idx) + + def _aligned_canvas_axes(self) -> tuple[np.ndarray, np.ndarray]: + r = self._alignment_result + n_rows, n_cols = r.canvas_shape + return (r.canvas_origin_mm[0] + np.arange(n_cols) * r.canvas_dx_mm, + r.canvas_origin_mm[1] + np.arange(n_rows) * r.canvas_dy_mm) + + def _get_aligned_display_image(self, raw_img: np.ndarray, angle_idx: int, + ch_idx: int) -> np.ndarray: + key = self._aligned_cache_key(angle_idx, ch_idx) + cached = self._aligned_cache.get(key) + if cached is None: + cached = apply_alignment(self._alignment_result, angle_idx, raw_img) + self._aligned_cache[key] = cached + return cached + + def _refresh_display(self): + """Show the image for the current angle/channel/threshold, using + cached data whenever possible and only falling back to a background + compute (with progress popup) when genuinely nothing is cached yet.""" + if self._sras is None: + return + angle_idx = self.spin_angle.value() + ch_idx = self.combo_channel.currentIndex() + + if ch_idx in CH1_DERIVED_MODES: + raw = self._fft_cache.get(self._fft_cache_key(angle_idx)) + if raw is not None: + self._show_image_now(self._scale_for_display(raw, ch_idx), + angle_idx, ch_idx) + return + else: + cached = self._dc_cache.get((angle_idx, ch_idx)) + if cached is not None: + self._show_image_now(cached, angle_idx, ch_idx) + return + + # Nothing cached for these settings — need a real compute. Changing + # the DC threshold changes *which* pixels get an FFT at all, so it + # can't be satisfied from the cache — but with the DC map already + # known, the recompute skips the FFT for masked-out pixels. + self._start_compute() + + def _show_image_now(self, img: np.ndarray, angle_idx: int, ch_idx: int): + """Display an already-available image with no compute involved.""" + self._current_image = img + self._current_angle = angle_idx + self._current_ch = ch_idx + self.btn_export_csv.setEnabled(ch_idx in CH1_DERIVED_MODES) + self._redraw_image(img) + self._update_roi_ui() + + def _redraw_image(self, img: np.ndarray): + s = self._sras + angle_idx = self._current_angle + ch_idx = self._current_ch + + aligned = (self.chk_aligned_view.isChecked() + and self._alignment_result is not None + and angle_idx in self._alignment_result.per_angle) + if aligned: + display_img = self._get_aligned_display_image(img, angle_idx, ch_idx) + x_axis, y_axis = self._aligned_canvas_axes() + else: + display_img = img + x_axis = s.x_axis_mm(angle_idx) + y_axis = s.y_positions_mm(angle_idx) + + dx = x_axis[1] - x_axis[0] if len(x_axis) > 1 else s.pixel_x_mm + dy = float(y_axis[1] - y_axis[0]) if len(y_axis) > 1 else 1.0 + extent = [x_axis[0] - dx / 2, x_axis[-1] + dx / 2, + y_axis[-1] + dy / 2, y_axis[0] - dy / 2] + + if self.chk_auto.isChecked(): + vmin, vmax = float(display_img.min()), float(display_img.max()) + for spin, val in ((self.spin_vmin, vmin), (self.spin_vmax, vmax)): + spin.blockSignals(True) + spin.setValue(val) + spin.blockSignals(False) + else: + vmin, vmax = self.spin_vmin.value(), self.spin_vmax.value() + + angle_deg = s.angles_deg[angle_idx] + mode_str, unit, colorbar_label = _CHANNEL_DISPLAY[ch_idx] + if ch_idx == VELOCITY_MODE_IDX: + ch_label = f"Velocity [grating={self.spin_grating_um.value():.2f} µm]" + else: + ch_label = CH_LABELS[ch_idx] + + title = f"{CH_NAMES[ch_idx]} | {mode_str} | {angle_deg:.1f}°" + if aligned: + title += " [Aligned]" + + self.image_canvas.show_image( + display_img, extent, + cmap=self.combo_cmap.currentText(), + vmin=vmin, vmax=vmax, + xlabel="X (mm)", ylabel="Y (mm)", + title=title, colorbar_label=colorbar_label, + ) + self.statusBar().showMessage( + f"{s.path.name} | {ch_label} @ {angle_deg:.1f}° " + f"| {display_img.shape[1]} × {display_img.shape[0]} px | {unit}" + f"{' | Aligned' if aligned else ''}" + ) + + # ------------------------------------------------------------------ + # Background compute (only reached on a genuine cache miss) + # ------------------------------------------------------------------ + + def _start_compute(self): + if self._sras is None or self._job_running("compute"): + return # re-checked when the running compute finishes + + angle_idx = self.spin_angle.value() + ch_idx = self.combo_channel.currentIndex() + is_fft = ch_idx in CH1_DERIVED_MODES + + self._pending_angle = angle_idx + self._pending_ch = ch_idx + self._pending_bg_sub = self.chk_bg_sub.isChecked() + self._pending_threshold = self.spin_threshold_mv.value() + self._pending_fft_pad_factor = self._fft_pad_factor + + worker = ComputeWorker( + self._sras, angle_idx, ch_idx, + apply_bg_sub=self._pending_bg_sub, + n_fft=self._current_n_fft(), + dc_threshold_mv=self._pending_threshold, + # Reuse the cached DC4 image (if the precompute has reached this + # angle) so the FFT skips masked-out pixels entirely and doesn't + # need to re-read the CH4 channel from disk. + dc4_mv=self._dc_cache.get((angle_idx, CH4_IDX)), + is_fft_mode=is_fft, + ) + if not self._run_worker( + "compute", worker, + connect=( + ("finished", self._on_compute_done), + ("error", lambda msg: self.statusBar().showMessage( + f"Compute error: {msg}")), + ), + on_done=self._after_compute): + return + + if is_fft: + self.statusBar().showMessage("Computing FFT…") + self._show_progress( + "main", + f"Computing FFT for angle {angle_idx}…\n" + "This can take a while on a large scan — result is cached " + "so revisiting this angle/mode/threshold will be instant.") + else: + self.statusBar().showMessage("Computing DC image…") + self._show_progress("main", f"Computing DC image for angle {angle_idx}…") + + def _after_compute(self): + """If settings changed while the compute was running, re-dispatch + through the cache-aware path — the now-current combination may + already be cached.""" + if (self.spin_angle.value(), self.combo_channel.currentIndex(), + self.chk_bg_sub.isChecked(), self.spin_threshold_mv.value(), + self._fft_pad_factor) != ( + self._pending_angle, self._pending_ch, self._pending_bg_sub, + self._pending_threshold, self._pending_fft_pad_factor): + self._refresh_display() + + def _on_compute_done(self, result): + self._close_progress("main") + if result is None: + return # cancelled mid-compute; the partial image must not cache + angle_idx = self._pending_angle + ch_idx = self._pending_ch + + if ch_idx in CH1_DERIVED_MODES: + self._fft_cache[(angle_idx, self._pending_bg_sub, + self._current_n_fft(), self._pending_threshold)] = result + img = self._scale_for_display(result, ch_idx) + else: + img = result + self._dc_cache[(angle_idx, ch_idx)] = img + + self._show_image_now(img, angle_idx, ch_idx) + + # ------------------------------------------------------------------ + # Background DC precompute (all angles, so switching is fluid) + # ------------------------------------------------------------------ + + def _start_dc_precompute(self): + if self._sras is None: + return + generation = self._dc_generation + n_angles = self._sras.n_angles + + worker = DcPrecomputeWorker(self._sras) + self._run_worker( + "dc_precompute", worker, + connect=( + ("angle_done", lambda a, dc3, dc4, g=generation: + self._on_dc_precompute_angle_done(g, a, dc3, dc4, n_angles)), + ("error", lambda msg: self.statusBar().showMessage( + f"DC precompute error: {msg}", 5000)), + ), + quit_on=("finished", "error"), + ) + + def _on_dc_precompute_angle_done(self, generation: int, angle_idx: int, + dc3_mv: np.ndarray, dc4_mv: np.ndarray, + n_angles: int): + if generation != self._dc_generation: + return # stale result from a previously-loaded file — discard + self._dc_cache[(angle_idx, CH3_IDX)] = dc3_mv + self._dc_cache[(angle_idx, CH4_IDX)] = dc4_mv + + done = sum(1 for a in range(n_angles) if (a, CH4_IDX) in self._dc_cache) + self.lbl_dc_precompute.setText( + f"Precomputing DC images: {done}/{n_angles} angles ready…" + if done < n_angles else "DC images ready for all angles.") + + # If we just finished the angle/channel the user is currently looking + # at and it wasn't shown yet (they switched here before the precompute + # caught up and are still waiting), show it now. + current_ch = self.combo_channel.currentIndex() + if (angle_idx == self.spin_angle.value() + and not self._job_running("compute") + and current_ch in (CH3_IDX, CH4_IDX) + and (self._current_angle != angle_idx or self._current_ch != current_ch)): + self._refresh_display() + + # ------------------------------------------------------------------ + # Pixel inspector + # ------------------------------------------------------------------ + + def _on_pixel_clicked(self, row_idx: int, frame_idx: int): + if self._sras is None or self._current_image is None: + return + angle_idx = self._current_angle + if (self.chk_aligned_view.isChecked() and self._alignment_result is not None + and angle_idx in self._alignment_result.per_angle): + # The click landed on the shared aligned canvas — invert the same + # canvas->raw affine used to display it back to a raw (row, frame) + # index before looking up the waveform. + t = self._alignment_result.per_angle[angle_idx] + raw = t.matrix @ np.array([row_idx, frame_idx], dtype=np.float64) + t.offset + row_idx, frame_idx = int(round(raw[0])), int(round(raw[1])) + n_rows_a, n_frames_a = self._sras.image_shape(angle_idx) + if not (0 <= row_idx < n_rows_a and 0 <= frame_idx < n_frames_a): + self.statusBar().showMessage( + "No source waveform here (padding region of the aligned canvas).") + return + + self.lbl_wave_hint.hide() + if self._current_ch in CH1_DERIVED_MODES: + self.wave_canvas.show_rf_waveform( + self._sras, angle_idx, row_idx, frame_idx, + apply_bg_sub=self.chk_bg_sub.isChecked()) + else: + self.wave_canvas.show_dc_waveform( + self._sras, angle_idx, self._current_ch, row_idx, frame_idx) + + # ------------------------------------------------------------------ + # Progress dialogs + # ------------------------------------------------------------------ + + def _show_progress(self, key: str, message: str, maximum: int = 0): + """Show (or relabel) the progress dialog under *key*. maximum=0 gives + an indeterminate busy indicator.""" + dlg = self._progress_dlgs.get(key) + if dlg is not None: + dlg.setLabelText(message) + return + dlg = QProgressDialog(message, "", 0, maximum, self) + dlg.setWindowTitle("Please wait…") + dlg.setCancelButton(None) + dlg.setWindowModality(Qt.WindowModality.WindowModal) + dlg.setMinimumDuration(300) # only appears if it takes > 300 ms + dlg.show() + self._progress_dlgs[key] = dlg + + def _set_progress(self, key: str, pct: int): + dlg = self._progress_dlgs.get(key) + if dlg is not None: + dlg.setValue(pct) + + def _close_progress(self, key: str): + dlg = self._progress_dlgs.pop(key, None) + if dlg is not None: + dlg.close() + + # ------------------------------------------------------------------ + # Convert menu: batch DC/FFT compute-and-store (v6 -> v7) + # ------------------------------------------------------------------ + + def _on_batch_compute(self, mode: str): + if self._job_running("batch"): + return + label = "DC" if mode == "dc" else "FFT" + paths, _ = QFileDialog.getOpenFileNames( + self, f"Select .sras files to batch-compute {label}", "", + "SRAS files (*.sras);;All files (*)") + if not paths: + return + + self._batch_errors = [] + worker = BatchCacheWorker(paths, mode, self.chk_bg_sub.isChecked()) + started = self._run_worker( + "batch", worker, + connect=( + ("progress", lambda pct: self._set_progress("batch", pct)), + ("file_done", self._on_batch_file_done), + ("finished", lambda p=paths: self._on_batch_finished(p)), + ), + on_done=self._after_batch, + ) + if not started: + return # a second trigger snuck in while the file dialog was open + + self._batch_dc_act.setEnabled(False) + self._batch_fft_act.setEnabled(False) + self._show_progress( + "batch", f"Batch computing {label} for {len(paths)} file(s)…", + maximum=100) + + def _on_batch_file_done(self, path: str, err: str): + if err: + self._batch_errors.append(f"{Path(path).name} — {err}") + self._show_progress("batch", f"Processed {Path(path).name}…") + + def _on_batch_finished(self, paths: list[str]): + self._close_progress("batch") + + n_total = len(paths) + n_failed = len(self._batch_errors) + n_ok = n_total - n_failed + if n_failed: + summary = (f"Batch store: {n_ok}/{n_total} file(s) updated, " + f"{n_failed} failed: {'; '.join(self._batch_errors)}") + else: + summary = f"Batch store: {n_ok}/{n_total} file(s) updated." + self.statusBar().showMessage(summary) + self._batch_errors = [] + + # If the currently-open file was in this batch, reload it so the GUI + # picks up the newly-written v7 cache instead of stale state. + if self._sras is not None and str(self._sras.path) in paths: + self._load_file(str(self._sras.path)) + + def _after_batch(self): + self._batch_dc_act.setEnabled(True) + self._batch_fft_act.setEnabled(True) + + # ------------------------------------------------------------------ + # Fusion: angle alignment + # ------------------------------------------------------------------ + + def _on_angle_alignment(self): + if self._sras is None or self._sras.n_angles <= 1: + return + ref_idx = 0 + threshold_mv = self.spin_threshold_mv.value() + generation = self._alignment_generation + + started = self._run_worker( + "align", AngleAlignmentWorker(self._sras, ref_idx, threshold_mv), + connect=( + ("progress", lambda pct: self._set_progress("main", pct)), + ("finished", lambda result, err, g=generation: + self._on_alignment_done(g, result, err)), + ), + on_done=lambda: self._update_controls_enabled(self._sras is not None), + ) + if not started: + return + + self._alignment_act.setEnabled(False) + self._show_progress( + "main", + f"Computing angle alignment ({self._sras.n_angles} angles, " + f"ref=angle 0, CH4 mask ≥ {threshold_mv:.3f} mV)…", + maximum=100) + + def _on_alignment_done(self, generation: int, result, error_msg: str): + self._close_progress("main") + if generation != self._alignment_generation: + return # a new file was loaded while this was computing — discard + if error_msg: + self.statusBar().showMessage(f"Angle alignment failed: {error_msg}") + return + self._alignment_result = result + self._aligned_cache = {} + self.chk_aligned_view.setEnabled(True) + self.chk_aligned_view.blockSignals(True) + self.chk_aligned_view.setChecked(True) + self.chk_aligned_view.blockSignals(False) + nr, nc = result.canvas_shape + self.statusBar().showMessage( + f"Angle alignment computed ({self._sras.n_angles} angles, " + f"canvas {nc}×{nr} px).") + self._refresh_display() + + # ------------------------------------------------------------------ + # Fusion: manual alignment + # ------------------------------------------------------------------ + + def _on_manual_alignment(self): + if self._sras is None or self._sras.n_angles <= 1: + return + if self._manual_align_dialog is not None: + self._manual_align_dialog.raise_() + self._manual_align_dialog.activateWindow() + return + + ref_idx = 0 + threshold_mv = self.spin_threshold_mv.value() + seed: dict[int, ManualAngleParams] = {} + # Seed only from a previously *saved manual* alignment (this dialog's + # own Save also writes this sidecar) -- never from self._alignment_result + # when it holds the automatic Fusion -> Angle Alignment's output. That + # path's translation comes from FFT phase correlation, which is the + # very thing manual mode exists to work around; inheriting it here + # would silently reintroduce the same bad translations under a + # "manual" label, on top of the (correct) analytic rotation, which is + # exactly what makes manual mode look like it "still does the same + # thing" the automatic one does. + sidecar = load_manual_alignment(self._sras) + if sidecar is not None and sidecar.ref_angle_idx == ref_idx: + seed = dict(sidecar.per_angle) + threshold_mv = sidecar.dc_threshold_mv + + cached_dc4 = {a: img for (a, ch), img in self._dc_cache.items() if ch == CH4_IDX} + dlg = ManualAlignmentDialog( + self, self._sras, ref_angle_idx=ref_idx, dc_threshold_mv=threshold_mv, + seed_per_angle=seed, cached_dc4_mv=cached_dc4) + dlg.alignment_saved.connect(self._on_manual_alignment_saved) + dlg.alignment_cleared.connect(self._on_manual_alignment_cleared) + dlg.finished.connect(self._on_manual_align_dialog_closed) + dlg.setAttribute(Qt.WidgetAttribute.WA_DeleteOnClose) + self._manual_align_dialog = dlg + dlg.show() + + def _on_manual_align_dialog_closed(self, _result_code: int): + self._manual_align_dialog = None + + def _on_manual_alignment_saved(self, result, sidecar_path_str: str): + self._alignment_result = result + self._aligned_cache = {} + self._alignment_generation += 1 + self.chk_aligned_view.setEnabled(True) + self.chk_aligned_view.blockSignals(True) + self.chk_aligned_view.setChecked(True) + self.chk_aligned_view.blockSignals(False) + self._update_controls_enabled(self._sras is not None) + self.statusBar().showMessage( + f"Manual alignment saved to {Path(sidecar_path_str).name}") + if self._current_image is not None: + self._refresh_display() + + def _on_manual_alignment_cleared(self): + self._alignment_result = None + self._aligned_cache = {} + self._alignment_generation += 1 + self.chk_aligned_view.blockSignals(True) + self.chk_aligned_view.setChecked(False) + self.chk_aligned_view.setEnabled(False) + self.chk_aligned_view.blockSignals(False) + self._update_controls_enabled(self._sras is not None) + self.statusBar().showMessage("Manual alignment cleared.") + if self._current_image is not None: + self._refresh_display() + + # ------------------------------------------------------------------ + # FFT Options + # ------------------------------------------------------------------ + + def _on_fft_options(self): + dlg = FftOptionsDialog( + self, + current_backend=compute.get_fft_backend(), + current_pad_factor=self._fft_pad_factor, + samples_per_frame=self._sras.samples_per_frame if self._sras else None, + sample_rate_hz=self._sras.sample_rate_hz if self._sras else None, + grating_um=self.spin_grating_um.value(), + ) + if dlg.exec() != QDialog.DialogCode.Accepted: + return + compute.set_fft_backend(dlg.get_backend()) + self._fft_pad_factor = dlg.get_pad_factor() + self._settings.setValue("fft/backend", compute.get_fft_backend()) + self._settings.setValue("fft/pad_factor", self._fft_pad_factor) + # Pad factor changes the FFT bin count, so it genuinely invalidates + # the cached raw FFT (part of the cache key) — _refresh_display() + # recomputes only on a cache miss. + if self._sras is not None and self.combo_channel.currentIndex() in CH1_DERIVED_MODES: + self._refresh_display() + + # ------------------------------------------------------------------ + + def closeEvent(self, event): + if self._manual_align_dialog is not None: + self._manual_align_dialog.close() + + # Signal every cancellable worker first, then wait. Waiting without + # signalling means sitting out whatever is in flight — on a large + # scan a single angle is ~40 s. + jobs = list(self._jobs.values()) + for _thread, worker, _on_done in jobs: + stop = getattr(worker, "stop", None) + if callable(stop): + stop() + for thread, _worker, _on_done in jobs: + thread.quit() + thread.wait(5000) + super().closeEvent(event) + + +# --------------------------------------------------------------------------- + +def main(): + app = QApplication(sys.argv) + window = SrasViewerWindow( + initial_path=sys.argv[1] if len(sys.argv) > 1 else None) + window.show() + sys.exit(app.exec()) + diff --git a/tests/test_gui.py b/tests/test_gui.py index cb0bc28..ada0b73 100644 --- a/tests/test_gui.py +++ b/tests/test_gui.py @@ -186,7 +186,7 @@ def test_roi_and_csv_export(ctx): assert win.btn_export_roi.isEnabled(), "Export ROI enabled" csv_path = ctx.tmpdir / "roi.csv" - with patch("sras_viewer.QFileDialog.getSaveFileName", + with patch("sras_viewer.main_window.QFileDialog.getSaveFileName", return_value=(str(csv_path), "")): win._on_export_roi_csv() assert csv_path.exists(), "ROI CSV written" @@ -195,7 +195,7 @@ def test_roi_and_csv_export(ctx): f"ROI CSV has {len(body)} lines for {npix} pixels (want header + one per pixel)" img_csv = ctx.tmpdir / "img.csv" - with patch("sras_viewer.QFileDialog.getSaveFileName", + with patch("sras_viewer.main_window.QFileDialog.getSaveFileName", return_value=(str(img_csv), "")): win._on_export_csv() assert img_csv.exists(), "image CSV written" @@ -426,7 +426,7 @@ def test_stale_schema_sidecar_ignored(ctx): def test_clear_with_confirmation(ctx): win, dlg, s = ctx.win, ctx.dlg, ctx.s - with patch("sras_viewer.QMessageBox.question", + with patch("sras_viewer.dialogs.QMessageBox.question", return_value=QMessageBox.StandardButton.Yes): dlg._on_clear() assert not ctx.sidecar.exists(), "sidecar file deleted"