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>
296 lines
11 KiB
Python
296 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
sras_edit_scans.py — Remove one or more angle scans from a .sras file.
|
|
|
|
A .sras file holds one or more "angles" (rotation positions); the viewer
|
|
cross-correlates each non-reference angle against the reference to align
|
|
them. If one angle's acquisition went wrong (stage glitch, bad trigger,
|
|
laser dropout, ...) it throws off that alignment for the whole file. This
|
|
tool drops the bad angle(s) and renumbers the rest, writing a new .sras file
|
|
with everything else — waveform samples, calibration preambles, background
|
|
waveform, row/geometry tables — carried over byte-for-byte.
|
|
|
|
Handles v2-v7. Any precomputed FFT/DC cache (v5 PREC tail, v7 CACH tail) is
|
|
dropped on write, since it's indexed by angle and would be stale/misaligned
|
|
after renumbering; the viewer just recomputes it next time the file opens.
|
|
|
|
Usage:
|
|
python sras_edit_scans.py input.sras --list
|
|
python sras_edit_scans.py input.sras output.sras --drop 2,5
|
|
python sras_edit_scans.py input.sras output.sras --keep 0,1,3,4,6
|
|
"""
|
|
|
|
import argparse
|
|
import struct
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
from sras_format import (
|
|
GEO_FMT_V6, GEO_SIZE_V6, HDR_FMT, HDR_FMT_V6, HDR_SIZE, HDR_SIZE_V6,
|
|
SrasFile,
|
|
)
|
|
|
|
_LEGACY_VERSIONS = (2, 3, 4, 5)
|
|
_V6_VERSIONS = (6, 7)
|
|
|
|
|
|
def parse_args():
|
|
p = argparse.ArgumentParser(
|
|
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
p.add_argument("input", help="Input .sras file")
|
|
p.add_argument("output", nargs="?", help="Output .sras file (omit with --list)")
|
|
p.add_argument("--list", action="store_true",
|
|
help="Print each angle's index/degrees/geometry and exit")
|
|
g = p.add_mutually_exclusive_group()
|
|
g.add_argument("--drop", metavar="I,J,...",
|
|
help="Comma-separated angle indices to remove")
|
|
g.add_argument("--keep", metavar="I,J,...",
|
|
help="Comma-separated angle indices to keep (all others dropped)")
|
|
return p.parse_args()
|
|
|
|
|
|
def _parse_index_list(s: str, n_angles: int) -> set[int]:
|
|
out = set()
|
|
for piece in s.split(","):
|
|
piece = piece.strip()
|
|
if not piece:
|
|
continue
|
|
i = int(piece)
|
|
if not (0 <= i < n_angles):
|
|
raise ValueError(f"angle index {i} out of range [0, {n_angles - 1}]")
|
|
out.add(i)
|
|
return out
|
|
|
|
|
|
def print_listing(sras: SrasFile):
|
|
print(f"\n{'idx':>4} {'angle_deg':>10} {'x_start_mm':>11} {'rows':>6} {'frames':>7}")
|
|
for a in range(sras.n_angles):
|
|
print(f"{a:>4} {sras.angles_deg[a]:>10.4f} {sras.x_start_mm[a]:>11.4f} "
|
|
f"{int(sras.n_rows[a]):>6} {int(sras.n_frames[a]):>7}")
|
|
|
|
|
|
def _copy_range(fin, fout, offset: int, nbytes: int, chunk: int = 64 * 1024 * 1024):
|
|
"""Stream *nbytes* raw bytes from *fin* at *offset* into *fout*, without
|
|
ever holding more than one chunk in memory (waveform blocks can be
|
|
hundreds of MB to low GB each)."""
|
|
fin.seek(offset)
|
|
remaining = nbytes
|
|
while remaining:
|
|
buf = fin.read(min(chunk, remaining))
|
|
if not buf:
|
|
raise IOError("unexpected EOF while copying waveform data")
|
|
fout.write(buf)
|
|
remaining -= len(buf)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Legacy (v2-v5): uniform geometry across angles, one flat waveform block
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _write_legacy(sras: SrasFile, keep: list[int], out_path: Path):
|
|
n_rows = int(sras.n_rows[0])
|
|
n_frames = int(sras.n_frames[0]) # uniform across angles for v2-v5
|
|
n_ch = sras.n_channels
|
|
spf = sras.samples_per_frame
|
|
bps = sras.bytes_per_sample
|
|
|
|
header = struct.pack(
|
|
HDR_FMT, b"SRAS", sras.version, len(keep), n_rows,
|
|
float(sras.x_start_mm[0]), float(sras.x_delta_mm),
|
|
sras.velocity_mm_s, sras.laser_freq_hz,
|
|
n_frames, spf, sras.sample_rate_hz, bps, n_ch,
|
|
)
|
|
|
|
# Row table + preambles + background sit right after the angle table and
|
|
# don't vary per angle — copy that whole span through unmodified.
|
|
angle_table_size = sras.n_angles * 4
|
|
with open(sras.path, "rb") as f:
|
|
f.seek(HDR_SIZE + angle_table_size)
|
|
shared_mid = f.read(sras._data_offset - (HDR_SIZE + angle_table_size))
|
|
|
|
angle_bytes = n_rows * n_ch * n_frames * spf * bps
|
|
|
|
with open(sras.path, "rb") as fin, open(out_path, "wb") as fout:
|
|
fout.write(header)
|
|
fout.write(sras.angles_deg[keep].astype(">f4").tobytes())
|
|
fout.write(shared_mid)
|
|
for a in keep:
|
|
_copy_range(fin, fout, sras._data_offset + a * angle_bytes, angle_bytes)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# v6/v7: per-angle geometry, ragged waveform blocks
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _read_v6_sections(path: Path) -> dict:
|
|
"""Raw, low-level read of everything before the waveform data.
|
|
|
|
SrasFile._parse_v6 reads and discards each angle's x_delta (it's not
|
|
part of the display-facing geometry it exposes), so a round-trip through
|
|
SrasFile would silently drop that field. Re-parsing here keeps every
|
|
byte of the Per-Angle Geometry Table intact.
|
|
"""
|
|
with open(path, "rb") as f:
|
|
hdr_raw = f.read(HDR_SIZE_V6)
|
|
(magic, ver, n_angles_declared, x_start_nom, y_start_nom, x_delta_nom,
|
|
y_delta_nom, row_spacing, vel, freq, spf, sr, bps,
|
|
n_ch) = struct.unpack(HDR_FMT_V6, hdr_raw)
|
|
|
|
angles_deg = list(struct.unpack(f">{n_angles_declared}f",
|
|
f.read(n_angles_declared * 4)))
|
|
|
|
geo = [struct.unpack(GEO_FMT_V6, f.read(GEO_SIZE_V6))
|
|
for _ in range(n_angles_declared)] # (x_start, x_delta, n_frames, n_rows)
|
|
|
|
row_table = [f.read(geo[a][3] * 4) for a in range(n_angles_declared)]
|
|
|
|
preamble_start = f.tell()
|
|
for _ in range(n_ch):
|
|
(length,) = struct.unpack(">H", f.read(2))
|
|
f.read(length)
|
|
preambles_raw = _reread_span(f, preamble_start)
|
|
|
|
bg_start = f.tell()
|
|
(n_bg,) = struct.unpack(">I", f.read(4))
|
|
f.read(n_bg)
|
|
background_raw = _reread_span(f, bg_start)
|
|
|
|
data_offset = f.tell()
|
|
|
|
return {
|
|
"version": ver, "n_angles_declared": n_angles_declared,
|
|
"x_start_nom": x_start_nom, "y_start_nom": y_start_nom,
|
|
"x_delta_nom": x_delta_nom, "y_delta_nom": y_delta_nom,
|
|
"row_spacing": row_spacing, "vel": vel, "freq": freq,
|
|
"spf": spf, "sr": sr, "bps": bps, "n_ch": n_ch,
|
|
"angles_deg": angles_deg, "geo": geo, "row_table": row_table,
|
|
"preambles_raw": preambles_raw, "background_raw": background_raw,
|
|
"data_offset": data_offset,
|
|
}
|
|
|
|
|
|
def _reread_span(f, start: int) -> bytes:
|
|
end = f.tell()
|
|
f.seek(start)
|
|
span = f.read(end - start)
|
|
f.seek(end)
|
|
return span
|
|
|
|
|
|
def _v6_angle_offsets(sections: dict, file_size: int) -> list[tuple[int, int]]:
|
|
"""(offset, nbytes) of each declared angle's waveform block, stopping at
|
|
the first angle whose data isn't fully on disk (aborted scan)."""
|
|
n_ch, spf, bps = sections["n_ch"], sections["spf"], sections["bps"]
|
|
offsets = []
|
|
offset = sections["data_offset"]
|
|
for xs, xd, nf, nr in sections["geo"]:
|
|
nbytes = nr * n_ch * nf * spf * bps
|
|
if offset + nbytes > file_size:
|
|
break
|
|
offsets.append((offset, nbytes))
|
|
offset += nbytes
|
|
return offsets
|
|
|
|
|
|
def _write_v6(in_path: Path, sections: dict, keep: list[int], out_path: Path):
|
|
file_size = in_path.stat().st_size
|
|
offsets = _v6_angle_offsets(sections, file_size)
|
|
geo = sections["geo"]
|
|
|
|
header = struct.pack(
|
|
HDR_FMT_V6, b"SRAS", sections["version"], len(keep),
|
|
sections["x_start_nom"], sections["y_start_nom"],
|
|
sections["x_delta_nom"], sections["y_delta_nom"],
|
|
sections["row_spacing"], sections["vel"], sections["freq"],
|
|
sections["spf"], sections["sr"], sections["bps"], sections["n_ch"],
|
|
)
|
|
|
|
with open(in_path, "rb") as fin, open(out_path, "wb") as fout:
|
|
fout.write(header)
|
|
fout.write(struct.pack(f">{len(keep)}f",
|
|
*[sections["angles_deg"][i] for i in keep]))
|
|
for i in keep:
|
|
fout.write(struct.pack(GEO_FMT_V6, *geo[i]))
|
|
for i in keep:
|
|
fout.write(sections["row_table"][i])
|
|
fout.write(sections["preambles_raw"])
|
|
fout.write(sections["background_raw"])
|
|
for i in keep:
|
|
off, nbytes = offsets[i]
|
|
_copy_range(fin, fout, off, nbytes)
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
in_path = Path(args.input)
|
|
if not in_path.exists():
|
|
print(f"Error: input file not found: {in_path}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
print(f"Reading {in_path} ...", flush=True)
|
|
try:
|
|
sras = SrasFile(str(in_path))
|
|
except ValueError as e:
|
|
print(f"Error: {e}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
if sras.version not in (*_LEGACY_VERSIONS, *_V6_VERSIONS):
|
|
print(f"Error: unsupported .sras version: {sras.version}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
aborted_note = " (scan aborted; trailing angle(s) already excluded)" if sras.scan_aborted else ""
|
|
print(f" Version : v{sras.version}", flush=True)
|
|
print(f" Angles : {sras.n_angles}{aborted_note}", flush=True)
|
|
|
|
if args.list:
|
|
print_listing(sras)
|
|
return
|
|
|
|
if not args.output:
|
|
print("Error: output path required unless --list is given.", file=sys.stderr)
|
|
sys.exit(1)
|
|
if not (args.drop or args.keep):
|
|
print("Error: specify --drop or --keep (see --list for indices).", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
out_path = Path(args.output)
|
|
if out_path.resolve() == in_path.resolve():
|
|
print("Error: output path must differ from input path.", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
try:
|
|
if args.drop:
|
|
drop = _parse_index_list(args.drop, sras.n_angles)
|
|
keep = [a for a in range(sras.n_angles) if a not in drop]
|
|
else:
|
|
keep = sorted(_parse_index_list(args.keep, sras.n_angles))
|
|
except ValueError as e:
|
|
print(f"Error: {e}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
if not keep:
|
|
print("Error: at least one angle must remain.", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
dropped = [a for a in range(sras.n_angles) if a not in keep]
|
|
print(f"\nDropping angle(s): {dropped}")
|
|
print(f"Keeping angle(s) : {keep} ({len(keep)} of {sras.n_angles})")
|
|
print(f"\nWriting {out_path} ...", flush=True)
|
|
|
|
if sras.version in _LEGACY_VERSIONS:
|
|
_write_legacy(sras, keep, out_path)
|
|
else:
|
|
_write_v6(in_path, _read_v6_sections(in_path), keep, out_path)
|
|
|
|
in_mb = in_path.stat().st_size / 1024**2
|
|
out_mb = out_path.stat().st_size / 1024**2
|
|
print(f" Input size : {in_mb:.1f} MB")
|
|
print(f" Output size: {out_mb:.1f} MB")
|
|
print("Done.")
|
|
print("Note: any precomputed FFT/DC cache was dropped (it's indexed by "
|
|
"angle); the viewer will recompute it next time this file opens.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|