Add scan-editing CLI and alignment tests; extend Manual Alignment correlation

Continues the Manual Alignment work: refines the FFT cross-correlation and
mask handling, adds sras_edit_scans.py (drop/renumber bad angle scans),
tools/test_alignment.py (registration ground-truth suite), and a rotating
test fixture in make_test_sras.py.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Thomas Ales
2026-08-06 09:35:30 -05:00
parent bbb075ff34
commit d5028db445
8 changed files with 1619 additions and 597 deletions
+223
View File
@@ -0,0 +1,223 @@
#!/usr/bin/env python3
"""Angle-alignment tests: does registration actually stack the scans?
Builds a synthetic scan in which one sample is imaged at several *known*
rotations and offsets (tools/make_test_sras.write_rotating) and checks that the
alignment path recovers them, that the shared canvas is angle 0's own pixel
grid extended, and that nothing in the result depends on any other angle's
stage coordinates.
No Qt — this exercises sras_compute directly. See tools/test_gui.py for the
dialog and Aligned-View plumbing.
Usage: python tools/test_alignment.py
"""
import sys
import tempfile
from pathlib import Path
import numpy as np
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
import sras_compute as compute # noqa: E402
from sras_format import CH4_IDX, SrasFile, adc_to_mv # noqa: E402
import tools.make_test_sras as gen # noqa: E402
# Registration is limited by how far a feature moves per degree: with this
# sample's ~1 mm radius and a ~16 µm registration pitch, a quarter degree is
# already sub-pixel, so it is the floor of what any metric can resolve here.
_ROT_TOL_DEG = 0.5
_SHIFT_TOL_MM = 0.02
_STACK_IOU_MIN = 0.90
_THRESHOLD_MV = 80.0
_failures: list[str] = []
def check(name: str, ok: bool, detail: str = ""):
print(f" {'PASS' if ok else 'FAIL'} {name}" + (f" — {detail}" if detail else ""))
if not ok:
_failures.append(name)
def dc4_images(sras: SrasFile) -> dict[int, np.ndarray]:
return {a: adc_to_mv(compute.compute_dc_image(sras, a, CH4_IDX), *sras.cal(CH4_IDX))
for a in range(sras.n_angles)}
def mm_transform(sras: SrasFile, result, angle_idx: int) -> np.ndarray:
"""Recover the pure mm-space rotation from a canvas->raw affine.
matrix == D @ R^T @ A_out, where A_out and D only carry the canvas and
per-angle pixel pitches; undoing both must leave something orthonormal, or
the transform is smuggling in a scale or a shear.
"""
dx_a, dy_a = compute._pixel_pitch_mm(sras, angle_idx)
A_out = np.array([[0.0, result.canvas_dx_mm], [result.canvas_dy_mm, 0.0]])
D = np.array([[0.0, 1.0 / dy_a], [1.0 / dx_a, 0.0]])
return np.linalg.inv(D) @ result.per_angle[angle_idx].matrix @ np.linalg.inv(A_out)
def main() -> int:
tmpdir = Path(tempfile.mkdtemp(prefix="sras_align_"))
path = tmpdir / "rotating.sras"
meta = gen.write_rotating(path, n_angles=5)
sras = SrasFile(str(path))
truth = meta["truth"]
print(f"\nrotating-sample scan: {sras.n_angles} angles, "
f"shapes {[sras.image_shape(a) for a in range(sras.n_angles)]}")
print("\nper-angle rigid registration (rotation + translation, no scale)")
dc4 = dc4_images(sras)
fits = {a: compute.register_angle_to_reference(
sras, a, 0, dc4, dc_threshold_mv=_THRESHOLD_MV)
for a in range(sras.n_angles)}
for a, fit in fits.items():
t_rot, t_shift = truth[a]
rot_err = abs(fit.rotation_deg - t_rot)
shift_err = float(np.hypot(fit.shift_mm[0] - t_shift[0],
fit.shift_mm[1] - t_shift[1]))
check(f"angle {a} rotation within {_ROT_TOL_DEG}° of truth",
rot_err <= _ROT_TOL_DEG,
f"got {fit.rotation_deg:.3f}°, truth {t_rot:.3f}° (err {rot_err:.3f}°)")
check(f"angle {a} translation within {_SHIFT_TOL_MM} mm of truth",
shift_err <= _SHIFT_TOL_MM, f"err {shift_err:.4f} mm")
check("reference angle registers as exact identity",
fits[0] == compute.RigidFit(0.0, (0.0, 0.0), 1.0, "reference"))
# The stage's rotational sense relative to this module's math-positive
# convention is not knowable from the file, and the old code hardcoded a
# guess. Flipping every reported angle must therefore change nothing: the
# search scores both signs and the images decide.
flipped = SrasFile(str(path))
flipped.angles_deg = -flipped.angles_deg
flipped_fits = {a: compute.register_angle_to_reference(
flipped, a, 0, dc4, dc_threshold_mv=_THRESHOLD_MV)
for a in range(1, flipped.n_angles)}
check("negating every reported stage angle changes no fit",
all(flipped_fits[a] == fits[a] for a in flipped_fits),
str({a: (flipped_fits[a].rotation_deg, fits[a].rotation_deg)
for a in flipped_fits if flipped_fits[a] != fits[a]}))
print("\nper-angle stage coordinates are not consulted")
# Move every non-reference angle's scan window somewhere else entirely.
# Only angle 0's coordinates may matter, so every fit must be untouched.
moved = SrasFile(str(path))
for a in range(1, moved.n_angles):
moved.x_start_mm[a] += 13.5 * a
moved._y_pos_per_angle[a] = moved._y_pos_per_angle[a] - 9.25 * a
moved_dc4 = dc4_images(moved)
moved_fits = {a: compute.register_angle_to_reference(
moved, a, 0, moved_dc4, dc_threshold_mv=_THRESHOLD_MV)
for a in range(1, moved.n_angles)}
check("relocating every other angle's scan window changes no fit",
all(moved_fits[a] == fits[a] for a in moved_fits),
str({a: (round(moved_fits[a].rotation_deg, 4), fits[a].rotation_deg)
for a in moved_fits if moved_fits[a] != fits[a]}))
print("\nshared canvas is angle 0's own pixel grid, extended")
result = compute.compute_angle_alignment(sras, 0, _THRESHOLD_MV)
t0 = result.per_angle[0]
check("angle 0's transform has no rotation, scale or shear",
np.allclose(t0.matrix, np.eye(2)), str(t0.matrix))
check("angle 0 lands on whole canvas pixels (no resampling of the reference)",
np.allclose(t0.offset, np.round(t0.offset)), str(t0.offset))
check("canvas pitch is angle 0's own pitch",
(result.canvas_dx_mm, result.canvas_dy_mm)
== compute._pixel_pitch_mm(sras, 0))
n_rows, n_cols = result.canvas_shape
x_axis = result.canvas_origin_mm[0] + np.arange(n_cols) * result.canvas_dx_mm
y_axis = result.canvas_origin_mm[1] + np.arange(n_rows) * result.canvas_dy_mm
row0, col0 = int(round(-t0.offset[0])), int(round(-t0.offset[1]))
a0_rows, a0_cols = sras.image_shape(0)
check("canvas X axis reproduces angle 0's own X coordinates",
np.allclose(x_axis[col0:col0 + a0_cols], sras.x_axis_mm(0)))
check("canvas Y axis reproduces angle 0's own Y coordinates",
np.allclose(y_axis[row0:row0 + a0_rows], sras.y_positions_mm(0)))
check("canvas covers every angle's footprint",
n_rows >= max(int(sras.n_rows[a]) for a in range(sras.n_angles))
and n_cols >= max(int(sras.n_frames[a]) for a in range(sras.n_angles)),
str(result.canvas_shape))
print("\nno scaling anywhere in the per-angle transforms")
for a in range(sras.n_angles):
R = mm_transform(sras, result, a)
check(f"angle {a}'s mm-space transform is a pure rotation",
np.allclose(R @ R.T, np.eye(2), atol=1e-9)
and abs(abs(np.linalg.det(R)) - 1.0) < 1e-9,
f"det={np.linalg.det(R):.6f}")
print("\nall angles stack into one shape")
aligned = {a: compute.apply_alignment(result, a, dc4[a])
for a in range(sras.n_angles)}
base = aligned[0] >= _THRESHOLD_MV
for a in range(1, sras.n_angles):
other = aligned[a] >= _THRESHOLD_MV
iou = float((base & other).sum()) / max(1, int((base | other).sum()))
check(f"angle {a}'s aligned sample overlaps angle 0's (IoU >= {_STACK_IOU_MIN})",
iou >= _STACK_IOU_MIN, f"IoU {iou:.4f}")
print("\ndownsampled preview lands where the full-resolution image does")
# ManualAlignmentDialog reprojects block-mean-downsampled masks, so the
# affine has to account for the factor. When it did not, every preview
# layer came out magnified by that factor and offset — the overlay showed a
# blown-up crop of each mask, which is not something you can align by eye.
pitch = (result.canvas_dx_mm, result.canvas_dy_mm)
a = sras.n_angles - 1
p = result.per_angle[a]
full_mask = (dc4[a] >= _THRESHOLD_MV).astype(np.float32)
full = compute.reproject_mask(
sras, a, 0, full_mask, p.rotation_deg, p.shift_mm, pitch,
result.canvas_origin_mm, result.canvas_shape)
fy, fx = 4, 16
small = compute.reproject_mask(
sras, a, 0, compute._block_mean_2d(full_mask, fy, fx),
p.rotation_deg, p.shift_mm, (pitch[0] * fx, pitch[1] * fy),
result.canvas_origin_mm,
(result.canvas_shape[0] // fy, result.canvas_shape[1] // fx),
src_downsample=(fy, fx))
# Compare in mm, via each layer's own center of mass.
def com_mm(layer, px, py):
rows, cols = np.nonzero(layer > 0.5)
return np.array([cols.mean() * px, rows.mean() * py])
d = com_mm(small, pitch[0] * fx, pitch[1] * fy) - com_mm(full, *pitch)
check("a downsampled preview layer lands within a preview pixel of the "
"full-resolution one",
abs(d[0]) <= abs(pitch[0] * fx) and abs(d[1]) <= abs(pitch[1] * fy),
f"offset {d[0]:+.4f}, {d[1]:+.4f} mm")
print("\nmanual path reproduces the same geometry")
params = {a: compute.ManualAngleParams(t.rotation_deg, t.shift_mm)
for a, t in result.per_angle.items()}
manual = compute.build_manual_alignment(sras, 0, _THRESHOLD_MV, params)
check("build_manual_alignment matches compute_angle_alignment for the same params",
manual.canvas_shape == result.canvas_shape
and np.allclose(manual.canvas_origin_mm, result.canvas_origin_mm)
and all(np.allclose(manual.per_angle[a].matrix, result.per_angle[a].matrix)
and np.allclose(manual.per_angle[a].offset, result.per_angle[a].offset)
for a in range(sras.n_angles)))
print("\nsidecar round-trip")
compute.save_manual_alignment(sras, 0, _THRESHOLD_MV, params)
loaded = compute.load_manual_alignment(sras)
check("sidecar reloads every angle's params",
loaded is not None
and all(np.isclose(loaded.per_angle[a].rotation_deg, params[a].rotation_deg)
and np.allclose(loaded.per_angle[a].shift_mm, params[a].shift_mm)
for a in range(sras.n_angles)))
check("sidecar deletes cleanly", compute.delete_manual_alignment(sras))
print()
if _failures:
print(f"{len(_failures)} FAILURE(S): " + ", ".join(_failures))
return 1
print("All alignment checks passed.")
return 0
if __name__ == "__main__":
sys.exit(main())