original program
This commit is contained in:
Thomas Ales
+251
@@ -0,0 +1,251 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user