Thomas Ales
3989c2a1b8
Add row-averaged FFT feature with configurable window size (row_avg_n parameter) for improved signal-to-noise ratio on noisy scans. Includes: - Gaussian-weighted same-row neighbor averaging (never crosses rows) - Masked/renormalized convolution handling for edge cases and masked samples - Cache format v2 with row_avg_n tracking to prevent silent cache mismatches - GUI dialog option for row-average window configuration - Comprehensive tests validating kernel properties, background subtraction invariance, and cache dispatch Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>
1447 lines
62 KiB
Python
1447 lines
62 KiB
Python
"""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, QSignalBlocker, Qt, QThread
|
||
from PyQt6.QtGui import QAction
|
||
from PyQt6.QtWidgets import (
|
||
QApplication, QCheckBox, QComboBox, QDialog, 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, _axes_extent, _CSS_HINT, _CSS_INFO, _CSS_MUTED, _CSS_WARN, _LEFT_PANEL_W,
|
||
_RIGHT_PANEL_W, Jobs, _form, _group, _make_dspin, _scroll_panel, _wrap_label,
|
||
)
|
||
from .dialogs import FftOptionsDialog, ManualAlignmentDialog, RowAverageFftOptionsDialog
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 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
|
||
try:
|
||
row_avg_n = int(self._settings.value("fft/row_avg_n", 3))
|
||
except (TypeError, ValueError):
|
||
row_avg_n = 3
|
||
self._pending_row_avg_n: int = max(1, min(50, row_avg_n))
|
||
|
||
# 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)
|
||
panel_layout.addWidget(self._build_file_group())
|
||
panel_layout.addWidget(self._build_info_group())
|
||
panel_layout.addWidget(self._build_view_group())
|
||
panel_layout.addWidget(self._build_roi_group())
|
||
panel_layout.addStretch()
|
||
return _scroll_panel(panel, _LEFT_PANEL_W)
|
||
|
||
def _build_file_group(self) -> QWidget:
|
||
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)
|
||
return grp_file
|
||
|
||
def _build_info_group(self) -> QWidget:
|
||
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)
|
||
return grp_info
|
||
|
||
def _build_view_group(self) -> QWidget:
|
||
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 = _make_dspin(-500.0, 500.0, 3, suffix=" mV",
|
||
value=50.0, step=0.025)
|
||
self.spin_threshold_mv.setEnabled(False)
|
||
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)
|
||
return grp_view
|
||
|
||
def _build_roi_group(self) -> QWidget:
|
||
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)
|
||
return grp_roi
|
||
|
||
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 = _make_dspin(0.1, 1000.0, 2, suffix=" µm",
|
||
value=25, step=0.5)
|
||
self.spin_grating_um.setEnabled(False)
|
||
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 = _make_dspin(-1e9, 1e9, 4)
|
||
spin.setEnabled(False)
|
||
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)
|
||
|
||
self._batch_fft_rowavg_act = QAction(
|
||
"Batch Compute Row-A&veraged FFT and Store…", self)
|
||
self._batch_fft_rowavg_act.setStatusTip(
|
||
"Select .sras files and compute+store a same-row, distance-weighted "
|
||
"smoothed FFT peak-frequency image for every angle — improves SNR "
|
||
"on noisy regions. DC images and raw waveform data are never "
|
||
"touched; the raw (unsmoothed) FFT is always a recompute away. "
|
||
"Converts v6 files to v7 in place.")
|
||
self._batch_fft_rowavg_act.triggered.connect(self._on_batch_compute_row_avg)
|
||
convert_menu.addAction(self._batch_fft_rowavg_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(
|
||
Jobs.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._apply_alignment_result(None, view_checked=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._apply_alignment_result(
|
||
build_manual_alignment(
|
||
sras, sidecar.ref_angle_idx, sidecar.dc_threshold_mv,
|
||
sidecar.per_angle),
|
||
view_checked=True)
|
||
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)
|
||
|
||
with QSignalBlocker(self.spin_angle):
|
||
self.spin_angle.setRange(0, max(0, sras.n_angles - 1))
|
||
self.spin_angle.setValue(0)
|
||
|
||
# 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.
|
||
with QSignalBlocker(self.combo_channel):
|
||
self.combo_channel.setCurrentIndex(CH4_IDX)
|
||
|
||
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)"
|
||
avg_note = (f", row-averaged n={s.precomputed_row_avg_n}"
|
||
if s.precomputed_row_avg_n else "")
|
||
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 ''}"
|
||
f"{avg_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(Jobs.BATCH)
|
||
self._batch_dc_act.setEnabled(can_batch)
|
||
self._batch_fft_act.setEnabled(can_batch)
|
||
self._batch_fft_rowavg_act.setEnabled(can_batch)
|
||
|
||
self._alignment_act.setEnabled(
|
||
has_file and s.n_angles > 1 and not self._job_running(Jobs.ALIGN))
|
||
self._manual_align_act.setEnabled(
|
||
has_file and s.n_angles > 1 and not self._job_running(Jobs.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._is_fft_mode():
|
||
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._is_fft_mode():
|
||
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.
|
||
with QSignalBlocker(self.btn_draw_roi):
|
||
self.btn_draw_roi.setChecked(active)
|
||
|
||
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 _is_fft_mode(self) -> bool:
|
||
"""Is the selected channel an FFT-derived (CH1/Velocity) mode?"""
|
||
return (self._sras is not None
|
||
and self.combo_channel.currentIndex() in CH1_DERIVED_MODES)
|
||
|
||
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 = _axes_extent(x_axis, y_axis, dx, dy)
|
||
|
||
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)):
|
||
with QSignalBlocker(spin):
|
||
spin.setValue(val)
|
||
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(Jobs.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(
|
||
Jobs.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(
|
||
Jobs.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(Jobs.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(Jobs.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(
|
||
Jobs.BATCH, worker,
|
||
connect=(
|
||
("progress", lambda pct: self._set_progress(Jobs.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._batch_fft_rowavg_act.setEnabled(False)
|
||
self._show_progress(
|
||
Jobs.BATCH, f"Batch computing {label} for {len(paths)} file(s)…",
|
||
maximum=100)
|
||
|
||
def _on_batch_compute_row_avg(self):
|
||
if self._job_running(Jobs.BATCH):
|
||
return
|
||
dlg = RowAverageFftOptionsDialog(
|
||
self, current_n=self._pending_row_avg_n,
|
||
current_threshold_mv=self.spin_threshold_mv.value(),
|
||
pixel_x_mm=self._sras.pixel_x_mm if self._sras is not None else None)
|
||
if dlg.exec() != QDialog.DialogCode.Accepted:
|
||
return
|
||
n, threshold_mv = dlg.get_half_width(), dlg.get_threshold_mv()
|
||
self._pending_row_avg_n = n
|
||
self._settings.setValue("fft/row_avg_n", n)
|
||
|
||
paths, _ = QFileDialog.getOpenFileNames(
|
||
self, "Select .sras files to batch-compute Row-Averaged FFT", "",
|
||
"SRAS files (*.sras);;All files (*)")
|
||
if not paths:
|
||
return
|
||
|
||
self._batch_errors = []
|
||
worker = BatchCacheWorker(paths, "fft_rowavg", self.chk_bg_sub.isChecked(),
|
||
dc_threshold_mv=threshold_mv, row_avg_n=n)
|
||
started = self._run_worker(
|
||
Jobs.BATCH, worker,
|
||
connect=(
|
||
("progress", lambda pct: self._set_progress(Jobs.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 a dialog was open
|
||
|
||
self._batch_dc_act.setEnabled(False)
|
||
self._batch_fft_act.setEnabled(False)
|
||
self._batch_fft_rowavg_act.setEnabled(False)
|
||
self._show_progress(
|
||
Jobs.BATCH,
|
||
f"Batch computing row-averaged FFT (n={n}, threshold={threshold_mv:.1f} mV) "
|
||
f"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(Jobs.BATCH, f"Processed {Path(path).name}…")
|
||
|
||
def _on_batch_finished(self, paths: list[str]):
|
||
self._close_progress(Jobs.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)
|
||
self._batch_fft_rowavg_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(
|
||
Jobs.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
|
||
# No generation bump: this result *is* the current generation's.
|
||
self._apply_alignment_result(result, view_checked=True,
|
||
bump_generation=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 _apply_alignment_result(self, result, *, view_checked: bool,
|
||
bump_generation: bool = True):
|
||
"""Install (or clear, with result=None) the active alignment: reset
|
||
the aligned-image cache and set the Aligned View checkbox without
|
||
firing its change signal."""
|
||
self._alignment_result = result
|
||
self._aligned_cache = {}
|
||
if bump_generation:
|
||
self._alignment_generation += 1
|
||
with QSignalBlocker(self.chk_aligned_view):
|
||
self.chk_aligned_view.setChecked(view_checked)
|
||
self.chk_aligned_view.setEnabled(result is not None)
|
||
|
||
def _manual_alignment_changed(self, result, message: str):
|
||
self._apply_alignment_result(result, view_checked=result is not None)
|
||
self._update_controls_enabled(self._sras is not None)
|
||
self.statusBar().showMessage(message)
|
||
if self._current_image is not None:
|
||
self._refresh_display()
|
||
|
||
def _on_manual_alignment_saved(self, result, sidecar_path_str: str):
|
||
self._manual_alignment_changed(
|
||
result, f"Manual alignment saved to {Path(sidecar_path_str).name}")
|
||
|
||
def _on_manual_alignment_cleared(self):
|
||
self._manual_alignment_changed(None, "Manual alignment cleared.")
|
||
|
||
# ------------------------------------------------------------------
|
||
# 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._is_fft_mode():
|
||
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())
|
||
|