Add batch export of DC/RF/Velocity map images

Add Export -> Batch Export Images..., which renders and saves one PNG
per angle for whichever of CH1 (RF), CH3 (DC-A), CH4 (DC-B), and
Velocity the user selects, for the currently open file. Each channel
gets its own fixed min/max colorbar range, entered once in the dialog
and held constant across every exported angle, so the resulting images
are directly comparable to each other instead of each auto-scaling to
its own data (today's live-view default).

BatchExportDialog (sras_viewer.py) collects the output folder, file
prefix, and per-channel checkbox + range, seeded from whatever's
already cached (session DC/FFT caches, or the file's own v7 cache
blocks) so opening the dialog never triggers a fresh compute. Export
runs on a background thread via a new BatchExportWorker
(sras_workers.py), which computes each angle's channel image with the
existing dc_image_mv/compute_rf_image helpers (reusing one FFT per
angle for both RF and Velocity) and renders it with a headless
matplotlib Agg canvas.

Also includes several small correctness/robustness fixes that were
already staged in the working tree: an int16-vs-float32 accumulator
mismatch in sras_average.py's row averaging, a DC4-mask cache reuse and
atomic sidecar write in sras_compute.py, a dc3/dc4 pairing guard in
sras_format.py's v7 cache writer, and matching updates to the
tools/ test fixtures.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
Thomas Ales [M S E]
2026-08-10 18:30:00 -05:00
parent b0bfe6e8c6
commit 6976cbf767
7 changed files with 507 additions and 43 deletions
+300 -15
View File
@@ -29,9 +29,10 @@ from PyQt6.QtCore import QObject, 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,
QDoubleSpinBox, QFileDialog, QFormLayout, QFrame, QGridLayout, QGroupBox,
QHBoxLayout, QLabel, QLineEdit, QMainWindow, QMessageBox, QProgressDialog,
QPushButton, QRadioButton, QScrollArea, QSizePolicy, QSpinBox, QSplitter,
QVBoxLayout, QWidget,
)
import sras_compute as compute
@@ -45,8 +46,9 @@ from sras_format import (
_FALLBACK_YMULT_MV, _FALLBACK_YOFF_ADC,
)
from sras_workers import (
AngleAlignmentWorker, BatchCacheWorker, Ch4MaskWorker, ComputeWorker,
CrossCorrelateWorker, DcPrecomputeWorker, LoadWorker,
AngleAlignmentWorker, BatchCacheWorker, BatchExportWorker, Ch4MaskWorker,
ComputeWorker, CrossCorrelateWorker, DcPrecomputeWorker, ExportChannel,
LoadWorker,
)
faulthandler.enable() # print a native stack trace on SIGSEGV/SIGABRT/etc.
@@ -743,6 +745,137 @@ class FftOptionsDialog(QDialog):
return max(1, self._spin_pad.value())
# ---------------------------------------------------------------------------
# Batch export dialog (Export -> Batch Export Images...)
# ---------------------------------------------------------------------------
class BatchExportDialog(QDialog):
"""Configure a batch PNG export of DC/RF/Velocity maps across every angle
of the currently open file.
Each channel's colorbar range is entered here and held fixed for every
exported angle (rather than auto-scaled per image, today's live-view
default) so the exported images are directly comparable to each other.
"""
# (ch_idx, is_velocity, label, colorbar unit, filename tag)
_ROWS = [
(CH1_IDX, False, CH_LABELS[CH1_IDX], _CHANNEL_DISPLAY[CH1_IDX][2],
CH_NAMES[CH1_IDX]),
(CH3_IDX, False, CH_LABELS[CH3_IDX], _CHANNEL_DISPLAY[CH3_IDX][2],
CH_NAMES[CH3_IDX]),
(CH4_IDX, False, CH_LABELS[CH4_IDX], _CHANNEL_DISPLAY[CH4_IDX][2],
CH_NAMES[CH4_IDX]),
(CH1_IDX, True, CH_LABELS[VELOCITY_MODE_IDX],
_CHANNEL_DISPLAY[VELOCITY_MODE_IDX][2], CH_NAMES[VELOCITY_MODE_IDX]),
]
# DC is cheap and precomputed on load; CH1/Velocity need a per-pixel FFT
# that can take minutes, so don't default to exporting them (same
# rationale as SrasViewerWindow._on_load_done's channel default).
_DEFAULT_CHECKED = {CH3_IDX, CH4_IDX}
def __init__(self, parent, *, default_dir: str, default_prefix: str,
default_ranges: dict[tuple[int, bool], tuple[float, float]]):
super().__init__(parent)
self.setWindowTitle("Batch Export Images")
self.setModal(True)
self.setMinimumWidth(460)
layout = QVBoxLayout(self)
# ---- Output ------------------------------------------------------
grp_out, ol = _group("Output")
dir_row = QHBoxLayout()
self._edit_dir = QLineEdit(default_dir)
btn_browse = QPushButton("Browse…")
btn_browse.clicked.connect(self._on_browse)
dir_row.addWidget(self._edit_dir)
dir_row.addWidget(btn_browse)
out_form = _form()
out_form.addRow("Folder:", dir_row)
self._edit_prefix = QLineEdit(default_prefix)
out_form.addRow("File prefix:", self._edit_prefix)
ol.addLayout(out_form)
layout.addWidget(grp_out)
# ---- Channels ------------------------------------------------------
grp_ch, cl = _group("Channels (fixed range, applied to every angle)")
grid = QGridLayout()
grid.setHorizontalSpacing(8)
grid.setVerticalSpacing(6)
grid.addWidget(_wrap_label("min:", _CSS_HINT), 0, 1)
grid.addWidget(_wrap_label("max:", _CSS_HINT), 0, 2)
self._rows: list[tuple[QCheckBox, QDoubleSpinBox, QDoubleSpinBox]] = []
for i, (ch_idx, is_velocity, label, unit, _tag) in enumerate(self._ROWS, 1):
chk = QCheckBox(label + (f" [{unit}]" if unit else ""))
chk.setChecked(not is_velocity and ch_idx in self._DEFAULT_CHECKED)
vmin, vmax = default_ranges.get((ch_idx, is_velocity), (0.0, 1.0))
spin_min, spin_max = QDoubleSpinBox(), QDoubleSpinBox()
for spin, val in ((spin_min, vmin), (spin_max, vmax)):
spin.setRange(-1e9, 1e9)
spin.setDecimals(4)
spin.setMinimumWidth(_SPIN_MIN_W)
spin.setValue(val)
grid.addWidget(chk, i, 0)
grid.addWidget(spin_min, i, 1)
grid.addWidget(spin_max, i, 2)
self._rows.append((chk, spin_min, spin_max))
cl.addLayout(grid)
layout.addWidget(grp_ch)
# ---- Buttons --------------------------------------------------
buttons = QDialogButtonBox()
buttons.addButton("Export", QDialogButtonBox.ButtonRole.AcceptRole
).clicked.connect(self.accept)
buttons.addButton("Cancel", QDialogButtonBox.ButtonRole.RejectRole
).clicked.connect(self.reject)
layout.addWidget(buttons)
def _on_browse(self):
d = QFileDialog.getExistingDirectory(
self, "Select Output Folder", self._edit_dir.text())
if d:
self._edit_dir.setText(d)
def accept(self):
"""Validate before closing — Cancel bypasses this entirely."""
if not self.get_prefix():
QMessageBox.warning(self, "Batch Export", "Enter a file prefix.")
return
if not self.get_output_dir():
QMessageBox.warning(self, "Batch Export", "Choose an output folder.")
return
selected = self.get_selected_channels()
if not selected:
QMessageBox.warning(
self, "Batch Export", "Select at least one channel to export.")
return
for ch in selected:
if not ch.vmin < ch.vmax:
QMessageBox.warning(
self, "Batch Export", f"{ch.label}: min must be less than max.")
return
super().accept()
def get_output_dir(self) -> str:
return self._edit_dir.text().strip()
def get_prefix(self) -> str:
return self._edit_prefix.text().strip()
def get_selected_channels(self) -> list[ExportChannel]:
result = []
for (chk, spin_min, spin_max), (ch_idx, is_velocity, label, unit, tag) in zip(
self._rows, self._ROWS):
if chk.isChecked():
result.append(ExportChannel(
ch_idx=ch_idx, is_velocity=is_velocity,
vmin=spin_min.value(), vmax=spin_max.value(),
label=label, unit=unit, tag=tag))
return result
# ---------------------------------------------------------------------------
# Manual alignment dialog (Fusion -> Manual Alignment...)
# ---------------------------------------------------------------------------
@@ -1399,6 +1532,10 @@ class SrasViewerWindow(QMainWindow):
# Convert menu: batch DC/FFT compute-and-store (v6 -> v7)
self._batch_errors: list[str] = []
# Export menu: batch image export
self._export_errors: list[str] = []
self._export_ok_count: int = 0
# 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
@@ -1412,6 +1549,8 @@ class SrasViewerWindow(QMainWindow):
self._fft_cache: dict[tuple[int, bool, int | None, float], np.ndarray] = {}
self._dc_precompute_worker: DcPrecomputeWorker | None = None
self._dc_generation: int = 0
self._compute_generation: int = 0
self._pending_compute_generation: int | None = None
# Angle alignment ("Fusion" menu)
self._alignment_result = None
@@ -1447,6 +1586,19 @@ class SrasViewerWindow(QMainWindow):
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.
3. The thread must quit once worker.run() returns, regardless of
which signal (if any) run() emitted on its way out. quit_on exists
so a worker whose real completion is signalled asynchronously
from *within* run() (e.g. angle_done fired from pool threads
before run() itself returns) can still opt into quitting earlier,
but relying on quit_on alone is fragile: a worker whose run()
emits `error` instead of `finished` on an exception path, at a
call site that forgot to add "error" to quit_on, would otherwise
never quit — silently wedging that job in self._jobs forever.
Quitting unconditionally after run() returns closes that gap for
every worker, present and future, without requiring every call
site to enumerate every signal run() might emit.
"""
if key in self._jobs:
return False
@@ -1454,7 +1606,11 @@ class SrasViewerWindow(QMainWindow):
thread = QThread()
self._jobs[key] = (thread, worker, on_done) # claim before anything pumps
worker.moveToThread(thread)
thread.started.connect(worker.run)
def _run_then_quit():
worker.run()
thread.quit()
thread.started.connect(_run_then_quit)
for signal_name, slot in connect:
getattr(worker, signal_name).connect(slot)
for signal_name in quit_on:
@@ -1792,6 +1948,15 @@ class SrasViewerWindow(QMainWindow):
self._batch_fft_act.triggered.connect(lambda: self._on_batch_compute("fft"))
convert_menu.addAction(self._batch_fft_act)
export_menu = menubar.addMenu("&Export")
self._batch_export_act = QAction("&Batch Export Images…", self)
self._batch_export_act.setStatusTip(
"Export DC/RF/Velocity map PNGs for every angle of the open "
"file, with a fixed colorbar range per channel.")
self._batch_export_act.setEnabled(False)
self._batch_export_act.triggered.connect(self._on_batch_export)
export_menu.addAction(self._batch_export_act)
# ------------------------------------------------------------------
# Drag-and-drop
# ------------------------------------------------------------------
@@ -1826,10 +1991,10 @@ class SrasViewerWindow(QMainWindow):
return
self.btn_open.setEnabled(False)
self.statusBar().showMessage(f"Loading {Path(path).name}…")
self._show_progress("main", f"Loading {Path(path).name}…")
self._show_progress("load", f"Loading {Path(path).name}…")
def _on_load_done(self, sras):
self._close_progress("main")
self._close_progress("load")
self.btn_open.setEnabled(True)
if sras is None:
return
@@ -1850,6 +2015,7 @@ class SrasViewerWindow(QMainWindow):
self._dc_cache = {}
self._fft_cache = {}
self._dc_generation += 1
self._compute_generation += 1
self.lbl_dc_precompute.setText("")
self._alignment_result = None
@@ -1988,6 +2154,8 @@ class SrasViewerWindow(QMainWindow):
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._batch_export_act.setEnabled(has_file and not self._job_running("export"))
self._update_roi_ui()
def _on_channel_changed(self):
@@ -2322,6 +2490,11 @@ class SrasViewerWindow(QMainWindow):
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
# Snapshot which file this compute belongs to, so a result that lands
# after a *different* file has since been loaded (nothing blocks
# opening a new file while a compute is in flight) gets discarded
# instead of being cached/displayed against the wrong SrasFile.
self._pending_compute_generation = self._compute_generation
worker = ComputeWorker(
self._sras, angle_idx, ch_idx,
@@ -2347,13 +2520,13 @@ class SrasViewerWindow(QMainWindow):
if is_fft:
self.statusBar().showMessage("Computing FFT…")
self._show_progress(
"main",
"compute",
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}…")
self._show_progress("compute", f"Computing DC image for angle {angle_idx}…")
def _after_compute(self):
"""If settings changed while the compute was running, re-dispatch
@@ -2367,7 +2540,10 @@ class SrasViewerWindow(QMainWindow):
self._refresh_display()
def _on_compute_done(self, result):
self._close_progress("main")
self._close_progress("compute")
generation, self._pending_compute_generation = self._pending_compute_generation, None
if generation != self._compute_generation:
return # a different file was loaded while this was computing — discard
if result is None:
return # cancelled mid-compute; the partial image must not cache
angle_idx = self._pending_angle
@@ -2554,6 +2730,106 @@ class SrasViewerWindow(QMainWindow):
self._batch_dc_act.setEnabled(True)
self._batch_fft_act.setEnabled(True)
# ------------------------------------------------------------------
# Export menu: batch image export
# ------------------------------------------------------------------
def _export_default_ranges(self) -> dict[tuple[int, bool], tuple[float, float]]:
"""Best-effort default min/max per exportable channel, from whatever
is already cached — this must never trigger a fresh compute just to
seed the dialog (CH1/Velocity in particular can be expensive)."""
s = self._sras
ranges: dict[tuple[int, bool], tuple[float, float]] = {}
for ch_idx, precomputed in ((CH3_IDX, s.precomputed_dc3_mv),
(CH4_IDX, s.precomputed_dc4_mv)):
arrays = [self._dc_cache.get((a, ch_idx)) for a in range(s.n_angles)]
arrays = [a for a in arrays if a is not None]
if not arrays:
arrays = [a for a in precomputed if a is not None]
if arrays:
ranges[(ch_idx, False)] = (
float(min(a.min() for a in arrays)),
float(max(a.max() for a in arrays)))
freq_arrays = list(self._fft_cache.values())
if not freq_arrays:
freq_arrays = [a for a in s.precomputed_freq_mhz if a is not None]
if freq_arrays:
fmin = float(min(a.min() for a in freq_arrays))
fmax = float(max(a.max() for a in freq_arrays))
ranges[(CH1_IDX, False)] = (fmin, fmax)
grating = self.spin_grating_um.value()
ranges[(CH1_IDX, True)] = (fmin * grating, fmax * grating)
return ranges
def _on_batch_export(self):
if self._sras is None or self._job_running("export"):
return
s = self._sras
dlg = BatchExportDialog(
self, default_dir=str(s.path.parent), default_prefix=s.path.stem,
default_ranges=self._export_default_ranges())
if dlg.exec() != QDialog.DialogCode.Accepted:
return
output_dir = dlg.get_output_dir()
try:
Path(output_dir).mkdir(parents=True, exist_ok=True)
except OSError as exc:
QMessageBox.warning(
self, "Batch Export", f"Could not create output folder:\n{exc}")
return
channels = dlg.get_selected_channels()
self._export_errors = []
self._export_ok_count = 0
worker = BatchExportWorker(
s, channels, output_dir, dlg.get_prefix(),
cmap=self.combo_cmap.currentText(),
apply_bg_sub=self.chk_bg_sub.isChecked(),
dc_threshold_mv=self.spin_threshold_mv.value(),
n_fft=self._current_n_fft(), grating_um=self.spin_grating_um.value())
started = self._run_worker(
"export", worker,
connect=(
("progress", lambda pct: self._set_progress("export", pct)),
("file_done", self._on_export_file_done),
("finished", self._on_export_finished),
),
on_done=self._after_export,
)
if not started:
return # a second trigger snuck in while the dialog was open
self._batch_export_act.setEnabled(False)
self._show_progress(
"export", f"Exporting images to {output_dir}…", maximum=100)
def _on_export_file_done(self, path: str, err: str):
if err:
self._export_errors.append(f"{Path(path).name if path else '?'} — {err}")
else:
self._export_ok_count += 1
self._show_progress("export", f"Wrote {Path(path).name}…")
def _on_export_finished(self):
self._close_progress("export")
n_ok = self._export_ok_count
n_failed = len(self._export_errors)
if n_failed:
summary = (f"Batch export: {n_ok} image(s) written, {n_failed} "
f"failed: {'; '.join(self._export_errors)}")
else:
summary = f"Batch export: {n_ok} image(s) written."
self.statusBar().showMessage(summary)
self._export_errors = []
def _after_export(self):
self._update_controls_enabled(self._sras is not None)
# ------------------------------------------------------------------
# Fusion: angle alignment
# ------------------------------------------------------------------
@@ -2568,7 +2844,7 @@ class SrasViewerWindow(QMainWindow):
started = self._run_worker(
"align", AngleAlignmentWorker(self._sras, ref_idx, threshold_mv),
connect=(
("progress", lambda pct: self._set_progress("main", pct)),
("progress", lambda pct: self._set_progress("align", pct)),
("finished", lambda result, err, g=generation:
self._on_alignment_done(g, result, err)),
),
@@ -2579,13 +2855,13 @@ class SrasViewerWindow(QMainWindow):
self._alignment_act.setEnabled(False)
self._show_progress(
"main",
"align",
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")
self._close_progress("align")
if generation != self._alignment_generation:
return # a new file was loaded while this was computing — discard
if error_msg:
@@ -2639,7 +2915,16 @@ class SrasViewerWindow(QMainWindow):
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)
# Deliberately NOT WA_DeleteOnClose: Ch4MaskWorker/CrossCorrelateWorker
# have no stop() and are owned by this window's _jobs (not the
# dialog), so a background mask-fetch/cross-correlate job can still be
# running after the dialog closes. Its signals stay connected to this
# dialog's own bound methods (_on_mask_angle_done, _finish_mask_prep,
# etc.) until it finishes — if Qt had already destroyed the dialog's
# C++ object, those late callbacks would raise "wrapped C/C++ object
# ... has been deleted". Leaving the object merely hidden (not
# destroyed) lets them land harmlessly; it is garbage-collected once
# the job's own references to it are released.
self._manual_align_dialog = dlg
dlg.show()