Furen Xiao
e8dfe0fe35
Update M6 manifest building to register T1c to CT with rigid Mattes MI on a CT-space head box, including multi-start initial tilts for large inter-session head motion. The inverse transform warps CT and GTV into T1c native space, with acceptance based on concentricity of the warped CT skull ring and MR head silhouette. Rerun preprocessing when a row has a label but the warped label file is missing. Normalize volume orientation in read_nii_img. Add registration benchmark documentation, benchmark script, and shared registration pipeline helpers.
118 lines
No EOL
5.1 KiB
Python
118 lines
No EOL
5.1 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")
|
|
def needs_proc(r):
|
|
if not os.path.exists(os.path.join(out_dir, r["key"] + ".nii.gz")):
|
|
return True
|
|
return bool(r.get("label")) and not os.path.exists(os.path.join(out_dir, r["key"] + "_label.nii.gz"))
|
|
|
|
todo = [(r["key"], r["img"], r.get("label"), out_dir, meta_dir) for r in rows if needs_proc(r)]
|
|
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() |