Add aligned/cropped .sras export and the compute pieces behind it

The alignment machinery never modified a scan: it produced an
AlignmentResult and every consumer resampled on the fly. That is right
for a viewer but means the aligned stack cannot leave the process, so no
other tool can read it and reopening the scan redoes the registration.

sras_align_export.write_aligned_sras bakes an alignment into a new v6
file: each angle is gathered onto the shared canvas by nearest neighbour,
so every output angle shares one grid and the file opens already aligned.
Registering the export against itself returns identity, which is the test
that pins the whole index chain.

Three details that are easy to get wrong and are now covered:

  * Rounding must be floor(x + 0.5), not np.rint. scipy's order=0 rounds
    halves away from zero, and the canvas is snapped to the reference's
    pixel grid, so exact halves are common rather than hypothetical.
  * Out-of-bounds must be tested on the fractional coordinate against
    [0, n-1], not on the rounded index, or a one-pixel rim gets real data
    everywhere the Aligned View shows padding.
  * ...but with a tolerance, because the mm-space affine chain lands an
    exactly-integer transform a few times 1e-13 off. A bare >= 0 drops
    the *reference* angle's entire first row and last column.

Padding is the per-channel ADC code nearest 0 mV, not zero: zero ADC
decodes to ~+100 mV on real calibration and would masquerade as sample.

Source rows are served from sliding in-RAM bands. A rotated angle maps
one output row to a diagonal across the source, so indexing a memmap in
output order refaults nearly the whole angle per row — terabytes of
paging for a gigabyte of data.

