tka/sras-viewer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Split into four modules, deduplicate, and parallelize the compute paths

Browse Source 
Structure
  sras_format.py   parsing/writing, calibration, axes  (numpy + struct)
  sras_compute.py  DC/FFT images, alignment, batch cache  (+ scipy)
  sras_workers.py  Qt background workers
  sras_viewer.py   ROI, canvases, dialog, main window

format+compute import in 0.59s with no Qt or matplotlib (vs 2.96s for the
full app), which is what makes a spawn-based process pool worth using.

Deduplication
  - SrasFile.cal() and compute.dc_image_mv() replace the ADC->mV
    calibration incantation that appeared at ten call sites.
  - _read_preambles/_read_background/_set_calibration are shared by the
    legacy and v6 parsers instead of duplicated.
  - v5 PREC images are normalised into the same ragged per-angle list
    v6/v7 uses, removing four isinstance() branches and shortening
    compute_rf_image's fast path.
  - _run_worker replaces five copies of the QThread setup and five
    near-identical teardown methods; the two lifetime hazards they
    guarded against are now documented once, authoritatively.
  - _on_channel_changed defers to _update_controls_enabled rather than
    re-deriving the same six enable rules.
  - _corners_in_ref_frame, _CHANNEL_DISPLAY dict, unified progress-dialog
    helper, dead _on_roi_angle_edited stub removed.
  - sras_average.py builds on SrasFile instead of carrying a second copy
    of the header format, and streams per angle instead of loading the
    whole file (was a ~3x file-size peak).
  Executable lines: 2390 -> 2254, despite adding all of the below.

Parallelism
  - compute_rf_image/compute_dc_image map row chunks over a thread pool.
    Worker count is derived from the memory budget rather than the core
    count: on a large scan chunk_rows is already floored at 1 row (~75 MB
    of float32 at 7507x2500), so only the worker count can bound peak RAM.
  - DcPrecomputeWorker computes angles on a pool, emitting each result
    from its own QThread as it lands.
  - BatchCacheWorker runs one process per file via cache_file(); only
    paths and scalars cross the boundary. Per-process thread counts are
    divided so the two levels don't oversubscribe. Drops the old pass 1,
    which fully parsed every file just to weight a progress bar.
  - dc_threshold_mv=None means "no mask" and skips the CH4 read entirely
    — the batch FFT job previously read all of CH4 to compare against a
    threshold of -1e9.
  - Alignment mask and correlation stages map over angles; fft2/ifft2 use
    workers=-1.

Fixes found on the way
  - A partially-cached v5 file showed an all-zero image for uncached
    angles: the fast-path check was file-wide, not per-angle.
  - v7 cached images were read-only big-endian views; now native float32.

Verified: tools/check_equivalence.py produces byte-identical hashes for
167 outputs (DC/FFT across channels, angles, bg-sub, pad factors and
thresholds, plus the full alignment result) against ed0eba4, on both
synthetic files and a real 496 GB 17-angle v6 scan.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Thomas Ales
2026-07-30 23:38:08 -05:00
parent 95c273dce5
commit 55a4c5e42f
5 changed files with 2218 additions and 2275 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sras_average.py
+78 -153
View File
@@ -15,25 +15,24 @@ Options:
Default: include a partial average for the last group.
"""
import sys
import struct
import argparse
import numpy as np
import shutil
import struct
import sys
from pathlib import Path
# ---------------------------------------------------------------------------
# Header format — must match sras_viewer.py exactly
# ---------------------------------------------------------------------------
import numpy as np
HDR_FMT = ">4sBHHffffIIdBB"
HDR_SIZE = struct.calcsize(HDR_FMT) # 43 bytes
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("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")
@@ -42,143 +41,68 @@ def parse_args():
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:
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)
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
}
return f.read(sras._data_offset - HDR_SIZE)
def average_frames(data: np.ndarray, n: int, discard_remainder: bool) -> np.ndarray:
"""Average every N frames along axis 3.
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).
data shape: (n_angles, n_rows, n_ch, n_frames, spf)
Returns array of shape (n_angles, 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 = data.shape[3]
n_full = n_frames // n
n_frames = block.shape[2]
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)
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 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
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_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]
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
# 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"],
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,
)
# 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:
with open(out_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)
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():
@@ -188,32 +112,35 @@ def main():
print("Error: --n must be at least 1.", file=sys.stderr)
sys.exit(1)
in_path = Path(args.input)
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)
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 = src["data"].shape[3]
print(f" Version : v{src['ver']}")
print(f" Angles : {src['n_angles']}")
print(f" Rows : {src['n_rows']}")
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 : {src['n_ch']}")
print(f" Samples/frame : {src['spf']}")
print(f" Bytes/sample : {src['bps']}")
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.")
import shutil
shutil.copy2(in_path, out_path)
print(f"Wrote {out_path}")
return
@@ -223,27 +150,25 @@ def main():
"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]
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
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:
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}")
print(f"\nWriting {out_path} ...", flush=True)
write_sras(out_path, src, data_out)
in_mb = in_path.stat().st_size / 1024**2
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(f" Output size: {out_mb:.1f} MB ({out_mb / in_mb * 100:.1f}% of input)")
print("Done.")