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]
@@ -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
|
||||
|
||||
+140
-4
@@ -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
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
+122
-5
@@ -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
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@@ -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"),
|
||||
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"),
|
||||
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"
|
||||
|
||||
|
||||
Reference in New Issue
Block a user