Thomas Ales
bbb075ff34
Manual mode previously offered only Auto De-rotate (rotation from known scan angles) plus keyboard nudging, so every angle's translation had to be found entirely by eye — fine for a small correction, unworkable for the tens-of-mm scatter a real many-angle scan can have between angles before any translation search runs at all. Add an "Auto Cross-Correlate (vs Reference)" action that, for every non-reference angle, sets rotation to the known analytic angle (same as Auto De-rotate) and translation to the FFT-phase-correlation best fit against angle #0 (always the reference/ground truth). This is meant to land every angle's outline roughly stacked on the others so keyboard nudging only has to make small corrections afterward, per the intended workflow: cross-correlate first, then nudge. Exposes two tunable options, since real data may correlate better one way than the other: "Correlate on" (raw CH4 signal, minus its own minimum -- the default, robust to a shared threshold not suiting every angle's real signal level -- or the thresholded binary mask), and "Search margin" (how far a shift the phase correlation searches before wraparound bias becomes a risk). Runs on a background thread (new CrossCorrelateWorker in sras_workers.py) since correlating many high-resolution angles can take long enough to visibly freeze the dialog otherwise. The underlying per-angle correlation math is factored out of compute_angle_alignment's Step 2 into a new, reusable correlate_translation_mm, so the existing automatic Fusion -> Angle Alignment action and the new manual button now share one implementation instead of two copies of the same phase-correlation logic. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
403 lines
16 KiB
Python
403 lines
16 KiB
Python
#!/usr/bin/env python3
|
||
"""Background workers for the SRAS viewer.
|
||
|
||
Every worker is a plain QObject moved onto its own QThread by
|
||
SrasViewerWindow._run_worker, exposing signals only. Workers must never touch
|
||
GUI-thread-owned state (the display caches in particular) — they take
|
||
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
|
||
|
||
import numpy as np
|
||
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
|
||
|
||
# Concurrency caps. Batch conversion runs one process per file, and each of
|
||
# those processes threads internally, so the two must be divided rather than
|
||
# both set to the core count. Files also commonly sit on one external drive,
|
||
# where a dozen concurrent readers is slower than a few — hence the low
|
||
# default, overridable from the environment.
|
||
_BATCH_MAX_PROCS = int(os.environ.get("SRAS_BATCH_PROCS", 0)) or min(
|
||
4, os.cpu_count() or 2)
|
||
|
||
# Spawning a pool costs roughly a second of interpreter startup (each child
|
||
# re-imports the entry module). That is noise against a multi-GB scan but
|
||
# dominates a batch of small files, where it would make the job *slower* —
|
||
# so below this total size the batch just runs in the worker thread.
|
||
_BATCH_POOL_MIN_BYTES = int(os.environ.get("SRAS_BATCH_POOL_MIN_MB", 512)) * 1024 * 1024
|
||
|
||
|
||
class CancellableWorker(QObject):
|
||
"""A worker whose compute polls stop() between row chunks.
|
||
|
||
Without this a shutdown has to wait out whatever is in flight, and on a
|
||
large scan a single angle is ~40 s — far too long to block closing the
|
||
window. Chunk-level polling bounds the wait to one chunk instead.
|
||
"""
|
||
|
||
def __init__(self):
|
||
super().__init__()
|
||
self._stop = False
|
||
|
||
def stop(self):
|
||
self._stop = True
|
||
|
||
def _stopped(self) -> bool:
|
||
return self._stop
|
||
|
||
|
||
class LoadWorker(QObject):
|
||
finished = pyqtSignal(object) # SrasFile | None
|
||
error = pyqtSignal(str)
|
||
|
||
def __init__(self, path: str):
|
||
super().__init__()
|
||
self._path = path
|
||
|
||
def run(self):
|
||
try:
|
||
self.finished.emit(SrasFile(self._path))
|
||
except Exception as exc:
|
||
self.error.emit(str(exc))
|
||
self.finished.emit(None)
|
||
|
||
|
||
class ComputeWorker(CancellableWorker):
|
||
"""Computes one displayable image for (angle, channel).
|
||
|
||
For CH1/Velocity (FFT-derived) channels, the FFT is only run for pixels
|
||
whose DC4 (Bias B) mean is at or above dc_threshold_mv — masked pixels are
|
||
left at 0 MHz without ever being FFT'd, since that's the expensive part of
|
||
a scan. If the DC4 image for this angle is already known, pass it in as
|
||
*dc4_mv* to skip re-reading the CH4 channel from disk entirely.
|
||
|
||
Emits a plain ``np.ndarray`` already in display units.
|
||
"""
|
||
finished = pyqtSignal(object)
|
||
error = pyqtSignal(str)
|
||
|
||
def __init__(self, sras: SrasFile, angle_idx: int, ch_idx: int,
|
||
apply_bg_sub: bool = True, n_fft: int | None = None,
|
||
dc_threshold_mv: float = 0.0,
|
||
dc4_mv: np.ndarray | None = None,
|
||
is_fft_mode: bool = False):
|
||
super().__init__()
|
||
self._sras = sras
|
||
self._angle = angle_idx
|
||
self._ch = ch_idx
|
||
self._apply_bg_sub = apply_bg_sub
|
||
self._n_fft = n_fft
|
||
self._dc_threshold = dc_threshold_mv
|
||
self._dc4_mv = dc4_mv
|
||
self._is_fft_mode = is_fft_mode
|
||
|
||
def run(self):
|
||
try:
|
||
if self._is_fft_mode:
|
||
img = compute_rf_image(
|
||
self._sras, self._angle, dc_threshold_mv=self._dc_threshold,
|
||
apply_bg_sub=self._apply_bg_sub, n_fft=self._n_fft,
|
||
dc4_mv=self._dc4_mv, should_stop=self._stopped)
|
||
else:
|
||
img = dc_image_mv(self._sras, self._angle, self._ch,
|
||
should_stop=self._stopped)
|
||
# On cancellation the image is only partly filled, so hand back
|
||
# None rather than something that would be cached as real. The
|
||
# signal still fires either way — it is what quits the thread.
|
||
self.finished.emit(None if self._stop else img)
|
||
except Exception as exc:
|
||
self.error.emit(str(exc))
|
||
|
||
|
||
class DcPrecomputeWorker(CancellableWorker):
|
||
"""Computes CH3/CH4 DC images for every angle in the background.
|
||
|
||
DC images are cheap (a per-waveform mean, no FFT) compared to the
|
||
CH1/Velocity FFT, so precomputing them for the whole file right after load
|
||
makes switching angles instant while on a DC channel, and also means the
|
||
FFT masking step (which needs a DC4 image) rarely has to wait on anything.
|
||
|
||
Angles are computed on a thread pool — the work is a pure mean over the
|
||
waveform block, so it is I/O- and bandwidth-bound and embarrassingly
|
||
parallel. Results are emitted one at a time as they land (out of angle
|
||
order), and always from this worker's own thread: nothing emits a Qt
|
||
signal from a pool thread.
|
||
"""
|
||
angle_done = pyqtSignal(int, np.ndarray, np.ndarray) # angle_idx, dc3_mv, dc4_mv
|
||
finished = pyqtSignal()
|
||
error = pyqtSignal(str)
|
||
|
||
def __init__(self, sras: SrasFile):
|
||
super().__init__()
|
||
self._sras = sras
|
||
|
||
def _one_angle(self, a: int) -> tuple[int, np.ndarray, np.ndarray]:
|
||
# max_workers=1 *and* a budget share: this call is one of several
|
||
# concurrent angles, and both the thread count and the buffer size
|
||
# have to be divided (see compute.plan_angle_level).
|
||
kw = dict(max_workers=1, budget=self._angle_budget,
|
||
should_stop=self._stopped)
|
||
return (a,
|
||
dc_image_mv(self._sras, a, CH3_IDX, **kw),
|
||
dc_image_mv(self._sras, a, CH4_IDX, **kw))
|
||
|
||
def run(self):
|
||
try:
|
||
n = self._sras.n_angles
|
||
n_workers, self._angle_budget = compute.plan_angle_level(self._sras)
|
||
pool = ThreadPoolExecutor(max_workers=n_workers)
|
||
try:
|
||
futures = {pool.submit(self._one_angle, a): a for a in range(n)}
|
||
for fut in as_completed(futures):
|
||
if self._stop:
|
||
break
|
||
a, dc3, dc4 = fut.result()
|
||
self.angle_done.emit(a, dc3, dc4)
|
||
finally:
|
||
# cancel_futures drops the queued angles; should_stop lets the
|
||
# in-flight ones bail within a chunk. Not waiting here is what
|
||
# keeps closing the window responsive on a large scan.
|
||
pool.shutdown(wait=not self._stop, cancel_futures=True)
|
||
self.finished.emit()
|
||
except Exception as exc:
|
||
self.error.emit(str(exc))
|
||
|
||
|
||
class BatchCacheWorker(QObject):
|
||
"""Batch-computes and stores DC or FFT images into each of *paths*'s v7
|
||
CACH tail, in place — converting v6 sources to v7 on first use, or updating
|
||
an existing v7 file's cache blocks without disturbing whatever the other
|
||
block already holds.
|
||
|
||
*mode* is ``"dc"`` (CH3/CH4 mean images) or ``"fft"`` (CH1 peak-frequency
|
||
images, unmasked — masking is applied at display time, same as v5's PREC
|
||
convention).
|
||
|
||
Files are processed one per subprocess: they are fully independent, each
|
||
opens its own memmap and writes only its own bytes, and only path strings
|
||
and scalars cross the process boundary. Emits ``progress(int)`` (0–100 by
|
||
files completed), ``file_done(str, str)`` (path, error message or "") so
|
||
one file's failure doesn't abort the batch, and ``finished()``.
|
||
"""
|
||
progress = pyqtSignal(int)
|
||
file_done = pyqtSignal(str, str)
|
||
finished = pyqtSignal()
|
||
|
||
def __init__(self, paths: list[str], mode: str, apply_bg_sub: bool):
|
||
super().__init__()
|
||
self._paths = paths
|
||
self._mode = mode
|
||
self._apply_bg_sub = apply_bg_sub
|
||
|
||
def _report(self, path: str, err: str, done: int, total: int):
|
||
self.file_done.emit(path, err)
|
||
self.progress.emit(int(done / max(1, total) * 100))
|
||
|
||
def _run_pooled(self, paths: list[str], n_procs: int) -> list[str]:
|
||
"""Process the batch across *n_procs* subprocesses. Returns the paths
|
||
that never got a real answer because the pool itself died, so the
|
||
caller can retry them in-process.
|
||
|
||
Under spawn each child re-imports the entry module, so the batch must
|
||
survive that going wrong (an unguarded __main__, a frozen build, a
|
||
sandbox that forbids subprocesses) rather than reporting every file as
|
||
failed — hence the retry list instead of a per-file error.
|
||
"""
|
||
# Each child threads internally; divide the machine rather than
|
||
# letting every process claim every core.
|
||
per_proc_workers = max(1, (os.cpu_count() or 4) // n_procs)
|
||
unresolved: list[str] = []
|
||
done = 0
|
||
|
||
with ProcessPoolExecutor(max_workers=n_procs) as executor:
|
||
futures = {
|
||
executor.submit(cache_file, p, self._mode, self._apply_bg_sub,
|
||
compute.get_fft_backend(), per_proc_workers): p
|
||
for p in paths
|
||
}
|
||
for fut in as_completed(futures):
|
||
path = futures[fut]
|
||
try:
|
||
err = fut.result()
|
||
except BrokenProcessPool:
|
||
unresolved.append(path)
|
||
continue
|
||
except Exception as exc:
|
||
err = str(exc)
|
||
done += 1
|
||
self._report(path, err, done, len(paths))
|
||
|
||
return unresolved
|
||
|
||
def _run_inline(self, paths: list[str], done: int, total: int):
|
||
"""Fallback / single-file path: compute in this thread. Still uses the
|
||
full core count internally, since nothing else is competing."""
|
||
for path in paths:
|
||
try:
|
||
err = cache_file(path, self._mode, self._apply_bg_sub,
|
||
compute.get_fft_backend(), compute._MAX_WORKERS)
|
||
except Exception as exc:
|
||
err = str(exc)
|
||
done += 1
|
||
self._report(path, err, done, total)
|
||
|
||
def _worth_pooling(self, paths: list[str]) -> bool:
|
||
if len(paths) < 2:
|
||
return False
|
||
total = 0
|
||
for p in paths:
|
||
try:
|
||
total += os.path.getsize(p)
|
||
except OSError:
|
||
pass # unreadable files are reported by cache_file
|
||
return total >= _BATCH_POOL_MIN_BYTES
|
||
|
||
def run(self):
|
||
paths = self._paths
|
||
n_procs = max(1, min(_BATCH_MAX_PROCS, len(paths)))
|
||
|
||
if not self._worth_pooling(paths):
|
||
self._run_inline(paths, 0, len(paths))
|
||
self.finished.emit()
|
||
return
|
||
|
||
try:
|
||
unresolved = self._run_pooled(paths, n_procs)
|
||
except Exception:
|
||
# The pool could not be created or collapsed wholesale.
|
||
unresolved = list(paths)
|
||
|
||
if unresolved:
|
||
self._run_inline(unresolved, len(paths) - len(unresolved), len(paths))
|
||
|
||
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.
|
||
Rotation is analytic (from sras.angles_deg); only translation is found by
|
||
phase correlation.
|
||
"""
|
||
progress = pyqtSignal(int) # 0–100
|
||
finished = pyqtSignal(object, str) # AlignmentResult|None, error ("" = success)
|
||
|
||
def __init__(self, sras: SrasFile, ref_angle_idx: int, dc_threshold_mv: float):
|
||
super().__init__()
|
||
self._sras = sras
|
||
self._ref = ref_angle_idx
|
||
self._threshold = dc_threshold_mv
|
||
|
||
def run(self):
|
||
try:
|
||
result = compute_angle_alignment(
|
||
self._sras, self._ref, self._threshold,
|
||
progress_cb=self.progress.emit)
|
||
self.finished.emit(result, "")
|
||
except Exception as exc:
|
||
self.finished.emit(None, str(exc))
|
||
|
||
|
||
class Ch4MaskWorker(QObject):
|
||
"""Fetches each requested angle's CH4 (Bias B) DC image in mV, for
|
||
ManualAlignmentDialog's initial threshold-mask overlay.
|
||
|
||
Reuses dc_image_mv, which prefers a stored v5/v7 cache over recomputing
|
||
from raw waveforms, so this only does real work for a file that hasn't
|
||
gone through the v7 "Convert" batch step and for angles the main
|
||
window's own DcPrecomputeWorker (which runs automatically right after
|
||
every file load) hasn't reached yet. In the common case — the user opens
|
||
Fusion -> Manual Alignment after DC precompute has already finished —
|
||
*angle_indices* is empty and this worker is never even constructed (see
|
||
ManualAlignmentDialog._start_mask_prep).
|
||
"""
|
||
angle_done = pyqtSignal(int, np.ndarray) # angle_idx, dc4_mv
|
||
finished = pyqtSignal()
|
||
error = pyqtSignal(str)
|
||
|
||
def __init__(self, sras: SrasFile, angle_indices: list[int]):
|
||
super().__init__()
|
||
self._sras = sras
|
||
self._angles = angle_indices
|
||
|
||
def run(self):
|
||
try:
|
||
n_workers, budget = compute.plan_angle_level(self._sras)
|
||
pool = ThreadPoolExecutor(max_workers=n_workers)
|
||
try:
|
||
futures = {
|
||
pool.submit(dc_image_mv, self._sras, a, CH4_IDX,
|
||
max_workers=1, budget=budget): a
|
||
for a in self._angles
|
||
}
|
||
for fut in as_completed(futures):
|
||
a = futures[fut]
|
||
self.angle_done.emit(a, fut.result())
|
||
finally:
|
||
pool.shutdown(wait=True)
|
||
self.finished.emit()
|
||
except Exception as exc:
|
||
self.error.emit(str(exc))
|
||
|
||
|
||
class CrossCorrelateWorker(QObject):
|
||
"""FFT phase-correlation translation for each of *angle_indices* against
|
||
*ref_angle_idx*, for ManualAlignmentDialog's Auto Cross-Correlate button.
|
||
|
||
Runs on a background thread — a real many-angle, high-resolution scan's
|
||
correlation (even at its downsampled working resolution) can take long
|
||
enough that doing all of them on the GUI thread would visibly freeze the
|
||
dialog. Rotation is set to the same analytic scan-angle delta Auto
|
||
De-rotate uses alongside the correlated shift, since a translation
|
||
search is only meaningful once both angles' content is already oriented
|
||
the same way. dc4_mv/pivot_mm are the dialog's own already-in-memory
|
||
per-angle images/pivots — this worker does no fetching of its own.
|
||
"""
|
||
angle_done = pyqtSignal(int, float, float, float) # angle_idx, rotation_deg, shift_x_mm, shift_y_mm
|
||
finished = pyqtSignal()
|
||
error = pyqtSignal(str)
|
||
|
||
def __init__(self, sras: SrasFile, ref_angle_idx: int, angle_indices: list[int],
|
||
dc4_mv: dict[int, np.ndarray], pivot_mm: dict[int, tuple[float, float]],
|
||
*, use_mask: bool, dc_threshold_mv: float, margin_frac: float):
|
||
super().__init__()
|
||
self._sras = sras
|
||
self._ref = ref_angle_idx
|
||
self._angles = angle_indices
|
||
self._dc4_mv = dc4_mv
|
||
self._pivot_mm = pivot_mm
|
||
self._use_mask = use_mask
|
||
self._threshold = dc_threshold_mv
|
||
self._margin = margin_frac
|
||
|
||
def _one(self, a: int) -> tuple[int, float, float, float]:
|
||
theta = compute._theta_deg(self._sras, a, self._ref)
|
||
dx, dy = compute.correlate_translation_mm(
|
||
self._sras, a, self._ref, self._dc4_mv, self._pivot_mm,
|
||
use_mask=self._use_mask, dc_threshold_mv=self._threshold,
|
||
margin_frac=self._margin)
|
||
return a, theta, dx, dy
|
||
|
||
def run(self):
|
||
try:
|
||
n_workers, _budget = compute.plan_angle_level(self._sras)
|
||
pool = ThreadPoolExecutor(max_workers=max(1, n_workers))
|
||
try:
|
||
futures = [pool.submit(self._one, a) for a in self._angles]
|
||
for fut in as_completed(futures):
|
||
a, theta, dx, dy = fut.result()
|
||
self.angle_done.emit(a, theta, dx, dy)
|
||
finally:
|
||
pool.shutdown(wait=True)
|
||
self.finished.emit()
|
||
except Exception as exc:
|
||
self.error.emit(str(exc))
|