#!/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()