diff --git a/sras_average.py b/sras_average.py index d997a15..0c8fc77 100644 --- a/sras_average.py +++ b/sras_average.py @@ -54,8 +54,13 @@ def average_rows(block: np.ndarray, n: int, discard_remainder: bool) -> np.ndarr """Average every N frames of one angle's (n_rows, n_ch, n_frames, spf) block. Returns int16 of shape (n_rows, n_ch, n_out, spf). - Averaging is done in float32 and rounded on cast, matching numpy's mean - followed by an int16 cast in the original implementation. + Cast to native int16 (not float32) before calling .mean(): numpy's mean() + uses a float64 accumulator by default for integer input, matching the + original implementation exactly (which kept the whole file as int16 and + called .mean() directly). A float32 cast here would use a float32 + accumulator instead — for large group sizes that can round the sum + differently than float64 and, after the int16 cast below, occasionally + land on a value 1 ADC count away from the original tool's output. """ n_frames = block.shape[2] n_full = n_frames // n @@ -63,11 +68,11 @@ def average_rows(block: np.ndarray, n: int, discard_remainder: bool) -> np.ndarr parts = [] if n_full: - full = block[:, :, :n_full * n, :].astype(np.float32) + full = block[:, :, :n_full * n, :].astype(np.int16) full = full.reshape(block.shape[0], block.shape[1], n_full, n, block.shape[3]) parts.append(full.mean(axis=3).astype(np.int16)) if remainder and not discard_remainder: - tail = block[:, :, n_full * n:, :].astype(np.float32) + tail = block[:, :, n_full * n:, :].astype(np.int16) parts.append(tail.mean(axis=2, keepdims=True).astype(np.int16)) if not parts: diff --git a/sras_compute.py b/sras_compute.py index ee34166..a0966d8 100644 --- a/sras_compute.py +++ b/sras_compute.py @@ -266,13 +266,19 @@ def compute_rf_image(sras: SrasFile, angle_idx: int, budget=budget) background = sras.background if (apply_bg_sub and sras.background is not None) else None cal4 = sras.cal(CH4_IDX) + # DC4 mask source, in the same priority order as the fast path above: + # already-cached DC block (skips reading CH4 raw waveforms entirely), + # then the caller-supplied image, else recompute per chunk below. + dc4_full = sras.cached_dc_mv(angle_idx, CH4_IDX) if dc_threshold_mv is not None else None + if dc4_full is None: + dc4_full = dc4_mv def chunk(r0: int, r1: int): if dc_threshold_mv is None: valid = None else: - if dc4_mv is not None: - dc4_chunk = dc4_mv[r0:r1] + if dc4_full is not None: + dc4_chunk = dc4_full[r0:r1] else: dc4_chunk = adc_to_mv( data[r0:r1, CH4_IDX, :, :].astype(np.float32).mean(axis=-1), *cal4) @@ -1007,7 +1013,13 @@ def save_manual_alignment(sras: SrasFile, ref_angle_idx: int, for a, p in per_angle.items() }, } - path.write_text(json.dumps(payload, indent=2)) + # Write to a temp file in the same directory then rename over the target + # — os.replace() is atomic on both POSIX and Windows, so a crash mid-write + # leaves either the old sidecar or the new one, never a truncated/corrupt + # file that load_manual_alignment would silently treat as "no alignment". + tmp_path = path.with_name(path.name + f".tmp{os.getpid()}") + tmp_path.write_text(json.dumps(payload, indent=2)) + os.replace(tmp_path, path) return path diff --git a/sras_format.py b/sras_format.py index 0b2715d..5c82f6a 100644 --- a/sras_format.py +++ b/sras_format.py @@ -523,7 +523,13 @@ class SrasFile: final_freq = new_freq_mhz if new_freq_mhz is not None else self.precomputed_freq_mhz final_bg_sub = new_bg_sub if new_bg_sub is not None else self.precomputed_bg_sub - dc_entries = [a for a in range(self.n_angles) if final_dc3[a] is not None] + # dc3/dc4 are always populated together by every current caller, but + # guard the per-angle pairing explicitly rather than assume it: an + # angle present in only one of the two arrays would otherwise crash + # below on final_dc4[a].astype(...) (or silently store the wrong + # dc3/dc4 pairing for that angle). + dc_entries = [a for a in range(self.n_angles) + if final_dc3[a] is not None and final_dc4[a] is not None] fft_entries = [a for a in range(self.n_angles) if final_freq[a] is not None] block_flags = ((CACH_FLAG_DC if dc_entries else 0) @@ -552,13 +558,23 @@ class SrasFile: f.truncate() f.flush() os.fsync(f.fileno()) - # Version-byte flip last: if the process dies before this point, - # the file is still readable as plain v6 (v6 parsing only - # bounds-checks per-angle offset+nbytes <= file_size, it never - # asserts exactly how many bytes follow the last angle) — so an - # interrupted write can never corrupt the file, only leave - # harmless trailing bytes that the next successful write - # overwrites via this same deterministic cache offset. + # Version-byte flip last: when converting a v6 source (self.version + # was 6 on entry), if the process dies before this point the file + # is still readable as plain v6 (v6 parsing only bounds-checks + # per-angle offset+nbytes <= file_size, it never asserts exactly + # how many bytes follow the last angle) — so an interrupted write + # can never corrupt the file, only leave harmless trailing bytes + # that the next successful write overwrites via this same + # deterministic cache offset. + # + # That guarantee does NOT extend to updating an already-v7 file: + # the version byte here is already 7 before this call, so a crash + # during the payload write above (before flush/fsync/truncate) + # can leave a cache tail that mixes a prefix of the new payload + # with a stale suffix of the old one, and this method has no + # protection against that case (no atomic rename — the tail is + # rewritten in place to avoid copying the, potentially huge, + # waveform data that precedes it). f.seek(4) f.write(struct.pack("B", 7)) f.flush() diff --git a/sras_viewer.py b/sras_viewer.py index fb7e8b5..a6b53c2 100644 --- a/sras_viewer.py +++ b/sras_viewer.py @@ -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() diff --git a/sras_workers.py b/sras_workers.py index 5696b0e..5fd0787 100644 --- a/sras_workers.py +++ b/sras_workers.py @@ -10,15 +10,19 @@ everything they need through their constructor and hand results back by signal. import os from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed from concurrent.futures.process import BrokenProcessPool +from dataclasses import dataclass +from pathlib import Path import numpy as np +from matplotlib.backends.backend_agg import FigureCanvasAgg +from matplotlib.figure import Figure from PyQt6.QtCore import QObject, pyqtSignal import sras_compute as compute from sras_compute import ( cache_file, compute_angle_alignment, compute_rf_image, dc_image_mv, ) -from sras_format import CH3_IDX, CH4_IDX, SrasFile +from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, SrasFile # Concurrency caps. Batch conversion runs one process per file, and each of # those processes threads internally, so the two must be divided rather than @@ -281,6 +285,140 @@ class BatchCacheWorker(QObject): self.finished.emit() +@dataclass +class ExportChannel: + """One row of a batch image export: which raw channel to read, whether + to apply the Velocity post-multiply, and the fixed display range/labels + to render it with. + + Kept free of sras_viewer's display constants (CH_LABELS, VELOCITY_MODE_IDX, + etc.) so BatchExportWorker has no dependency on the GUI module — the + caller resolves labels/units/filename tags once, up front. + """ + ch_idx: int # CH1_IDX, CH3_IDX, or CH4_IDX -- which raw data to read + is_velocity: bool # True only for the derived Velocity map (post-multiply of CH1 freq) + vmin: float + vmax: float + label: str # e.g. "CH3 -- Bias A (DC mean)" + unit: str # colorbar units, e.g. "mV" + tag: str # filename tag: "CH1", "CH3", "CH4", "VEL" + + +def _render_map_png(img: np.ndarray, extent: list[float], *, cmap: str, + vmin: float, vmax: float, title: str, colorbar_label: str, + out_path: str): + """Render one map image to *out_path* with a fixed vmin/vmax, using a + headless Agg canvas so this never touches the GUI thread's interactive + matplotlib backend. Layout mirrors ImageCanvas.show_image.""" + fig = Figure(figsize=(7, 5), tight_layout=True) + FigureCanvasAgg(fig) + ax = fig.add_subplot(111) + im = ax.imshow(img, aspect="auto", origin="upper", extent=extent, + cmap=cmap, vmin=vmin, vmax=vmax, interpolation="nearest") + cb = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + if colorbar_label: + cb.set_label(colorbar_label) + ax.set_xlabel("X (mm)") + ax.set_ylabel("Y (mm)") + ax.set_title(title) + fig.savefig(out_path, dpi=150) + + +class BatchExportWorker(CancellableWorker): + """Renders and saves one PNG per (angle, selected channel) for an + already-open SrasFile, with each channel's vmin/vmax held fixed across + every angle so the colorbar is directly comparable image to image. + + Takes the SrasFile object directly (like ComputeWorker/DcPrecomputeWorker) + rather than a path -- this runs against the file already open in the GUI, + not an arbitrary batch of files, so there is no need to reopen it in a + subprocess the way BatchCacheWorker does. + + Emits progress(int) (0-100 over all angle x channel pairs), file_done(str, + str) (output path, error message or ""), and finished() -- same shape as + BatchCacheWorker. + """ + progress = pyqtSignal(int) + file_done = pyqtSignal(str, str) + finished = pyqtSignal() + + def __init__(self, sras: SrasFile, channels: list[ExportChannel], + output_dir: str, prefix: str, *, cmap: str, + apply_bg_sub: bool, dc_threshold_mv: float, + n_fft: int | None, grating_um: float): + super().__init__() + self._sras = sras + self._channels = channels + self._output_dir = Path(output_dir) + self._prefix = prefix + self._cmap = cmap + self._apply_bg_sub = apply_bg_sub + self._dc_threshold_mv = dc_threshold_mv + self._n_fft = n_fft + self._grating_um = grating_um + + def _angle_extent(self, angle_idx: int) -> list[float]: + s = self._sras + 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 + return [x_axis[0] - dx / 2, x_axis[-1] + dx / 2, + y_axis[-1] + dy / 2, y_axis[0] - dy / 2] + + def run(self): + try: + s = self._sras + n_angles = s.n_angles + total = n_angles * len(self._channels) + done = 0 + needs_freq = any(c.ch_idx == CH1_IDX for c in self._channels) + + for angle_idx in range(n_angles): + if self._stop: + break + extent = self._angle_extent(angle_idx) + angle_deg = s.angles_deg[angle_idx] + + freq_mhz = None + if needs_freq: + freq_mhz = compute_rf_image( + s, angle_idx, dc_threshold_mv=self._dc_threshold_mv, + apply_bg_sub=self._apply_bg_sub, n_fft=self._n_fft, + should_stop=self._stopped) + + for channel in self._channels: + if self._stop: + break + out_path = (self._output_dir + / f"{self._prefix}_angle{angle_idx:02d}_{channel.tag}.png") + try: + if channel.ch_idx == CH1_IDX: + img = (freq_mhz * self._grating_um if channel.is_velocity + else freq_mhz) + else: + img = dc_image_mv(s, angle_idx, channel.ch_idx, + should_stop=self._stopped) + title = f"{channel.tag} | {angle_deg:.1f}°" + colorbar_label = (f"{channel.label} ({channel.unit})" + if channel.unit else channel.label) + _render_map_png( + img, extent, cmap=self._cmap, + vmin=channel.vmin, vmax=channel.vmax, + title=title, colorbar_label=colorbar_label, + out_path=str(out_path)) + self.file_done.emit(str(out_path), "") + except Exception as exc: + self.file_done.emit(str(out_path), str(exc)) + done += 1 + self.progress.emit(int(done / max(1, total) * 100)) + + self.finished.emit() + except Exception as exc: + self.file_done.emit("", str(exc)) + self.finished.emit() + + class AngleAlignmentWorker(QObject): """Computes rotation+translation alignment for every angle in *sras*, referenced to *ref_angle_idx*, from each angle's binarized CH4 mask. diff --git a/tools/make_test_sras.py b/tools/make_test_sras.py index 07386c3..a9bf3e2 100644 --- a/tools/make_test_sras.py +++ b/tools/make_test_sras.py @@ -11,12 +11,16 @@ Usage: import argparse import struct +import sys from pathlib import Path import numpy as np -HDR_FMT_V6 = ">4sBHfffffffIdBB" -GEO_FMT_V6 = ">ffIH" +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +# Packed straight from sras_format.py's own constants (not a hand-copy) so +# this generator can never silently drift from what the real parser expects. +from sras_format import GEO_FMT_V6, HDR_FMT as HDR_FMT_LEGACY, HDR_FMT_V6 # noqa: E402 # Per-angle (n_rows, n_frames) — deliberately different per angle so ragged # geometry handling is actually exercised. @@ -125,9 +129,6 @@ def write(path: Path, n_angles: int = 3, seed: int = 0, return meta -HDR_FMT_LEGACY = ">4sBHHffffIIdBB" - - def write_legacy(path: Path, version: int = 4, n_angles: int = 2, n_rows: int = 4, n_frames: int = 10, samples_per_frame: int = 32, seed: int = 0) -> dict: diff --git a/tools/test_refactor.py b/tools/test_refactor.py index 0cc50f6..cbcb93e 100644 --- a/tools/test_refactor.py +++ b/tools/test_refactor.py @@ -288,23 +288,30 @@ def test_legacy_parse_and_average(scratch: Path): check("background preserved", np.array_equal(avg.background, SrasFile(str(src)).background)) 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) check("first averaged group equals the mean of its 4 source frames", np.array_equal(np.asarray(avg.data[0])[:, :, 0, :], expect0)) # Remainder handling: 12 frames / 5 -> 2 full groups + 1 partial. dst2 = scratch / "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) + check("sras_average ran (--n 5)", proc2.returncode == 0, + (proc2.stderr or proc2.stdout).strip()[-200:]) if dst2.exists(): check("partial trailing group kept by default", list(SrasFile(str(dst2)).n_frames) == [3, 3], f"{list(SrasFile(str(dst2)).n_frames)}") dst3 = scratch / "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) + check("sras_average ran (--n 5 --discard-remainder)", proc3.returncode == 0, + (proc3.stderr or proc3.stdout).strip()[-200:]) if dst3.exists(): check("--discard-remainder drops the partial group", list(SrasFile(str(dst3)).n_frames) == [2, 2],