From aca501ae5c08db8167dad3b612365742760655b9 Mon Sep 17 00:00:00 2001 From: "Thomas Ales [M S E]" Date: Mon, 10 Aug 2026 18:58:56 -0500 Subject: [PATCH] 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. --- sras_viewer/common.py | 1 + sras_viewer/dialogs.py | 144 +++++++++++++++++++++++++++++++++++-- sras_viewer/main_window.py | 127 ++++++++++++++++++++++++++++++-- tests/test_compute.py | 19 +++-- 4 files changed, 275 insertions(+), 16 deletions(-) diff --git a/sras_viewer/common.py b/sras_viewer/common.py index ad587e9..f543b39 100644 --- a/sras_viewer/common.py +++ b/sras_viewer/common.py @@ -60,6 +60,7 @@ class Jobs: COMPUTE = "compute" DC_PRECOMPUTE = "dc_precompute" BATCH = "batch" + EXPORT = "export" # The alignment wizard's three background steps: fetching each angle's CH4 # image for the mask stack, registering the angles, and writing the aligned # export. Separate keys because a retry of one must not be blocked by diff --git a/sras_viewer/dialogs.py b/sras_viewer/dialogs.py index 87379eb..e85bc7e 100644 --- a/sras_viewer/dialogs.py +++ b/sras_viewer/dialogs.py @@ -8,15 +8,20 @@ mask-overlay editor plus the crop and export steps. from pathlib import Path from PyQt6.QtWidgets import ( - QButtonGroup, QCheckBox, QDialog, QDialogButtonBox, QFileDialog, - QGroupBox, QHBoxLayout, QLabel, QLineEdit, QPushButton, QRadioButton, - QScrollArea, QSpinBox, QVBoxLayout, QWidget, + QButtonGroup, QCheckBox, QDialog, QDialogButtonBox, QDoubleSpinBox, + QFileDialog, QGridLayout, QGroupBox, QHBoxLayout, QLabel, QLineEdit, + QMessageBox, QPushButton, QRadioButton, QScrollArea, QSpinBox, + QVBoxLayout, QWidget, ) from sras_compute import PYFFTW_AVAILABLE from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, CH_NAMES +from sras_workers import ExportChannel -from .common import CH_LABELS, VELOCITY_MODE_IDX, _CSS_HINT, _CSS_WARN, _group, _make_dspin, _wrap_label +from .common import ( + CH_LABELS, VELOCITY_MODE_IDX, _CHANNEL_DISPLAY, _CSS_HINT, _CSS_WARN, + _SPIN_MIN_W, _form, _group, _make_dspin, _wrap_label, +) # --------------------------------------------------------------------------- @@ -240,6 +245,137 @@ class RowAverageFftOptionsDialog(QDialog): return self._spin_threshold.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 + + # --------------------------------------------------------------------------- # Export Fused ROI dialog # --------------------------------------------------------------------------- diff --git a/sras_viewer/main_window.py b/sras_viewer/main_window.py index b02d6e7..441199d 100644 --- a/sras_viewer/main_window.py +++ b/sras_viewer/main_window.py @@ -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 # ------------------------------------------------------------------ diff --git a/tests/test_compute.py b/tests/test_compute.py index 3c6e400..f702299 100644 --- a/tests/test_compute.py +++ b/tests/test_compute.py @@ -327,21 +327,26 @@ def test_sras_average(tmp_path): assert np.array_equal(avg.background, SrasFile(str(src)).background), \ "background preserved" src_data = meta["data"] - expect0 = src_data[0][:, :, 0:4, :].astype(np.float32).mean(axis=2).astype(np.int16) + # int16 (not float32) before .mean(): matches average_rows' own + # float64-accumulator behavior for integer input, so this doesn't + # drift from what average_rows actually guarantees. + expect0 = src_data[0][:, :, 0:4, :].astype(np.int16).mean(axis=2).astype(np.int16) assert np.array_equal(np.asarray(avg.data[0])[:, :, 0, :], expect0), \ "first averaged group equals the mean of its 4 source frames" # Remainder handling: 12 frames / 5 -> 2 full groups + 1 partial. dst2 = tmp_path / "legacy_v4_avg5.sras" - subprocess.run([sys.executable, str(REPO / "sras_average.py"), - str(src), str(dst2), "--n", "5"], - capture_output=True, text=True, cwd=REPO) + proc2 = subprocess.run([sys.executable, str(REPO / "sras_average.py"), + str(src), str(dst2), "--n", "5"], + capture_output=True, text=True, cwd=REPO) + assert proc2.returncode == 0, (proc2.stderr or proc2.stdout).strip()[-200:] assert list(SrasFile(str(dst2)).n_frames) == [3, 3], \ "partial trailing group kept by default" dst3 = tmp_path / "legacy_v4_avg5d.sras" - subprocess.run([sys.executable, str(REPO / "sras_average.py"), - str(src), str(dst3), "--n", "5", "--discard-remainder"], - capture_output=True, text=True, cwd=REPO) + proc3 = subprocess.run([sys.executable, str(REPO / "sras_average.py"), + str(src), str(dst3), "--n", "5", "--discard-remainder"], + capture_output=True, text=True, cwd=REPO) + assert proc3.returncode == 0, (proc3.stderr or proc3.stdout).strip()[-200:] assert list(SrasFile(str(dst3)).n_frames) == [2, 2], \ "--discard-remainder drops the partial group"