File size: 4,425 Bytes
b360028 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | """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 # noqa: E402
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: # a bad study must not kill the shard
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)
# Pool.imap's feeder queues without bound and each item is a ~150 MB zip,
# so gate it on completed results to keep the tar out of RAM.
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())
|