Furen Xiao
b6fa62a763
Add .gitignore, AGENTS.md, scripts directory, and src directory to initialize the repository.
114 lines
No EOL
5 KiB
Python
114 lines
No EOL
5 KiB
Python
"""Preprocess source volumes (T1c + optional label) into a uniform 1mm cropped/normalized form.
|
|
|
|
Usage:
|
|
python scripts/preprocess.py --manifest data/manifests/train_labeled.jsonl --workers 16
|
|
Each row: {key, subject, date, img, label?, source}
|
|
Outputs: data/proc/<key>.nii.gz, data/proc/<key>_label.nii.gz, data/procmeta/<key>.json
|
|
"""
|
|
import os
|
|
import sys
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
import json
|
|
import argparse
|
|
import numpy as np
|
|
import SimpleITK as sitk
|
|
from concurrent.futures import ProcessPoolExecutor, as_completed
|
|
from src.common import ROOT, d, read_nii_img, head_mask_from_image, crop_box, normalize_volume, write_arr, save_jsonl
|
|
|
|
SPACING = (1.0, 1.0, 1.0)
|
|
|
|
|
|
def resample_volume(itk, spacing, interp):
|
|
size = [max(1, int(round(s * sp / nsp))) for s, sp, nsp in zip(itk.GetSize(), itk.GetSpacing(), spacing)]
|
|
ptype = sitk.sitkFloat32 if interp != sitk.sitkNearestNeighbor else itk.GetPixelID()
|
|
return sitk.Resample(itk, size, sitk.Transform(), interp,
|
|
itk.GetOrigin(),
|
|
tuple(float(s) for s in spacing),
|
|
itk.GetDirection(),
|
|
0.0, ptype)
|
|
|
|
|
|
def worker(args):
|
|
key, img_path, label_path, out_dir, meta_dir = args
|
|
try:
|
|
img_itk = read_nii_img(img_path)
|
|
img1 = resample_volume(img_itk, SPACING, sitk.sitkLinear)
|
|
a = sitk.GetArrayFromImage(img1)
|
|
label1 = None
|
|
if label_path and os.path.exists(label_path):
|
|
lbl = read_nii_img(label_path)
|
|
label1 = sitk.GetArrayFromImage(resample_volume(lbl, SPACING, sitk.sitkNearestNeighbor)).astype(np.uint8)
|
|
label1 = (label1 > 0).astype(np.uint8)
|
|
# ensure label1 grid == a grid
|
|
if label1 is not None and label1.shape != a.shape:
|
|
ref = sitk.GetImageFromArray(np.zeros(a.shape, dtype=np.uint8))
|
|
ref.CopyInformation(img1)
|
|
lbl2 = resample_volume(read_nii_img(label_path), SPACING, sitk.sitkNearestNeighbor)
|
|
label1 = sitk.GetArrayFromImage(sitk.Resample(lbl2, ref, sitk.Transform(), sitk.sitkNearestNeighbor)).astype(np.uint8)
|
|
label1 = (label1 > 0).astype(np.uint8)
|
|
if label1.shape != a.shape:
|
|
raise RuntimeError(f"label/image grid mismatch {label1.shape} vs {a.shape}")
|
|
if a.max() <= a.min() + 1e-6:
|
|
raise RuntimeError("empty volume")
|
|
box = crop_box(a, label=label1, margin=12, cap=216, spacing=1.0)
|
|
(s0, s1), (s2, s3), (s4, s5) = box
|
|
ac_arr = a[s0:s1, s2:s3, s4:s5]
|
|
lab_arr = label1[s0:s1, s2:s3, s4:s5] if label1 is not None else None
|
|
m = head_mask_from_image(ac_arr)
|
|
if m.sum() < 5000:
|
|
raise RuntimeError("head mask too small")
|
|
norm, (lo, hi) = normalize_volume(ac_arr, m)
|
|
norm = norm.astype(np.float32)
|
|
out = os.path.join(out_dir, key + ".nii.gz")
|
|
write_arr(norm, out)
|
|
meta = {"key": key, "img_spacing_in": list(img_itk.GetSpacing()), "img_size_in": list(img_itk.GetSize()),
|
|
"origin": list(img_itk.GetOrigin()), "direction": list(img_itk.GetDirection()),
|
|
"crop_vox": [int(s0), int(s2), int(s4)],
|
|
"norm_lo": float(lo), "norm_hi": float(hi),
|
|
"shape": list(norm.shape)}
|
|
with open(os.path.join(meta_dir, key + ".json"), "w") as f:
|
|
json.dump(meta, f)
|
|
if lab_arr is not None:
|
|
outl = os.path.join(out_dir, key + "_label.nii.gz")
|
|
write_arr(lab_arr.astype(np.uint8), outl)
|
|
meta["n_tumor_vox"] = int(lab_arr.sum())
|
|
with open(os.path.join(meta_dir, key + ".json"), "w") as f:
|
|
json.dump(meta, f)
|
|
return key, None
|
|
except Exception as e: # noqa
|
|
return key, repr(e)
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--manifest", required=True)
|
|
ap.add_argument("--workers", type=int, default=16)
|
|
ap.add_argument("--subset", type=int, default=0, help="only first N rows (0=all)")
|
|
args = ap.parse_args()
|
|
rows = [json.loads(l) for l in open(args.manifest) if l.strip()]
|
|
if args.subset:
|
|
rows = rows[: args.subset]
|
|
out_dir = d("data/proc")
|
|
meta_dir = d("data/procmeta")
|
|
todo = [(r["key"], r["img"], r.get("label"), out_dir, meta_dir) for r in rows
|
|
if not os.path.exists(os.path.join(out_dir, r["key"] + ".nii.gz"))]
|
|
print(f"rows={len(rows)} todo={len(todo)}")
|
|
errs = []
|
|
with ProcessPoolExecutor(max_workers=args.workers) as ex:
|
|
futs = {ex.submit(worker, t): t[0] for t in todo}
|
|
done = 0
|
|
for fu in as_completed(futs):
|
|
k, err = fu.result()
|
|
done += 1
|
|
if err:
|
|
errs.append((k, err))
|
|
print("ERR", k, err, flush=True)
|
|
if done % 50 == 0:
|
|
print(f" {done}/{len(todo)} done, {len(errs)} errs", flush=True)
|
|
print(f"finished. errors: {len(errs)}")
|
|
for k, e in errs[:20]:
|
|
print(" ", k, e)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |