#!/usr/bin/env python3 """ sras_average.py — Waveform-averaging utility for .sras files. Reduces memory footprint by coherently averaging every N consecutive frames along the acquisition axis, writing a new .sras file with n_frames / N frames. Usage: python sras_average.py input.sras output.sras --n 10 python sras_average.py input.sras output.sras --n 10 --discard-remainder Options: --n INT Number of frames to average into one (required). --discard-remainder Drop trailing frames that don't fill a complete group. Default: include a partial average for the last group. """ import argparse import shutil import struct import sys from pathlib import Path import numpy as np from sras_format import HDR_FMT, HDR_SIZE, SrasFile _SUPPORTED = (2, 3, 4) def parse_args(): p = argparse.ArgumentParser( description="Average every N waveforms in a .sras file and write a new file." ) p.add_argument("input", help="Input .sras file") p.add_argument("output", help="Output .sras file") p.add_argument("--n", type=int, required=True, metavar="N", help="Number of consecutive frames to average into one") p.add_argument("--discard-remainder", action="store_true", help="Drop trailing frames that don't fill a complete group of N") return p.parse_args() def read_header_sections(sras: SrasFile) -> bytes: """The raw bytes between the header and the waveform data (angle table, row table, preambles, background), copied through verbatim so nothing is lost in a re-encode.""" with open(sras.path, "rb") as f: f.seek(HDR_SIZE) return f.read(sras._data_offset - HDR_SIZE) def average_rows(block: np.ndarray, n: int, discard_remainder: bool) -> np.ndarray: """Average every N frames of one angle's (n_rows, n_ch, n_frames, spf) block. Returns int16 of shape (n_rows, n_ch, n_out, spf). Averaging is done in float32 and rounded on cast, matching numpy's mean followed by an int16 cast in the original implementation. """ n_frames = block.shape[2] n_full = n_frames // n remainder = n_frames % n parts = [] if n_full: full = block[:, :, :n_full * n, :].astype(np.float32) full = full.reshape(block.shape[0], block.shape[1], n_full, n, block.shape[3]) parts.append(full.mean(axis=3).astype(np.int16)) if remainder and not discard_remainder: tail = block[:, :, n_full * n:, :].astype(np.float32) parts.append(tail.mean(axis=2, keepdims=True).astype(np.int16)) if not parts: return np.empty((*block.shape[:2], 0, block.shape[3]), dtype=np.int16) return parts[0] if len(parts) == 1 else np.concatenate(parts, axis=2) def write_averaged(out_path: Path, sras: SrasFile, mid_sections: bytes, n: int, discard_remainder: bool) -> int: """Stream each angle through the averager, writing as we go so peak RAM stays at one angle's block rather than the whole file.""" bps = sras.bytes_per_sample n_frames_in = int(sras.n_frames[0]) n_out = n_frames_in // n if n_frames_in % n and not discard_remainder: n_out += 1 header = struct.pack( HDR_FMT, b"SRAS", sras.version, sras.n_angles, int(sras.n_rows[0]), float(sras.x_start_mm[0]), float(sras.x_delta_mm), sras.velocity_mm_s, sras.laser_freq_hz, n_out, # updated frame count sras.samples_per_frame, sras.sample_rate_hz, bps, sras.n_channels, ) with open(out_path, "wb") as f: f.write(header) f.write(mid_sections) for a in range(sras.n_angles): averaged = average_rows(sras.data[a], n, discard_remainder) if bps == 1: f.write(np.clip(averaged, -128, 127).astype(np.int8).tobytes()) else: f.write(averaged.astype(">i2").tobytes()) return n_out def main(): args = parse_args() if args.n < 1: print("Error: --n must be at least 1.", file=sys.stderr) sys.exit(1) in_path = Path(args.input) out_path = Path(args.output) if not in_path.exists(): print(f"Error: input file not found: {in_path}", file=sys.stderr) sys.exit(1) if out_path.resolve() == in_path.resolve(): print("Error: output path must differ from input path.", file=sys.stderr) sys.exit(1) print(f"Reading {in_path} ...", flush=True) sras = SrasFile(str(in_path)) if sras.version not in _SUPPORTED: print(f"Error: unsupported .sras version: {sras.version} " f"(this tool handles v{'/v'.join(map(str, _SUPPORTED))})", file=sys.stderr) sys.exit(1) n_frames_in = int(sras.n_frames[0]) print(f" Version : v{sras.version}") print(f" Angles : {sras.n_angles}") print(f" Rows : {int(sras.n_rows[0])}") print(f" Frames (actual): {n_frames_in}") print(f" Channels : {sras.n_channels}") print(f" Samples/frame : {sras.samples_per_frame}") print(f" Bytes/sample : {sras.bytes_per_sample}") if args.n == 1: print("--n 1: no averaging needed; copying file as-is.") shutil.copy2(in_path, out_path) print(f"Wrote {out_path}") return if args.n > n_frames_in: print(f"Warning: --n ({args.n}) exceeds available frames ({n_frames_in}). " "The entire dataset will be averaged into a single frame.") print(f"\nAveraging every {args.n} frames ...", flush=True) print(f"\nWriting {out_path} ...", flush=True) mid_sections = read_header_sections(sras) n_frames_out = write_averaged(out_path, sras, mid_sections, args.n, args.discard_remainder) n_full = n_frames_in // args.n remainder = n_frames_in % args.n if remainder and not args.discard_remainder: status = f"({n_full} full groups + 1 partial group of {remainder})" elif remainder: status = f"({n_full} full groups, {remainder} trailing frames discarded)" else: status = f"({n_full} full groups)" print(f" {n_frames_in} frames -> {n_frames_out} frames {status}") 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 ({out_mb / in_mb * 100:.1f}% of input)") print("Done.") if __name__ == "__main__": main()