Thomas Ales
d5028db445
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>
224 lines
10 KiB
Python
224 lines
10 KiB
Python
#!/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())
|