| """Stream one MR-RATE atlas batch tar -> one uint8 memmap shard + index CSV. |
| |
| One shard per batch rather than one file per study keeps the file count small. |
| |
| Written to --out: |
| <batch>.u8 raw uint8, shape (n_studies, 4, *geometry output shape), C order |
| <batch>.csv row,study_uid,has_t1w,has_t2w,has_flair,has_swi,n_series |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import csv |
| import io |
| import os |
| import sys |
| import tarfile |
| import time |
| from multiprocessing import Pool |
| from threading import Semaphore |
|
|
| import numpy as np |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from preprocess import MODALITIES, process_zip |
|
|
| GEOMETRY = "d15" |
|
|
|
|
| def _init(geometry): |
| global GEOMETRY |
| GEOMETRY = geometry |
|
|
|
|
| def _work(job): |
| uid, raw = job |
| try: |
| per, n = process_zip(io.BytesIO(raw), (GEOMETRY,)) |
| vol, presence = per[GEOMETRY] |
| return uid, vol, presence, n, "" |
| except Exception as exc: |
| return uid, None, None, 0, f"{type(exc).__name__}: {exc}" |
|
|
|
|
| def members(tar_path: str, wanted: set[str] | None): |
| """Yield (uid, zip_bytes) streaming forward through the tar exactly once.""" |
| with tarfile.open(tar_path, "r|") as tf: |
| for m in tf: |
| if not m.isfile() or not m.name.endswith(".zip"): |
| continue |
| uid = os.path.basename(m.name)[: -len(".zip")] |
| if uid.endswith("_atlas"): |
| uid = uid[: -len("_atlas")] |
| if wanted is not None and uid not in wanted: |
| continue |
| if m.size == 0: |
| continue |
| data = tf.extractfile(m).read() |
| yield uid, data |
|
|
|
|
| def main() -> int: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--tar", required=True) |
| ap.add_argument("--batch", required=True) |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--geometry", default="d15", choices=("d15", "d2")) |
| ap.add_argument("--workers", type=int, default=16) |
| ap.add_argument("--uid-list", default=None, help="optional file of study_uids to keep") |
| ap.add_argument("--limit", type=int, default=0) |
| args = ap.parse_args() |
|
|
| wanted = None |
| if args.uid_list: |
| with open(args.uid_list) as f: |
| wanted = {ln.strip() for ln in f if ln.strip()} |
|
|
| os.makedirs(args.out, exist_ok=True) |
| bin_path = os.path.join(args.out, f"{args.batch}.u8") |
| csv_path = os.path.join(args.out, f"{args.batch}.csv") |
| tmp_bin = bin_path + ".part" |
| tmp_csv = csv_path + ".part" |
|
|
| t0 = time.time() |
| written = 0 |
| failed = 0 |
| src = members(args.tar, wanted) |
| if args.limit: |
| src = (x for i, x in enumerate(src) if i < args.limit) |
|
|
| |
| |
| inflight = Semaphore(args.workers * 3) |
|
|
| def throttled(it): |
| for item in it: |
| inflight.acquire() |
| yield item |
|
|
| src = throttled(src) |
|
|
| with open(tmp_bin, "wb", buffering=1024 * 1024) as fbin, open(tmp_csv, "w", newline="") as fcsv: |
| w = csv.writer(fcsv) |
| w.writerow(["row", "study_uid", *[f"has_{m}" for m in MODALITIES], "n_series"]) |
| with Pool(args.workers, maxtasksperchild=64, initializer=_init, initargs=(args.geometry,)) as pool: |
| for uid, vol, presence, n, err in pool.imap(_work, src, chunksize=1): |
| inflight.release() |
| if vol is None or n == 0: |
| failed += 1 |
| if failed <= 20: |
| print(f"[skip] {uid} {err or 'no usable series'}", flush=True) |
| continue |
| fbin.write(vol.tobytes()) |
| w.writerow([written, uid, *presence.tolist(), n]) |
| written += 1 |
| if written % 200 == 0: |
| el = time.time() - t0 |
| print(f"[{args.batch}] {written} studies {el/60:.1f} min " |
| f"({written/max(el,1e-9):.2f}/s, {failed} skipped)", flush=True) |
|
|
| os.replace(tmp_bin, bin_path) |
| os.replace(tmp_csv, csv_path) |
| el = time.time() - t0 |
| print(f"[{args.batch}] DONE {written} studies, {failed} skipped, " |
| f"{el/60:.1f} min, {os.path.getsize(bin_path)/2**30:.1f} GiB", flush=True) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
