#!/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 sys import struct import argparse import numpy as np from pathlib import Path # --------------------------------------------------------------------------- # Header format — must match sras_viewer.py exactly # --------------------------------------------------------------------------- HDR_FMT = ">4sBHHffffIIdBB" HDR_SIZE = struct.calcsize(HDR_FMT) # 43 bytes 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_sras(path: Path): """Read all sections of a .sras file and return them as a dict.""" with open(path, "rb") as f: header_bytes = f.read(HDR_SIZE) fields = struct.unpack(HDR_FMT, header_bytes) (magic, ver, n_angles, n_rows, x_start, x_delta, vel, freq, n_frames_hdr, spf, sr, bps, n_ch) = fields if magic != b"SRAS": raise ValueError(f"Not a .sras file (bad magic: {magic!r})") if ver not in (2, 3, 4): raise ValueError(f"Unsupported .sras version: {ver}") with open(path, "rb") as f: f.seek(HDR_SIZE) angles = f.read(n_angles * 4) # big-endian float32 array, raw bytes y_pos = f.read(n_rows * 4) # big-endian float32 array, raw bytes preambles = [] # list of raw bytes (length-prefixed strings) if ver >= 3: for _ in range(n_ch): (length,) = struct.unpack(">H", f.read(2)) preambles.append(f.read(length)) background = b"" # raw bytes for the v4 background block if ver >= 4: (n_bg,) = struct.unpack(">I", f.read(4)) background = f.read(n_bg) raw = f.read() # all waveform data # ----------------------------------------------------------------------- # Determine actual frame count from file size (the header value can be # wrong — the viewer does the same correction) # ----------------------------------------------------------------------- total_samples = len(raw) // bps samples_per_pixel = n_ch * spf # samples in one (angle, row, frame) cell samples_per_full = n_angles * n_rows * samples_per_pixel actual_n_frames = total_samples // (n_angles * n_rows * samples_per_pixel) good_bytes = actual_n_frames * samples_per_full * bps # Decode waveform data dtype = np.int8 if bps == 1 else ">i2" data = np.frombuffer(raw[:good_bytes], dtype=dtype) data = data.reshape(n_angles, n_rows, n_ch, actual_n_frames, spf) # Work in int16 (safe intermediate for both int8 and int16 inputs) data = data.astype(np.int16) return { "ver": ver, "n_angles": n_angles, "n_rows": n_rows, "x_start": x_start, "x_delta": x_delta, "vel": vel, "freq": freq, "n_frames_hdr": n_frames_hdr, "spf": spf, "sr": sr, "bps": bps, "n_ch": n_ch, "angles_raw": angles, "y_pos_raw": y_pos, "preambles": preambles, "background": background, "data": data, # shape: (n_angles, n_rows, n_ch, n_frames, spf), int16 } def average_frames(data: np.ndarray, n: int, discard_remainder: bool) -> np.ndarray: """Average every N frames along axis 3. data shape: (n_angles, n_rows, n_ch, n_frames, spf) Returns array of shape (n_angles, n_rows, n_ch, n_out, spf). """ n_frames = data.shape[3] n_full = n_frames // n remainder = n_frames % n # Average full groups using reshape-trick (no Python loop) if n_full > 0: full = data[:, :, :, :n_full * n, :] # trim to full groups full = full.reshape(data.shape[0], data.shape[1], data.shape[2], n_full, n, data.shape[4]) # (..., n_out, n, spf) averaged = full.mean(axis=4).astype(np.int16) # (..., n_out, spf) else: averaged = np.empty((*data.shape[:3], 0, data.shape[4]), dtype=np.int16) if remainder > 0 and not discard_remainder: tail = data[:, :, :, n_full * n:, :] # shape (..., remainder, spf) tail_avg = tail.mean(axis=3, keepdims=True).astype(np.int16) averaged = np.concatenate([averaged, tail_avg], axis=3) return averaged def write_sras(path: Path, src: dict, data_out: np.ndarray): """Write a new .sras file with the averaged waveform data.""" ver = src["ver"] bps = src["bps"] n_out = data_out.shape[3] # Pack header — update only n_frames_hdr; everything else stays the same header = struct.pack( HDR_FMT, b"SRAS", ver, src["n_angles"], src["n_rows"], src["x_start"], src["x_delta"], src["vel"], src["freq"], n_out, # updated frame count src["spf"], src["sr"], bps, src["n_ch"], ) # Encode waveform data back to original dtype if bps == 1: raw_out = np.clip(data_out, -128, 127).astype(np.int8).tobytes() else: # big-endian int16 raw_out = data_out.astype(">i2").tobytes() with open(path, "wb") as f: f.write(header) f.write(src["angles_raw"]) f.write(src["y_pos_raw"]) if ver >= 3: for preamble_bytes in src["preambles"]: f.write(struct.pack(">H", len(preamble_bytes))) f.write(preamble_bytes) if ver >= 4: bg = src["background"] f.write(struct.pack(">I", len(bg))) f.write(bg) f.write(raw_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) src = read_sras(in_path) n_frames_in = src["data"].shape[3] print(f" Version : v{src['ver']}") print(f" Angles : {src['n_angles']}") print(f" Rows : {src['n_rows']}") print(f" Frames (actual): {n_frames_in}") print(f" Channels : {src['n_ch']}") print(f" Samples/frame : {src['spf']}") print(f" Bytes/sample : {src['bps']}") if args.n == 1: print("--n 1: no averaging needed; copying file as-is.") import shutil 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) data_out = average_frames(src["data"], args.n, args.discard_remainder) n_frames_out = data_out.shape[3] 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 and args.discard_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}") print(f"\nWriting {out_path} ...", flush=True) write_sras(out_path, src, data_out) 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()