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]
+139
-1
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user