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>
1056 lines
48 KiB
Python
1056 lines
48 KiB
Python
#!/usr/bin/env python3
|
|
"""Image computation and angle alignment for .sras scans.
|
|
|
|
Depends only on numpy/scipy (+ optional pyfftw) and sras_format, so a
|
|
multiprocessing child can import it without loading Qt or matplotlib —
|
|
which matters because Python 3.14 on macOS spawns rather than forks.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import scipy.fft as scipy_fft
|
|
import scipy.ndimage as scipy_ndimage
|
|
|
|
from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, SrasFile, adc_to_mv
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# FFT backend
|
|
# ---------------------------------------------------------------------------
|
|
|
|
try:
|
|
import pyfftw
|
|
pyfftw.interfaces.cache.enable()
|
|
PYFFTW_AVAILABLE = True
|
|
except ImportError:
|
|
PYFFTW_AVAILABLE = False
|
|
|
|
_fft_backend = "numpy" # "numpy" or "pyfftw"; set via set_fft_backend()
|
|
|
|
|
|
def set_fft_backend(name: str):
|
|
"""Select the rfft implementation. Module-level state, so it must be set
|
|
explicitly inside each multiprocessing child — it does not survive a
|
|
spawn."""
|
|
global _fft_backend
|
|
_fft_backend = name if (name != "pyfftw" or PYFFTW_AVAILABLE) else "numpy"
|
|
|
|
|
|
def get_fft_backend() -> str:
|
|
return _fft_backend
|
|
|
|
|
|
def _do_rfft(x: np.ndarray, n: int | None = None, axis: int = -1,
|
|
workers: int = 1) -> np.ndarray:
|
|
"""Dispatch rfft to the selected backend with optional multithreading."""
|
|
if _fft_backend == "pyfftw" and PYFFTW_AVAILABLE:
|
|
return pyfftw.interfaces.numpy_fft.rfft(x, n=n, axis=axis, threads=workers)
|
|
return scipy_fft.rfft(x, n=n, axis=axis, workers=workers)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Chunking / parallel budget
|
|
# ---------------------------------------------------------------------------
|
|
#
|
|
# Rows are batched so the float32 working buffers for one chunk stay under a
|
|
# memory budget. A fixed row count (the original design) works fine for small
|
|
# legacy scans but is catastrophic for a v6 scan with a large per-angle
|
|
# frame/sample count — e.g. a 7500-frame x 2500-sample angle needs ~2.4 GB for
|
|
# a single 32-row chunk.
|
|
#
|
|
# With chunks running concurrently the budget has to cover *all* live chunks at
|
|
# once. Note that on a large scan chunk_rows is already clamped to its floor of
|
|
# 1 row (one row alone is ~75 MB of float32 at 7507x2500), so shrinking the
|
|
# per-chunk size cannot buy more concurrency — the worker count must be derived
|
|
# from the budget instead. See _plan_chunks.
|
|
|
|
# 1024 MB is the measured knee on a 16-core machine against a 7507-frame x
|
|
# 2500-sample angle: 512 MB leaves ~20% of the FFT speedup on the table, and
|
|
# 1536+ MB costs ~0.4 GB more resident for no further gain. Override with
|
|
# SRAS_MEM_BUDGET_MB on a smaller machine.
|
|
_TOTAL_BYTES_BUDGET = int(os.environ.get("SRAS_MEM_BUDGET_MB", 1024)) * 1024 * 1024
|
|
_CHUNK_ROWS_MAX = 32 # cap for small scans (original behavior)
|
|
_MAX_WORKERS = int(os.environ.get("SRAS_MAX_WORKERS", 0)) or (os.cpu_count() or 4)
|
|
|
|
# An rfft chunk holds, live at once: the float32 input, the complex64
|
|
# transform, and the float32 power spectrum — roughly 3x the input buffer.
|
|
_FFT_LIVE_MULTIPLIER = 3
|
|
|
|
|
|
def _chunk_rows_for(n_frames: int, samples_per_frame: int,
|
|
budget: int = _TOTAL_BYTES_BUDGET) -> int:
|
|
bytes_per_row = max(1, n_frames * samples_per_frame * 4) # float32
|
|
return int(max(1, min(_CHUNK_ROWS_MAX, budget // bytes_per_row)))
|
|
|
|
|
|
def _plan_chunks(n_rows: int, n_frames: int, samples_per_frame: int,
|
|
live_multiplier: int = 1,
|
|
max_workers: int | None = None,
|
|
budget: int | None = None) -> tuple[int, int]:
|
|
"""(chunk_rows, n_workers) such that all concurrent chunks together fit
|
|
the budget: n_workers x live_multiplier x chunk_rows x n_frames x
|
|
samples x 4 bytes <= budget.
|
|
|
|
The worker count is chosen *first* and the chunk sized to it. Sizing the
|
|
chunk first is the trap: _chunk_rows_for spends whatever budget it is
|
|
given, so a single chunk would always consume the lot and leave room for
|
|
exactly one worker — no concurrency, on precisely the large scans that
|
|
need it most.
|
|
|
|
A caller that is itself running several of these concurrently must pass
|
|
*both* max_workers=1 and its share of the budget. Capping the workers
|
|
alone is not enough: the chunk would still be sized against the whole
|
|
budget, and N concurrent callers would each allocate all of it.
|
|
"""
|
|
total = _TOTAL_BYTES_BUDGET if budget is None else max(1, budget)
|
|
cap = max_workers if max_workers is not None else _MAX_WORKERS
|
|
bytes_per_row = max(1, live_multiplier * n_frames * samples_per_frame * 4)
|
|
|
|
# A chunk is at least one row, so that alone caps how many can be live.
|
|
n_workers = int(max(1, min(cap, n_rows, total // bytes_per_row)))
|
|
chunk_rows = _chunk_rows_for(n_frames, samples_per_frame,
|
|
max(1, total // (n_workers * live_multiplier)))
|
|
# With chunk_rows known, more workers than chunks buys nothing.
|
|
n_workers = int(max(1, min(n_workers, -(-n_rows // chunk_rows))))
|
|
return chunk_rows, n_workers
|
|
|
|
|
|
def _map_row_chunks(n_rows: int, chunk_rows: int, n_workers: int, fn,
|
|
should_stop=None):
|
|
"""Apply fn(r0, r1) over row chunks, in parallel when it pays.
|
|
|
|
Chunks write to disjoint output slices, so no locking is needed. numpy
|
|
ufuncs, scipy's pocketfft, and memmap page faults all release the GIL, so
|
|
threads give real parallelism here — and on a very large file they also
|
|
keep many more page-fault requests in flight, which is what the I/O path
|
|
wants.
|
|
|
|
*should_stop* is polled per chunk so a cancelled job abandons work within
|
|
one chunk rather than one whole angle — on a large scan an angle is
|
|
~40 s, which is far too long to block application shutdown.
|
|
"""
|
|
bounds = [(r0, min(r0 + chunk_rows, n_rows))
|
|
for r0 in range(0, n_rows, chunk_rows)]
|
|
|
|
def run(b):
|
|
if should_stop is not None and should_stop():
|
|
return
|
|
fn(*b)
|
|
|
|
if n_workers <= 1 or len(bounds) == 1:
|
|
for b in bounds:
|
|
run(b)
|
|
return
|
|
with ThreadPoolExecutor(max_workers=min(n_workers, len(bounds))) as pool:
|
|
list(pool.map(run, bounds))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Image computation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def plan_angle_level(sras: SrasFile,
|
|
live_multiplier: int = 1) -> tuple[int, int]:
|
|
"""(n_angle_workers, per_angle_budget) for a caller that parallelises
|
|
over angles instead of over rows.
|
|
|
|
Callers that parallelise over angles must not also let each angle
|
|
parallelise over rows: the two levels multiply, in threads and in memory.
|
|
Each angle is therefore given max_workers=1 and the returned budget share,
|
|
which together keep total live buffers within _TOTAL_BYTES_BUDGET.
|
|
"""
|
|
biggest = max(range(sras.n_angles), key=lambda a: int(sras.n_frames[a]))
|
|
_, n_workers = _plan_chunks(int(sras.n_rows[biggest]),
|
|
int(sras.n_frames[biggest]),
|
|
sras.samples_per_frame,
|
|
live_multiplier=live_multiplier)
|
|
n_workers = max(1, min(n_workers, sras.n_angles))
|
|
return n_workers, max(1, _TOTAL_BYTES_BUDGET // n_workers)
|
|
|
|
|
|
def compute_dc_image(sras: SrasFile, angle_idx: int, ch_idx: int,
|
|
max_workers: int | None = None,
|
|
budget: int | None = None,
|
|
should_stop=None) -> np.ndarray:
|
|
"""Mean of each waveform → (n_rows, n_frames) float32, in ADC counts.
|
|
|
|
Pass max_workers=1 when the caller is already parallelising over angles.
|
|
If *should_stop* ever returns True the result is incomplete — the caller
|
|
is expected to be abandoning it.
|
|
"""
|
|
n_rows, n_frames = sras.image_shape(angle_idx)
|
|
data = sras.data[angle_idx]
|
|
chunk_rows, n_workers = _plan_chunks(n_rows, n_frames, sras.samples_per_frame,
|
|
max_workers=max_workers, budget=budget)
|
|
img = np.empty((n_rows, n_frames), dtype=np.float32)
|
|
|
|
def chunk(r0: int, r1: int):
|
|
img[r0:r1] = data[r0:r1, ch_idx, :, :].astype(np.float32).mean(axis=-1)
|
|
|
|
_map_row_chunks(n_rows, chunk_rows, n_workers, chunk, should_stop=should_stop)
|
|
return img
|
|
|
|
|
|
def dc_image_mv(sras: SrasFile, angle_idx: int, ch_idx: int,
|
|
max_workers: int | None = None, budget: int | None = None,
|
|
should_stop=None) -> np.ndarray:
|
|
"""DC image for (angle, channel) in mV, preferring a stored v5/v7 cache."""
|
|
cached = sras.cached_dc_mv(angle_idx, ch_idx)
|
|
if cached is not None:
|
|
return cached
|
|
return adc_to_mv(
|
|
compute_dc_image(sras, angle_idx, ch_idx, max_workers=max_workers,
|
|
budget=budget, should_stop=should_stop),
|
|
*sras.cal(ch_idx))
|
|
|
|
|
|
def compute_rf_image(sras: SrasFile, angle_idx: int,
|
|
dc_threshold_mv: float | None,
|
|
apply_bg_sub: bool = True,
|
|
n_fft: int | None = None,
|
|
dc4_mv: np.ndarray | None = None,
|
|
max_workers: int | None = None,
|
|
budget: int | None = None,
|
|
should_stop=None) -> np.ndarray:
|
|
"""FFT of each CH1 waveform; pixel = peak frequency in MHz.
|
|
|
|
Pixels where CH4_dc < dc_threshold_mv are set to 0 — and the FFT is
|
|
never run for them, since that's the expensive part. The masked-out
|
|
CH1 samples are also never *read*: the boolean mask is applied to the
|
|
raw memmap slice before any dtype conversion, so numpy only pages in
|
|
the bytes for pixels that pass the threshold (an untouched memmap page
|
|
is never read from disk).
|
|
|
|
*dc_threshold_mv* of None means "mask nothing", and skips reading and
|
|
averaging CH4 entirely — worth a third of the I/O when caching a whole
|
|
file, where the mask is applied later at display time.
|
|
|
|
If the DC4 image for this angle is already known (e.g. from the
|
|
DC-channel precompute cache), pass it as *dc4_mv* (mV, shape
|
|
(n_rows, n_frames)) to reuse it instead of re-reading CH4 here.
|
|
|
|
Fast path: if the file has a precomputed peak-frequency image for this
|
|
angle (v5 PREC or v7 CACH), zero-padding is off, and the bg-sub flag
|
|
matches, the stored image is used directly — no FFT is run.
|
|
"""
|
|
n_rows, n_frames = sras.image_shape(angle_idx)
|
|
data = sras.data[angle_idx]
|
|
|
|
# ---- Fast path: precomputed image (v5 PREC or v7 CACH) ----------------
|
|
cached_freq = sras.precomputed_freq_mhz[angle_idx]
|
|
if (cached_freq is not None
|
|
and n_fft is None # no custom zero-padding
|
|
and sras.precomputed_bg_sub == (apply_bg_sub and sras.background is not None)):
|
|
freq_img = cached_freq.copy()
|
|
if dc_threshold_mv is not None:
|
|
# DC4 mask, in priority order: already-cached DC block, caller-
|
|
# supplied image, or a fresh (cheap — no FFT) recompute.
|
|
dc4_img = sras.cached_dc_mv(angle_idx, CH4_IDX)
|
|
if dc4_img is None:
|
|
dc4_img = dc4_mv if dc4_mv is not None else adc_to_mv(
|
|
compute_dc_image(sras, angle_idx, CH4_IDX), *sras.cal(CH4_IDX))
|
|
freq_img[dc4_img < dc_threshold_mv] = 0.0
|
|
return freq_img
|
|
|
|
# ---- Chunked FFT path --------------------------------------------------
|
|
freq_axis = sras.freq_axis_mhz(n_fft)
|
|
img = np.zeros((n_rows, n_frames), dtype=np.float32)
|
|
n_fft_bins = n_fft if n_fft is not None else sras.samples_per_frame
|
|
chunk_rows, n_workers = _plan_chunks(
|
|
n_rows, n_frames, max(sras.samples_per_frame, n_fft_bins),
|
|
live_multiplier=_FFT_LIVE_MULTIPLIER, max_workers=max_workers,
|
|
budget=budget)
|
|
background = sras.background if (apply_bg_sub and sras.background is not None) else None
|
|
cal4 = sras.cal(CH4_IDX)
|
|
|
|
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]
|
|
else:
|
|
dc4_chunk = adc_to_mv(
|
|
data[r0:r1, CH4_IDX, :, :].astype(np.float32).mean(axis=-1), *cal4)
|
|
valid = dc4_chunk >= dc_threshold_mv # True = run the FFT
|
|
if not valid.any():
|
|
return
|
|
|
|
# Index the raw memmap slice with the boolean mask *before*
|
|
# converting dtype — this is a lazy view until touched, so only the
|
|
# selected elements are actually read from disk; masked-out pixels'
|
|
# pages are never paged in at all.
|
|
raw = data[r0:r1, CH1_IDX, :, :]
|
|
waves = (raw[valid] if valid is not None else raw).astype(np.float32)
|
|
|
|
if background is not None:
|
|
waves -= background # background is 1-D (spf,)
|
|
|
|
# Parallelism comes from the outer chunk loop, so keep the inner
|
|
# transform single-threaded to avoid oversubscribing the machine.
|
|
spectrum = _do_rfft(waves, n=n_fft, axis=-1, workers=1)
|
|
del waves
|
|
# |z|^2 without np.abs()'s extra full-size temporary.
|
|
power = spectrum.real ** 2
|
|
power += spectrum.imag ** 2
|
|
del spectrum
|
|
# [..., 0] not [:, 0]: the unmasked path keeps the (rows, frames,
|
|
# bins) shape, where [:, 0] would blank a whole frame.
|
|
power[..., 0] = 0.0 # suppress DC bin
|
|
peak_bins = np.argmax(power, axis=-1)
|
|
del power
|
|
|
|
if valid is not None:
|
|
img[r0:r1][valid] = freq_axis[peak_bins]
|
|
else:
|
|
img[r0:r1] = freq_axis[peak_bins]
|
|
|
|
_map_row_chunks(n_rows, chunk_rows, n_workers, chunk, should_stop=should_stop)
|
|
return img
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Batch cache (Convert menu) — process-pool entry point
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def cache_file(path: str, mode: str, apply_bg_sub: bool,
|
|
fft_backend: str = "numpy", max_workers: int = 0) -> str:
|
|
"""Compute and store DC or FFT images for every angle of one file,
|
|
converting v6 → v7 in place. Returns "" on success or an error message.
|
|
|
|
Module-level and picklable so it can run in a ProcessPoolExecutor. The
|
|
FFT backend and worker cap are passed explicitly because module globals
|
|
do not survive a spawn.
|
|
"""
|
|
global _MAX_WORKERS
|
|
try:
|
|
if mode not in ("dc", "fft"):
|
|
return f"unknown cache mode {mode!r} (expected 'dc' or 'fft')"
|
|
set_fft_backend(fft_backend)
|
|
if max_workers:
|
|
_MAX_WORKERS = max_workers
|
|
|
|
sras = SrasFile(path)
|
|
if sras.version not in (6, 7):
|
|
return (f"unsupported version {sras.version} — only v6/v7 files "
|
|
"can be batch-cached")
|
|
|
|
n = sras.n_angles
|
|
if mode == "dc":
|
|
dc3 = [adc_to_mv(compute_dc_image(sras, a, CH3_IDX), *sras.cal(CH3_IDX))
|
|
for a in range(n)]
|
|
dc4 = [adc_to_mv(compute_dc_image(sras, a, CH4_IDX), *sras.cal(CH4_IDX))
|
|
for a in range(n)]
|
|
sras.write_v7_cache(new_dc3_mv=dc3, new_dc4_mv=dc4)
|
|
else:
|
|
effective_bg = apply_bg_sub and sras.background is not None
|
|
# dc_threshold_mv=None: store unmasked images and mask at display
|
|
# time (same convention as v5's PREC block). Skipping the mask
|
|
# also skips reading CH4 entirely.
|
|
freq = [compute_rf_image(sras, a, dc_threshold_mv=None,
|
|
apply_bg_sub=effective_bg)
|
|
for a in range(n)]
|
|
sras.write_v7_cache(new_freq_mhz=freq, new_bg_sub=effective_bg)
|
|
return ""
|
|
except Exception as exc:
|
|
return str(exc)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Angle alignment (Fusion menu)
|
|
#
|
|
# Puts every angle's images onto one shared, zero-padded pixel grid using a
|
|
# rigid transform only (rotation + translation, never scale). Rotation for
|
|
# angle `a` is the *known* scan-angle delta relative to a reference angle —
|
|
# never searched. Only the residual translation is found, via FFT phase
|
|
# correlation of each angle's binarized CH4 ("dc-mask") image.
|
|
#
|
|
# Rotation is done in physical mm space rather than on raw pixel indices:
|
|
# the x-pixel pitch (SrasFile.pixel_x_mm) is file-wide constant but the
|
|
# y-pixel pitch (row spacing) can differ from it, and for v6 files can even
|
|
# vary per angle. Rotating the raw index grid directly would implicitly
|
|
# assume square pixels and shear a non-square-pixel image — an unwanted
|
|
# effective anisotropic scale. Instead each angle gets one affine that maps
|
|
# shared-canvas pixel index -> mm -> undo rotation/shift -> that angle's own
|
|
# local mm -> that angle's own raw pixel index, matching the output->input
|
|
# convention scipy.ndimage.affine_transform expects.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@dataclass
|
|
class AngleTransform:
|
|
rotation_deg: float
|
|
shift_mm: tuple[float, float] # (dx_mm, dy_mm) found by phase correlation
|
|
matrix: np.ndarray # (2,2): canvas (row,col) -> this angle's raw (row,col)
|
|
offset: np.ndarray # (2,)
|
|
|
|
|
|
@dataclass
|
|
class AlignmentResult:
|
|
ref_angle_idx: int
|
|
dc_threshold_mv: float
|
|
canvas_shape: tuple[int, int] # (n_rows, n_cols)
|
|
canvas_dx_mm: float
|
|
canvas_dy_mm: float
|
|
canvas_origin_mm: tuple[float, float] # mm at canvas pixel index (0, 0)
|
|
per_angle: dict[int, AngleTransform]
|
|
|
|
|
|
@dataclass
|
|
class ManualAngleParams:
|
|
"""One angle's manual-alignment state, independent of any canvas.
|
|
|
|
rotation_deg/shift_mm are exactly AngleTransform's non-derived fields —
|
|
the pair a canvas-bound AngleTransform's matrix/offset get built from
|
|
once a canvas is decided (build_manual_alignment). Defaults to identity
|
|
(no rotation, no shift): a fresh angle with no prior alignment is shown
|
|
raw, exactly as scanned — the same "fully unaligned" state Clear
|
|
Alignment resets back to.
|
|
"""
|
|
rotation_deg: float = 0.0
|
|
shift_mm: tuple[float, float] = (0.0, 0.0)
|
|
|
|
|
|
def _pixel_pitch_mm(sras: SrasFile, angle_idx: int) -> tuple[float, float]:
|
|
"""(dx, dy) mm/pixel for one angle: dx is the file-wide constant
|
|
pixel_x_mm; dy is this angle's own row spacing (assumed uniform, the same
|
|
assumption _redraw_image makes when it builds the display extent)."""
|
|
y = sras.y_positions_mm(angle_idx)
|
|
return sras.pixel_x_mm, (float(y[1] - y[0]) if len(y) > 1 else 1.0)
|
|
|
|
|
|
def _bbox_center_mm(sras: SrasFile, angle_idx: int) -> tuple[float, float]:
|
|
x = sras.x_axis_mm(angle_idx)
|
|
y = sras.y_positions_mm(angle_idx)
|
|
return float((x[0] + x[-1]) / 2.0), float((y[0] + y[-1]) / 2.0)
|
|
|
|
|
|
def _bbox_corners_mm(sras: SrasFile, angle_idx: int) -> np.ndarray:
|
|
"""4 corners (x, y) of this angle's raw mm bounding box, shape (4, 2)."""
|
|
x = sras.x_axis_mm(angle_idx)
|
|
y = sras.y_positions_mm(angle_idx)
|
|
return np.array([[xx, yy] for xx in (x[0], x[-1]) for yy in (y[0], y[-1])])
|
|
|
|
|
|
def _rotation_matrix(theta_deg: float) -> np.ndarray:
|
|
t = np.radians(theta_deg)
|
|
c, s = np.cos(t), np.sin(t)
|
|
return np.array([[c, -s], [s, c]]) # CCW rotation acting on (x, y)
|
|
|
|
|
|
def _theta_deg(sras: SrasFile, angle_idx: int, ref_idx: int) -> float:
|
|
"""CCW rotation, in degrees and in _rotation_matrix's convention, that
|
|
maps angle_idx's own local mm frame onto ref_idx's.
|
|
|
|
This is the *negative* of the raw angles_deg delta: the GR rotation
|
|
stage's reported angle increases in the opposite rotational sense from
|
|
this module's math-positive (CCW, x toward y) convention in scan mm
|
|
space. Rotating by +(angles_deg[a] - angles_deg[ref]) therefore turns
|
|
misalignment the wrong way — confirmed empirically (Auto De-rotate made
|
|
real scans worse, not better, before this negation).
|
|
"""
|
|
return -float(sras.angles_deg[angle_idx] - sras.angles_deg[ref_idx])
|
|
|
|
|
|
def _signal_centroid_mm(sras: SrasFile, angle_idx: int,
|
|
dc4_mv: np.ndarray) -> tuple[float, float]:
|
|
"""Intensity-weighted centroid (mean x, mean y, weighted by CH4 signal
|
|
after subtracting this angle's own minimum) in angle_idx's own local mm
|
|
frame — the alignment pivot used in place of the raw scan-window bbox
|
|
center (see compute_pivot_points_mm).
|
|
|
|
Weighting by the continuous DC signal, rather than a binary >=
|
|
dc_threshold_mv mask, means the pivot never depends on how well one
|
|
shared threshold happens to suit this particular angle: real signal
|
|
levels vary scan to scan, so a threshold tuned for one angle can leave
|
|
another angle's binary mask empty — and a centroid of an empty mask has
|
|
nothing to fall back to *except* the raw bbox center, silently
|
|
reproducing the exact "aligned to the scan window, not the sample"
|
|
problem this pivot exists to avoid. Falls back to the bbox center only
|
|
in the fully-degenerate case of a perfectly flat signal (nothing to
|
|
weight by at all).
|
|
"""
|
|
weights = dc4_mv - dc4_mv.min()
|
|
total = float(weights.sum())
|
|
if total <= 0.0:
|
|
return _bbox_center_mm(sras, angle_idx)
|
|
x = sras.x_axis_mm(angle_idx)
|
|
y = sras.y_positions_mm(angle_idx)
|
|
cx = float((x * weights.sum(axis=0)).sum() / total)
|
|
cy = float((y * weights.sum(axis=1)).sum() / total)
|
|
return cx, cy
|
|
|
|
|
|
def compute_pivot_points_mm(sras: SrasFile,
|
|
dc4_mv: dict[int, np.ndarray] | None = None
|
|
) -> dict[int, tuple[float, float]]:
|
|
"""Per-angle alignment pivot, in each angle's own local mm frame: the
|
|
CH4-signal-weighted centroid of its own footprint (see
|
|
_signal_centroid_mm), rather than the raw scan-window bbox center.
|
|
|
|
Pivoting on each angle's own content — instead of on wherever its
|
|
scanned window happened to sit in microscope/global XY space — is what
|
|
makes alignment purely relative *between scans* rather than to a global
|
|
coordinate system: the rotation stage's true mechanical axis need not
|
|
coincide with the scan window's geometric center, and the sample need
|
|
not be perfectly centered on that axis either, so a bbox-center pivot
|
|
leaves a residual orbital motion between angles that a content-centroid
|
|
pivot does not.
|
|
|
|
Deliberately independent of dc_threshold_mv (the RF-mask / overlay
|
|
threshold): that value is a display/masking choice and must never
|
|
silently change where alignment pivots.
|
|
|
|
*dc4_mv* lets a caller that has already computed each angle's CH4 mV
|
|
image (compute_angle_alignment's Step 1, or ManualAlignmentDialog's own
|
|
cache) reuse it instead of recomputing; angles missing from it are
|
|
computed fresh via dc_image_mv, which prefers a stored v5/v7 cache over
|
|
recomputing from raw waveforms.
|
|
"""
|
|
dc4_mv = dc4_mv or {}
|
|
pivots: dict[int, tuple[float, float]] = {}
|
|
for a in range(sras.n_angles):
|
|
img = dc4_mv.get(a)
|
|
if img is None:
|
|
img = dc_image_mv(sras, a, CH4_IDX)
|
|
pivots[a] = _signal_centroid_mm(sras, a, img)
|
|
return pivots
|
|
|
|
|
|
def _corners_in_ref_frame(sras: SrasFile, angle_idx: int, ref_idx: int,
|
|
pivot_mm: dict[int, tuple[float, float]],
|
|
shift_mm=(0.0, 0.0),
|
|
theta_deg: float | None = None) -> np.ndarray:
|
|
"""Angle *angle_idx*'s bbox corners, rotated about its own alignment
|
|
pivot (pivot_mm[angle_idx] — see compute_pivot_points_mm) into the
|
|
reference frame and translated by *shift_mm*. Shape (4, 2).
|
|
|
|
theta_deg overrides the analytic angles_deg-derived rotation used by
|
|
default — the manual-alignment path (union_canvas_mm) passes a
|
|
user-chosen rotation here (which may differ from the known scan-angle
|
|
delta) without needing a parallel code path.
|
|
"""
|
|
theta = _theta_deg(sras, angle_idx, ref_idx) if theta_deg is None else theta_deg
|
|
R = _rotation_matrix(theta)
|
|
c_a = np.array(pivot_mm[angle_idx])
|
|
c_ref = np.array(pivot_mm[ref_idx])
|
|
shift = np.asarray(shift_mm, dtype=np.float64)
|
|
return np.array([R @ (corner - c_a) + c_ref + shift
|
|
for corner in _bbox_corners_mm(sras, angle_idx)])
|
|
|
|
|
|
def _build_affine_canvas_to_raw(sras: SrasFile, angle_idx: int, ref_idx: int,
|
|
shift_mm: tuple[float, float],
|
|
canvas_dx: float, canvas_dy: float,
|
|
canvas_origin_mm: tuple[float, float],
|
|
pivot_mm: dict[int, tuple[float, float]],
|
|
theta_deg: float | None = None
|
|
) -> tuple[np.ndarray, np.ndarray]:
|
|
"""matrix, offset s.t. raw_index = matrix @ [row_out, col_out] + offset,
|
|
matching scipy.ndimage.affine_transform's output->input convention.
|
|
|
|
Pipeline (all mm unless noted):
|
|
[X;Y] = A_out @ [row_out;col_out] + b_out # canvas idx -> ref-frame mm
|
|
[lx;ly] = R(theta)^T @ ([X;Y]-c_ref-shift) + c_a # undo rotation+shift -> angle a's local mm
|
|
[row;col] = D @ ([lx;ly] - [x_start_a; y0_a]) # local mm -> angle a's raw idx
|
|
|
|
where theta is _theta_deg(angle_idx, ref_idx) unless *theta_deg*
|
|
overrides it (see _corners_in_ref_frame's note, used by the manual-
|
|
alignment path), c_ref/c_a are each angle's own alignment pivot
|
|
(pivot_mm — see compute_pivot_points_mm; the CH4-signal-weighted
|
|
centroid of its own footprint, not the raw scan-window bbox center —
|
|
this keeps the sample itself centered post-rotation, minimizing
|
|
required canvas padding and, more importantly, keeping alignment
|
|
relative to the sample rather than to wherever the scan window sat in
|
|
microscope/global XY space), and A_out/D are the index<->mm scaling
|
|
matrices for the canvas pitch and this angle's own native pitch
|
|
respectively.
|
|
"""
|
|
theta = _theta_deg(sras, angle_idx, ref_idx) if theta_deg is None else theta_deg
|
|
Rinv = _rotation_matrix(theta).T
|
|
cx_a, cy_a = pivot_mm[angle_idx]
|
|
cx_ref, cy_ref = pivot_mm[ref_idx]
|
|
dx_a, dy_a = _pixel_pitch_mm(sras, angle_idx)
|
|
x0_a = float(sras.x_start_mm[angle_idx])
|
|
y0_a = float(sras.y_positions_mm(angle_idx)[0])
|
|
|
|
A_out = np.array([[0.0, canvas_dx], [canvas_dy, 0.0]]) # [row,col] -> [X,Y]
|
|
b_out = np.array(canvas_origin_mm, dtype=np.float64)
|
|
D = np.array([[0.0, 1.0 / dy_a], [1.0 / dx_a, 0.0]]) # [x,y] -> [row,col]
|
|
shift = np.array(shift_mm, dtype=np.float64)
|
|
c_ref_v = np.array([cx_ref, cy_ref])
|
|
c_a_v = np.array([cx_a, cy_a])
|
|
origin_a = np.array([x0_a, y0_a])
|
|
|
|
matrix = D @ Rinv @ A_out
|
|
offset = D @ Rinv @ (b_out - c_ref_v - shift) + D @ (c_a_v - origin_a)
|
|
return matrix, offset
|
|
|
|
|
|
def apply_alignment(result: AlignmentResult, angle_idx: int, img: np.ndarray,
|
|
order: int = 0) -> np.ndarray:
|
|
"""Resample any already-computed 2D image for `angle_idx` (same shape as
|
|
that angle's raw (n_rows, n_frames)) onto the shared alignment canvas.
|
|
order=0 (nearest) avoids blending real data with zero-padding or with
|
|
masked-out (0-valued) CH1/velocity pixels at mask edges. Channel-
|
|
agnostic: the same per-angle transform (found from the CH4 mask) works
|
|
for any channel's image of that angle."""
|
|
t = result.per_angle[angle_idx]
|
|
return scipy_ndimage.affine_transform(
|
|
img.astype(np.float32, copy=False), t.matrix, offset=t.offset,
|
|
output_shape=result.canvas_shape, order=order,
|
|
mode="constant", cval=0.0)
|
|
|
|
|
|
def _block_mean_downsample(img: np.ndarray, factor: int) -> np.ndarray:
|
|
if factor <= 1:
|
|
return img
|
|
h, w = img.shape
|
|
h2, w2 = (h // factor) * factor, (w // factor) * factor
|
|
trimmed = img[:h2, :w2]
|
|
return trimmed.reshape(h2 // factor, factor, w2 // factor, factor).mean(axis=(1, 3))
|
|
|
|
|
|
def _phase_correlate_shift(ref_img: np.ndarray, mov_img: np.ndarray) -> tuple[int, int]:
|
|
"""FFT normalized cross-power-spectrum phase correlation. Returns the
|
|
integer (dr, dc) pixel shift of mov_img relative to ref_img; both must
|
|
be the same shape. Risk: if the true shift is near +/- half the array
|
|
size, wraparound can bias the peak — mitigated by the generous
|
|
margin_frac padding in _working_canvas_for_pair, which keeps the true
|
|
residual shift small relative to the correlation canvas."""
|
|
F1 = scipy_fft.fft2(ref_img.astype(np.float64), workers=-1)
|
|
F2 = scipy_fft.fft2(mov_img.astype(np.float64), workers=-1)
|
|
R = F1 * np.conj(F2)
|
|
R /= np.maximum(np.abs(R), 1e-12)
|
|
corr = scipy_fft.ifft2(R, workers=-1).real
|
|
dr, dc = np.unravel_index(np.argmax(corr), corr.shape)
|
|
h, w = corr.shape
|
|
if dr > h // 2:
|
|
dr -= h
|
|
if dc > w // 2:
|
|
dc -= w
|
|
return int(dr), int(dc)
|
|
|
|
|
|
def _working_canvas_for_pair(sras: SrasFile, ref_idx: int, a_idx: int,
|
|
dx: float, dy: float,
|
|
pivot_mm: dict[int, tuple[float, float]],
|
|
margin_frac: float = 0.3
|
|
) -> tuple[tuple[float, float], tuple[int, int]]:
|
|
"""Union of the reference's own raw bbox and angle a's raw bbox rotated
|
|
(about its own alignment pivot) into the ref frame with zero shift,
|
|
padded by margin_frac on each side — sized generously so the true
|
|
phase-correlation shift lands well inside the canvas (see
|
|
_phase_correlate_shift's wraparound note)."""
|
|
pts = np.vstack([_bbox_corners_mm(sras, ref_idx),
|
|
_corners_in_ref_frame(sras, a_idx, ref_idx, pivot_mm)])
|
|
x_min, y_min = pts.min(axis=0)
|
|
x_max, y_max = pts.max(axis=0)
|
|
pad_x, pad_y = (x_max - x_min) * margin_frac, (y_max - y_min) * margin_frac
|
|
x_min, x_max = x_min - pad_x, x_max + pad_x
|
|
y_min, y_max = y_min - pad_y, y_max + pad_y
|
|
n_cols = int(np.ceil((x_max - x_min) / dx)) + 1
|
|
n_rows = int(np.ceil((y_max - y_min) / abs(dy))) + 1
|
|
origin = (x_min, y_min if dy > 0 else y_max)
|
|
return origin, (n_rows, n_cols)
|
|
|
|
|
|
def correlate_translation_mm(sras: SrasFile, angle_idx: int, ref_angle_idx: int,
|
|
signal_mv: dict[int, np.ndarray],
|
|
pivot_mm: dict[int, tuple[float, float]], *,
|
|
use_mask: bool = False,
|
|
dc_threshold_mv: float = 0.0,
|
|
margin_frac: float = 0.3,
|
|
max_corr_dim: int = 1024) -> tuple[float, float]:
|
|
"""FFT phase-correlation translation (shift_x_mm, shift_y_mm) that best
|
|
lines up angle_idx's CH4 content onto ref_angle_idx's, given each angle's
|
|
own alignment pivot (pivot_mm — see compute_pivot_points_mm) already
|
|
rotated about that pivot by the analytic scan-angle delta with zero
|
|
shift. Returns (0.0, 0.0) unconditionally for the reference angle.
|
|
|
|
By default (use_mask=False) correlates on each angle's own raw CH4
|
|
signal minus its own minimum — the same weighting _signal_centroid_mm
|
|
uses — rather than a dc_threshold_mv binary mask: a shared threshold
|
|
that doesn't suit every angle's real signal level can empty or distort
|
|
one angle's mask (the same failure mode the alignment pivot was fixed
|
|
to avoid), and correlating on the raw signal also lets phase
|
|
correlation lock onto real internal sample structure rather than just
|
|
the scan window's silhouette. Subtracting each angle's own minimum
|
|
(rather than using the signal as-is) keeps the zero-padding surrounding
|
|
the rotated image from reading as a spurious high-contrast edge against
|
|
a nonzero DC baseline. use_mask=True switches to the binary
|
|
>= dc_threshold_mv mask instead (compute_angle_alignment's original
|
|
behavior), for cases where raw-signal correlation locks onto noise.
|
|
"""
|
|
if angle_idx == ref_angle_idx:
|
|
return (0.0, 0.0)
|
|
|
|
def corr_img(a: int) -> np.ndarray:
|
|
img = signal_mv[a]
|
|
if use_mask:
|
|
return (img >= dc_threshold_mv).astype(np.float32)
|
|
return (img - img.min()).astype(np.float32)
|
|
|
|
img_a = corr_img(angle_idx)
|
|
img_ref = corr_img(ref_angle_idx)
|
|
|
|
dx_ref, dy_ref = _pixel_pitch_mm(sras, ref_angle_idx)
|
|
max_dim = max(img_a.shape + img_ref.shape)
|
|
factor = max(1, int(np.ceil(max_dim / max_corr_dim)))
|
|
dx_c, dy_c = dx_ref * factor, dy_ref * factor
|
|
small_a = _block_mean_downsample(img_a, factor)
|
|
small_ref = _block_mean_downsample(img_ref, factor)
|
|
|
|
work_origin, work_shape = _working_canvas_for_pair(
|
|
sras, ref_angle_idx, angle_idx, dx_c, dy_c, pivot_mm, margin_frac=margin_frac)
|
|
|
|
m_a, o_a = _build_affine_canvas_to_raw(
|
|
sras, angle_idx, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin, pivot_mm)
|
|
m_ref, o_ref = _build_affine_canvas_to_raw(
|
|
sras, ref_angle_idx, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin, pivot_mm)
|
|
|
|
rotated_a = scipy_ndimage.affine_transform(
|
|
small_a, m_a / factor, offset=o_a / factor,
|
|
output_shape=work_shape, order=0, mode="constant", cval=0.0)
|
|
embedded_ref = scipy_ndimage.affine_transform(
|
|
small_ref, m_ref / factor, offset=o_ref / factor,
|
|
output_shape=work_shape, order=0, mode="constant", cval=0.0)
|
|
|
|
dr, dc = _phase_correlate_shift(embedded_ref, rotated_a)
|
|
return (dc * dx_c, dr * dy_c)
|
|
|
|
|
|
def _parallel_map(fn, items, n_workers: int) -> list:
|
|
"""fn over items, in order, threaded when it pays."""
|
|
items = list(items)
|
|
if n_workers <= 1 or len(items) <= 1:
|
|
return [fn(x) for x in items]
|
|
with ThreadPoolExecutor(max_workers=min(n_workers, len(items))) as pool:
|
|
return list(pool.map(fn, items))
|
|
|
|
|
|
def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int,
|
|
dc_threshold_mv: float,
|
|
progress_cb=None) -> AlignmentResult:
|
|
"""Top-level alignment driver. Runs on a background thread (see
|
|
AngleAlignmentWorker) — deliberately recomputes CH4 DC images from
|
|
scratch rather than reading the GUI-thread _dc_cache dict, since
|
|
background-thread workers must not touch GUI-thread-owned caches."""
|
|
n = sras.n_angles # already the *complete*-angle count for aborted v6 scans
|
|
dx_ref, dy_ref = _pixel_pitch_mm(sras, ref_angle_idx)
|
|
# Parallel over angles, serial within each — see plan_angle_level.
|
|
n_workers, angle_budget = plan_angle_level(sras)
|
|
|
|
# progress_cb fires from pool threads, so the counter behind it must be
|
|
# atomic. list.append is, and len() of a list is a consistent read.
|
|
_ticks: list[int] = []
|
|
|
|
def tick(base: int, span: int):
|
|
if progress_cb:
|
|
_ticks.append(1)
|
|
progress_cb(base + int(min(len(_ticks), n) / n * span))
|
|
|
|
# ---- Step 1: CH4 DC image + binarized mask per angle, native grid -----
|
|
def dc4_for(a: int) -> np.ndarray:
|
|
dc4 = adc_to_mv(
|
|
compute_dc_image(sras, a, CH4_IDX, max_workers=1, budget=angle_budget),
|
|
*sras.cal(CH4_IDX))
|
|
tick(0, 25)
|
|
return dc4
|
|
|
|
dc4_mv = dict(enumerate(_parallel_map(dc4_for, range(n), n_workers)))
|
|
|
|
# Each angle's own alignment pivot — the CH4-signal-weighted centroid,
|
|
# not the raw scan-window bbox center (see compute_pivot_points_mm).
|
|
# Reuses the dc4_mv images just computed above, so this is free, and is
|
|
# deliberately independent of dc_threshold_mv (see that function's
|
|
# docstring) so a threshold that happens to leave some angle's binary
|
|
# mask empty can't silently degrade the pivot back to the bbox center.
|
|
pivot_mm = compute_pivot_points_mm(sras, dc4_mv=dc4_mv)
|
|
|
|
# ---- Step 2: coarse translation via FFT phase correlation, on each
|
|
# angle's binarized CH4 mask against the reference's (see
|
|
# correlate_translation_mm — also the engine behind ManualAlignmentDialog's
|
|
# Auto Cross-Correlate button, there defaulting to the raw signal instead
|
|
# of a mask).
|
|
def shift_for(a: int) -> tuple[float, float]:
|
|
shift = correlate_translation_mm(
|
|
sras, a, ref_angle_idx, dc4_mv, pivot_mm,
|
|
use_mask=True, dc_threshold_mv=dc_threshold_mv)
|
|
tick(25, 50)
|
|
return shift
|
|
|
|
shifts_mm = dict(enumerate(_parallel_map(shift_for, range(n), n_workers)))
|
|
|
|
# ---- Step 3: union bounding box over all angles (rotation+shift applied)
|
|
corners = np.vstack([
|
|
_corners_in_ref_frame(sras, a, ref_angle_idx, pivot_mm, shifts_mm[a])
|
|
for a in range(n)])
|
|
x_min, y_min = corners.min(axis=0)
|
|
x_max, y_max = corners.max(axis=0)
|
|
n_cols = int(np.ceil((x_max - x_min) / dx_ref)) + 1
|
|
n_rows = int(np.ceil((y_max - y_min) / abs(dy_ref))) + 1
|
|
canvas_origin_mm = (float(x_min), float(y_min if dy_ref > 0 else y_max))
|
|
|
|
# ---- Step 4: final per-angle full-resolution affine (canvas -> raw idx)
|
|
per_angle: dict[int, AngleTransform] = {}
|
|
for a in range(n):
|
|
matrix, offset = _build_affine_canvas_to_raw(
|
|
sras, a, ref_angle_idx, shifts_mm[a], dx_ref, dy_ref, canvas_origin_mm,
|
|
pivot_mm)
|
|
per_angle[a] = AngleTransform(
|
|
_theta_deg(sras, a, ref_angle_idx), shifts_mm[a], matrix, offset)
|
|
if progress_cb:
|
|
progress_cb(75 + int((a + 1) / n * 25))
|
|
|
|
return AlignmentResult(ref_angle_idx, dc_threshold_mv, (n_rows, n_cols),
|
|
dx_ref, dy_ref, canvas_origin_mm, per_angle)
|
|
|
|
|
|
# Back-compat alias for the pre-split private name (used by tooling).
|
|
_compute_angle_alignment = compute_angle_alignment
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Manual alignment (Fusion menu -> Manual Alignment... dialog)
|
|
#
|
|
# Skips compute_angle_alignment's mask + phase-correlation search entirely:
|
|
# every angle's rotation_deg/shift_mm is supplied directly by the caller
|
|
# (nudged by eye against a live multi-angle mask overlay, or pre-seeded from
|
|
# a previous compute_angle_alignment run or a saved sidecar). Building the
|
|
# final AlignmentResult from already-known per-angle parameters is pure
|
|
# closed-form matrix math (union_canvas_mm + _build_affine_canvas_to_raw)
|
|
# with no per-pixel image work at all, so build_manual_alignment is cheap
|
|
# enough to call synchronously on the GUI thread on every edit. The only
|
|
# genuinely expensive per-pixel operation anywhere in this flow is
|
|
# reproject_mask, and only ManualAlignmentDialog's own downsampled preview
|
|
# calls that per keystroke — see that class's docstring for how it limits
|
|
# each nudge to reprojecting only the actively-edited angle.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def union_canvas_mm(sras: SrasFile, ref_angle_idx: int, dx: float, dy: float,
|
|
per_angle_params: dict[int, ManualAngleParams],
|
|
pivot_mm: dict[int, tuple[float, float]],
|
|
margin_frac: float = 0.0
|
|
) -> tuple[tuple[float, float], tuple[int, int]]:
|
|
"""Shared-canvas origin (mm) and (n_rows, n_cols) at pitch (dx, dy) that
|
|
contains every angle's footprint after applying its own rotation+shift —
|
|
the generalisation of compute_angle_alignment's Step 3 to arbitrary (not
|
|
just _theta_deg-analytic) per-angle rotation. Angles missing from
|
|
per_angle_params default to identity (e.g. a sidecar saved before a
|
|
rescan added more angles).
|
|
|
|
margin_frac pads the box on every side: 0 for a final canvas (this then
|
|
reproduces compute_angle_alignment's own Step-3 math exactly, when every
|
|
angle's rotation_deg equals the analytic delta and shift_mm matches);
|
|
nonzero for ManualAlignmentDialog's downsampled preview canvas, which
|
|
needs headroom so an ordinary translation nudge of the active angle never
|
|
has to trigger a full canvas resize (see that class's docstring — an
|
|
extreme nudge can still, in principle, push content past this padding;
|
|
accepted as a known edge case, same as _phase_correlate_shift's
|
|
wraparound risk note).
|
|
"""
|
|
n = sras.n_angles
|
|
corners = np.vstack([
|
|
_corners_in_ref_frame(
|
|
sras, a, ref_angle_idx, pivot_mm,
|
|
per_angle_params.get(a, ManualAngleParams()).shift_mm,
|
|
theta_deg=per_angle_params.get(a, ManualAngleParams()).rotation_deg)
|
|
for a in range(n)])
|
|
x_min, y_min = corners.min(axis=0)
|
|
x_max, y_max = corners.max(axis=0)
|
|
if margin_frac:
|
|
pad_x, pad_y = (x_max - x_min) * margin_frac, (y_max - y_min) * margin_frac
|
|
x_min, x_max = x_min - pad_x, x_max + pad_x
|
|
y_min, y_max = y_min - pad_y, y_max + pad_y
|
|
n_cols = int(np.ceil((x_max - x_min) / dx)) + 1
|
|
n_rows = int(np.ceil((y_max - y_min) / abs(dy))) + 1
|
|
origin = (float(x_min), float(y_min if dy > 0 else y_max))
|
|
return origin, (n_rows, n_cols)
|
|
|
|
|
|
def reproject_mask(sras: SrasFile, angle_idx: int, ref_angle_idx: int,
|
|
mask: np.ndarray, rotation_deg: float,
|
|
shift_mm: tuple[float, float],
|
|
canvas_dx: float, canvas_dy: float,
|
|
canvas_origin_mm: tuple[float, float],
|
|
canvas_shape: tuple[int, int],
|
|
pivot_mm: dict[int, tuple[float, float]]) -> np.ndarray:
|
|
"""Resample one angle's binary/float mask onto an arbitrary canvas via an
|
|
explicit rotation+shift — the single building block
|
|
ManualAlignmentDialog's live preview repeatedly calls (once per
|
|
keystroke, for only the actively-nudged angle), since it bypasses
|
|
compute_angle_alignment's phase-correlation search entirely and just
|
|
takes rotation_deg/shift_mm as given. order=0 (nearest) matches
|
|
apply_alignment's own reasoning: a binary mask must never be blended with
|
|
zero-padding.
|
|
"""
|
|
matrix, offset = _build_affine_canvas_to_raw(
|
|
sras, angle_idx, ref_angle_idx, shift_mm, canvas_dx, canvas_dy,
|
|
canvas_origin_mm, pivot_mm, theta_deg=rotation_deg)
|
|
return scipy_ndimage.affine_transform(
|
|
mask.astype(np.float32, copy=False), matrix, offset=offset,
|
|
output_shape=canvas_shape, order=0, mode="constant", cval=0.0)
|
|
|
|
|
|
def build_manual_alignment(sras: SrasFile, ref_angle_idx: int,
|
|
dc_threshold_mv: float,
|
|
per_angle_params: dict[int, ManualAngleParams],
|
|
pivot_mm: dict[int, tuple[float, float]] | None = None
|
|
) -> AlignmentResult:
|
|
"""Build a full, full-resolution AlignmentResult from user-supplied
|
|
per-angle rotation+shift — the Manual Alignment counterpart to
|
|
compute_angle_alignment, skipping its mask/phase-correlation search
|
|
entirely (every angle's transform here is exactly what the caller
|
|
supplied). Pure matrix/bbox math once pivot_mm is known, so it is cheap
|
|
enough to call synchronously on the GUI thread. The reference angle's
|
|
params are always forced to identity, regardless of what
|
|
per_angle_params holds for it — it defines the shared origin and must
|
|
never be transformed.
|
|
|
|
Pass an already-computed *pivot_mm* (e.g. ManualAlignmentDialog's own
|
|
cache, built once from its live CH4 images) to skip recomputing every
|
|
angle's DC image here; otherwise it's computed fresh via
|
|
compute_pivot_points_mm, which is still fine for a one-off Save or
|
|
sidecar restore (just not free on a very large, not-yet-cached scan).
|
|
dc_threshold_mv itself plays no part in the pivot (see that function's
|
|
docstring) — it's stored on the returned AlignmentResult purely as a
|
|
record of the RF-mask threshold in effect at the time.
|
|
"""
|
|
n = sras.n_angles
|
|
dx_ref, dy_ref = _pixel_pitch_mm(sras, ref_angle_idx)
|
|
params = {a: per_angle_params.get(a, ManualAngleParams()) for a in range(n)}
|
|
params[ref_angle_idx] = ManualAngleParams()
|
|
if pivot_mm is None:
|
|
pivot_mm = compute_pivot_points_mm(sras)
|
|
|
|
canvas_origin_mm, canvas_shape = union_canvas_mm(
|
|
sras, ref_angle_idx, dx_ref, dy_ref, params, pivot_mm, margin_frac=0.0)
|
|
|
|
per_angle: dict[int, AngleTransform] = {}
|
|
for a in range(n):
|
|
p = params[a]
|
|
matrix, offset = _build_affine_canvas_to_raw(
|
|
sras, a, ref_angle_idx, p.shift_mm, dx_ref, dy_ref,
|
|
canvas_origin_mm, pivot_mm, theta_deg=p.rotation_deg)
|
|
per_angle[a] = AngleTransform(p.rotation_deg, p.shift_mm, matrix, offset)
|
|
|
|
return AlignmentResult(ref_angle_idx, dc_threshold_mv, canvas_shape,
|
|
dx_ref, dy_ref, canvas_origin_mm, per_angle)
|
|
|
|
|
|
# ---- Sidecar persistence (<name>.sras.align.json) -------------------------
|
|
#
|
|
# Lives here, not sras_format.py: sras_format.py is scoped to the versioned
|
|
# binary .sras spec itself (see scan_format.md); a manual alignment is a
|
|
# viewer-computed *derived* artifact, analogous in kind to AlignmentResult —
|
|
# so it belongs with the alignment math it serialises, which already lives
|
|
# in this module. json + pathlib are both stdlib, so this doesn't add a new
|
|
# dependency to a module whose only load-bearing constraint is staying free
|
|
# of Qt/matplotlib for cheap multiprocessing-child imports.
|
|
|
|
@dataclass
|
|
class ManualAlignmentSidecar:
|
|
ref_angle_idx: int
|
|
dc_threshold_mv: float
|
|
per_angle: dict[int, ManualAngleParams]
|
|
|
|
|
|
def sidecar_path(sras_path) -> Path:
|
|
"""<name>.sras.align.json next to the scan file. A thin, independently
|
|
testable function since save/load/delete and the GUI's status messages
|
|
all need the identical path."""
|
|
p = Path(sras_path)
|
|
return p.with_name(p.name + ".align.json")
|
|
|
|
|
|
# Bumped from 1 -> 2 when the rotation pivot changed from the raw scan-
|
|
# window bbox center to a content-derived centroid, *and* _theta_deg's sign
|
|
# convention was corrected — either change alone makes a version-1 file's
|
|
# stored rotation_deg/shift_mm numbers describe a different (and, for the
|
|
# pivot bug, actively wrong) transform than they would today. Loading one
|
|
# unchanged would silently reproduce exactly the "scans show up everywhere"
|
|
# symptom these fixes address, so version-1 sidecars are treated as absent
|
|
# rather than migrated.
|
|
_SIDECAR_SCHEMA_VERSION = 2
|
|
|
|
|
|
def save_manual_alignment(sras: SrasFile, ref_angle_idx: int,
|
|
dc_threshold_mv: float,
|
|
per_angle: dict[int, ManualAngleParams]) -> Path:
|
|
"""Write the sidecar JSON for sras.path (overwriting any existing one)
|
|
and return the path written.
|
|
|
|
Schema (schema_version 2):
|
|
{
|
|
"schema_version": 2,
|
|
"ref_angle_idx": <int>,
|
|
"dc_threshold_mv": <float>,
|
|
"per_angle": {
|
|
"<angle_idx>": {"rotation_deg": <float>, "shift_mm": [<dx_mm>, <dy_mm>]},
|
|
...
|
|
}
|
|
}
|
|
Angle indices are JSON object keys, so they round-trip as strings —
|
|
load_manual_alignment converts them back to int.
|
|
"""
|
|
path = sidecar_path(sras.path)
|
|
payload = {
|
|
"schema_version": _SIDECAR_SCHEMA_VERSION,
|
|
"ref_angle_idx": ref_angle_idx,
|
|
"dc_threshold_mv": dc_threshold_mv,
|
|
"per_angle": {
|
|
str(a): {"rotation_deg": p.rotation_deg, "shift_mm": list(p.shift_mm)}
|
|
for a, p in per_angle.items()
|
|
},
|
|
}
|
|
path.write_text(json.dumps(payload, indent=2))
|
|
return path
|
|
|
|
|
|
def load_manual_alignment(sras: SrasFile) -> ManualAlignmentSidecar | None:
|
|
"""Read <sras.path>'s sidecar JSON if present, else None — never raises:
|
|
a missing, corrupt, foreign, or version-mismatched JSON file must not
|
|
block opening the .sras file itself (see SrasViewerWindow._on_load_done).
|
|
Per-angle entries for an angle index no longer present in *sras* (e.g.
|
|
re-scanned with fewer angles) are silently dropped.
|
|
"""
|
|
path = sidecar_path(sras.path)
|
|
if not path.exists():
|
|
return None
|
|
try:
|
|
raw = json.loads(path.read_text())
|
|
if raw.get("schema_version") != _SIDECAR_SCHEMA_VERSION:
|
|
return None
|
|
per_angle = {
|
|
int(a): ManualAngleParams(
|
|
float(v["rotation_deg"]),
|
|
(float(v["shift_mm"][0]), float(v["shift_mm"][1])))
|
|
for a, v in raw.get("per_angle", {}).items()
|
|
if int(a) < sras.n_angles
|
|
}
|
|
return ManualAlignmentSidecar(
|
|
ref_angle_idx=int(raw.get("ref_angle_idx", 0)),
|
|
dc_threshold_mv=float(raw.get("dc_threshold_mv", 0.0)),
|
|
per_angle=per_angle)
|
|
except (OSError, ValueError, KeyError, TypeError, IndexError,
|
|
json.JSONDecodeError):
|
|
return None
|
|
|
|
|
|
def delete_manual_alignment(sras: SrasFile) -> bool:
|
|
"""Delete the sidecar if present. Returns whether a file actually existed
|
|
to delete, so Clear Alignment's status message can say so. Genuine I/O
|
|
errors (permission denied, read-only share) propagate — the caller
|
|
(ManualAlignmentDialog._on_clear) surfaces them rather than silently
|
|
pretending the destructive action succeeded."""
|
|
path = sidecar_path(sras.path)
|
|
try:
|
|
path.unlink()
|
|
return True
|
|
except FileNotFoundError:
|
|
return False
|