Also in compute: crop_alignment_result (a crop is pure index translation,
so it folds into the affine's offset rather than becoming a second
transform), overlap_stats, largest_rect_at_least for a crop that stays
inside the overlap region, and seed/sign/refine knobs on
register_angle_to_reference whose defaults leave existing behaviour and
tests byte-identical.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Thomas Ales
2026-08-07 22:34:18 -05:00
parent 3989c2a1b8
commit 6d0e30b9ce
8 changed files with 1389 additions and 36 deletions
+155 -16
View File
@@ -1235,13 +1235,22 @@ def default_max_workers() -> int:
def _rotation_candidates(nominal_deg: float, search_deg: float,
step_deg: float) -> list[float]:
"""Coarse rotation candidates: a window around *both* signs of the stage's
reported angle change. Scoring both is what makes the stage's sign
convention a non-issue — the images decide which way the stage turns, and
a file whose stage reports the opposite sense registers just as well."""
step_deg: float,
signs: tuple[int, ...] = (-1, 1)) -> list[float]:
"""Coarse rotation candidates: a window around each requested sign of the
stage's reported angle change. Scoring both signs (the default) is what
makes the stage's sign convention a non-issue — the images decide which way
the stage turns, and a file whose stage reports the opposite sense
registers just as well.
*signs* exists only so a user who has established which way their own stage
turns can halve the coarse sweep from the alignment wizard. It is an
override, not an inference: nothing in the file says which sign is right.
"""
out: list[float] = []
for center in (-nominal_deg, nominal_deg):
step_deg = max(abs(step_deg), 1e-9) # a 0 step would divide by zero below
for sign in signs:
center = sign * nominal_deg
k = int(np.floor(search_deg / step_deg))
for i in range(-k, k + 1):
out.append(center + i * step_deg)
@@ -1351,15 +1360,18 @@ def register_angle_to_reference(
dc_threshold_mv: float = 0.0,
sources: tuple[str, ...] = ("signal", "mask"),
coarse_dim: int = 256, fine_dim: int = _DEFAULT_FINE_DIM,
search_deg: float = 6.0, coarse_step_deg: float = 2.0) -> RigidFit:
search_deg: float = 6.0, coarse_step_deg: float = 2.0,
seed_deg: float | None = None,
seed_signs: tuple[int, ...] = (-1, 1),
refine: bool = True) -> RigidFit:
"""Rigid (rotation + translation, never scale) fit of *angle_idx* onto
*ref_angle_idx*, found entirely by cross-correlating image content.
Two stages:
1. Coarse sweep at *coarse_dim* over a ±*search_deg* window around both
signs of the stage's reported angle change (see _rotation_candidates),
for each requested source image, scored by _overlap_ncc. Only has to
pick the right basin.
1. Coarse sweep at *coarse_dim* over a ±*search_deg* window around the
requested signs of the stage's reported angle change (see
_rotation_candidates), for each requested source image, scored by
_overlap_ncc. Only has to pick the right basin.
2. Hill-climbing refinement of the winning (source, rotation) at
*fine_dim* with sub-pixel translation folded into every score
(_refine_rotation, _score_rotation), down to 0.05°.
@@ -1368,6 +1380,22 @@ def register_angle_to_reference(
local mm to ref mm (q = R @ l + shift); score is the final NCC, which the
caller can surface so a bad scan is visible rather than silently fused in.
Returns identity for the reference angle itself.
The last three arguments only exist to let the alignment wizard expose the
rotation search; every default reproduces the search this function has
always done.
*seed_deg* replaces the stage's reported angle change as the center of the
coarse sweep — the pre-rotation the search starts from. None (the default)
means the stage angle from the file's Angle Table, which is nearly always
what you want: it puts the sweep within a couple of degrees of the answer.
Pass 0.0 to search around no rotation at all, which is the honest choice for
a file whose stage angles are known to be wrong.
*refine=False* stops after the coarse sweep, leaving rotation on the
*coarse_step_deg* grid. Combined with search_deg=0.0 and a single seed sign
it pins rotation to exactly the seed and searches translation only — for a
scan whose stage angles are trusted more than its image content.
"""
if angle_idx == ref_angle_idx:
return RigidFit(0.0, (0.0, 0.0), 1.0, "reference")
@@ -1377,8 +1405,10 @@ def register_angle_to_reference(
if not candidates:
return RigidFit(0.0, (0.0, 0.0), -1.0, "none")
nominal = nominal_delta_deg(sras, angle_idx, ref_angle_idx)
thetas = _rotation_candidates(nominal, search_deg, coarse_step_deg)
nominal = (nominal_delta_deg(sras, angle_idx, ref_angle_idx)
if seed_deg is None else float(seed_deg))
thetas = _rotation_candidates(nominal, search_deg, coarse_step_deg,
seed_signs)
# ---- Stage 1: coarse sweep, every source ------------------------------
pitch_c, n_c = _reg_pitch_and_size(sras, coarse_dim)
@@ -1407,9 +1437,10 @@ def register_angle_to_reference(
theta = best[1]
score, shift = _score_rotation(ref_img, ref_valid, mov_reg, pitch_f, n_f,
theta, subpixel=True)
theta, score, shift = _refine_rotation(
ref_img, ref_valid, mov_reg, pitch_f, n_f, theta, score, shift,
step_deg=coarse_step_deg)
if refine:
theta, score, shift = _refine_rotation(
ref_img, ref_valid, mov_reg, pitch_f, n_f, theta, score, shift,
step_deg=coarse_step_deg)
dr, dc = shift
return RigidFit(float(theta), (float(dc * pitch_f), float(dr * pitch_f)),
@@ -1540,6 +1571,114 @@ def _result_from_params(sras: SrasFile, ref_angle_idx: int,
pitch[0], pitch[1], canvas_origin_mm, per_angle)
def crop_alignment_result(result: AlignmentResult, row0: int, col0: int,
n_rows: int, n_cols: int) -> AlignmentResult:
"""The same alignment restricted to a rectangular window of its canvas —
canvas pixel (row0, col0) becomes the cropped canvas's (0, 0).
Cropping folds into each angle's existing affine instead of becoming a
second transform, because a canvas crop is *pure index translation*. From
_affine_out_to_src, matrix = D @ Rinv @ A_out depends only on pitch and
rotation, and A_out @ [row0, col0] is exactly the mm displacement of the new
origin, so:
matrix @ [r', c'] + (offset + matrix @ [row0, col0])
== matrix @ [r' + row0, c' + col0] + offset
i.e. shifting the offset by matrix @ [row0, col0] reproduces the original
mapping at the shifted indices, exactly. That equality is what lets
apply_alignment, reproject_mask and the aligned .sras exporter all keep
working on a cropped result with no special-casing — and it is why the crop
preview a user approves is guaranteed to be the same pixels the exporter
writes.
Bounds are the caller's responsibility (the wizard's ROI page clamps to the
canvas): an out-of-range window is geometrically well-defined here and
simply resamples padding.
"""
if n_rows <= 0 or n_cols <= 0:
raise ValueError(f"empty crop: {n_rows} x {n_cols}")
delta = np.array([float(row0), float(col0)])
per_angle = {
a: AngleTransform(t.rotation_deg, t.shift_mm, t.matrix,
t.offset + t.matrix @ delta, t.score, t.source)
for a, t in result.per_angle.items()
}
origin = (result.canvas_origin_mm[0] + col0 * result.canvas_dx_mm,
result.canvas_origin_mm[1] + row0 * result.canvas_dy_mm)
return AlignmentResult(result.ref_angle_idx, result.dc_threshold_mv,
(int(n_rows), int(n_cols)),
result.canvas_dx_mm, result.canvas_dy_mm,
origin, per_angle)
def overlap_stats(counts: np.ndarray, n_angles: int) -> dict:
"""Summarize a per-pixel "how many angles cover this pixel" image — the
number the alignment wizard's mask-stack view is colored by.
Separated from the drawing code because it is the actual judgement the user
makes on that screen ("do the angles land on top of each other?"), and a
plain array-in/dict-out function can be tested without Qt.
"""
counts = np.asarray(counts)
union = int(np.count_nonzero(counts))
full = int(np.count_nonzero(counts >= n_angles))
return {
"union_px": union,
"full_px": full,
"full_frac": (full / union) if union else 0.0,
"mean_count": float(counts[counts > 0].mean()) if union else 0.0,
"max_count": int(counts.max()) if counts.size else 0,
"empty": union == 0,
}
def largest_rect_at_least(counts: np.ndarray, min_count: int
) -> tuple[int, int, int, int] | None:
"""Largest axis-aligned rectangle whose every pixel has counts >= min_count,
as (row0, col0, n_rows, n_cols), or None if no pixel qualifies.
Backs the wizard's "fit to full overlap" button. A *bounding box* of the
qualifying pixels would be the obvious thing and is wrong: the full-overlap
region of several rotated scans is roughly a disc, whose bounding box has
corners no angle covers at all. Offering that as the crop would hand the
user padding they explicitly asked to avoid, so this finds a rectangle that
is entirely inside the region.
Standard largest-rectangle-in-a-histogram sweep — O(rows * cols) on a
preview-sized array, so it is instant at interactive rates.
"""
good = np.asarray(counts) >= min_count
if not good.any():
return None
n_rows, n_cols = good.shape
best = (0, 0, 0, 0) # area, row0, col0, ...
best_area = 0
heights = np.zeros(n_cols, dtype=np.int64)
for r in range(n_rows):
heights = np.where(good[r], heights + 1, 0)
# Sentinel column of height 0 flushes the stack at the end of the row.
stack: list[tuple[int, int]] = [] # (start col, height)
for c in range(n_cols + 1):
h = int(heights[c]) if c < n_cols else 0
start = c
while stack and stack[-1][1] >= h:
s, sh = stack.pop()
area = sh * (c - s)
if area > best_area:
best_area = area
best = (s, sh, c - s, r)
start = s
if h:
stack.append((start, h))
col0, height, width, row_end = best
return (row_end - height + 1, col0, height, width)
def _parallel_map(fn, items, n_workers: int) -> list:
"""fn over items, in order, threaded when it pays."""
items = list(items)