Port batch image export onto the sras_viewer package split

origin/main split sras_viewer.py into the sras_viewer/ package (dialogs,
main_window, common, canvases) after the batch-export feature was built
against the old monolith, so merging origin/main into main lost the feature
to the sras_viewer.py deletion. Re-add it in the new structure:
BatchExportDialog in dialogs.py, a Jobs.EXPORT job key in common.py, and
the menu action/worker wiring in main_window.py, following the Jobs-key
and dialog-class conventions the split already established rather than the
old ad hoc "main" progress-key naming the original commit used.

Also carries forward two test fixes bundled in the original batch-export
commit that the parallel pytest conversion (007089d) had ported from
tools/test_refactor.py without them: the int16-vs-float32 accumulator
expectation in test_sras_average, and asserting the sras_average.py
subprocess calls actually succeed.
This commit is contained in:
Thomas Ales [M S E]
2026-08-10 18:58:56 -05:00
parent 105b9514a5
commit aca501ae5c
4 changed files with 275 additions and 16 deletions
+122 -5
View File
@@ -9,8 +9,8 @@ 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,
QHBoxLayout, QLabel, QMainWindow, QMessageBox, QProgressDialog,
QPushButton, QSizePolicy, QSpinBox, QSplitter, QVBoxLayout, QWidget,
)
import sras_compute as compute
@@ -19,11 +19,12 @@ from sras_compute import (
load_manual_alignment, save_manual_alignment, sidecar_path,
)
from sras_format import (
CH3_IDX, CH4_IDX, CH_NAMES, SrasFile, mv_to_adc,
CH1_IDX, CH3_IDX, CH4_IDX, CH_NAMES, SrasFile, mv_to_adc,
_FALLBACK_YMULT_MV, _FALLBACK_YOFF_ADC,
)
from sras_workers import (
BatchCacheWorker, ComputeWorker, DcPrecomputeWorker, LoadWorker,
BatchCacheWorker, BatchExportWorker, ComputeWorker, DcPrecomputeWorker,
LoadWorker,
)
from .canvases import ImageCanvas, WaveformCanvas
@@ -35,7 +36,8 @@ from .common import (
)
from .align_wizard import AlignmentWizard
from .dialogs import (
FftOptionsDialog, FusedRoiExportDialog, RowAverageFftOptionsDialog,
BatchExportDialog, FftOptionsDialog, FusedRoiExportDialog,
RowAverageFftOptionsDialog,
)
# ---------------------------------------------------------------------------
@@ -85,6 +87,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
@@ -480,6 +486,15 @@ class SrasViewerWindow(QMainWindow):
self._batch_fft_rowavg_act.triggered.connect(self._on_batch_compute_row_avg)
convert_menu.addAction(self._batch_fft_rowavg_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
# ------------------------------------------------------------------
@@ -702,6 +717,8 @@ class SrasViewerWindow(QMainWindow):
self._wizard_act.setEnabled(
has_file and s.n_angles > 1 and self._align_wizard is None)
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(Jobs.EXPORT))
self._update_roi_ui()
def _on_channel_changed(self):
@@ -1538,6 +1555,106 @@ class SrasViewerWindow(QMainWindow):
self._batch_fft_act.setEnabled(True)
self._batch_fft_rowavg_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(Jobs.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(
Jobs.EXPORT, worker,
connect=(
("progress", lambda pct: self._set_progress(Jobs.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(
Jobs.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(Jobs.EXPORT, f"Wrote {Path(path).name}…")
def _on_export_finished(self):
self._close_progress(Jobs.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
# ------------------------------------------------------------------