MRI2CT

Imported from minsuk00/MRI2CT (AGENTS.md). Install upstream with npx skills add minsuk00/MRI2CT. Copyright stays with the author.

CLAUDE.md

This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.

Project Overview

MRI2CT synthesizes CT images from MRI scans using deep learning. The main approach uses a U-Net translator with a frozen pre-trained Anatomix feature extractor for perceptual/segmentation losses. Baselines include a standard U-Net, a MAISI diffusion ControlNet, a BabyUNet segmentation teacher, and the KoalAI SynthRAD2025 winner (baselines/koalAI/, nnsyn fork — see baselines/koalAI/TODO.md).

Environment

micromamba activate mrct      # main repo + amix/unet/maisi/baby_unet + TotalSegmentator
micromamba activate koalai    # nnsyn fork (only for KoalAI baseline: convert/preprocess/train/predict)
micromamba activate pyradplan # pyRadPlan dose-eval downstream task ONLY (evaluation/dosimetry/ dir)

The two main envs are kept separate because both register a package named nnunetv2: mrct has upstream 2.5.2 (required by TotalSegmentator), koalai has the nnsyn fork. Don't pip-install across them.

pyradplan is a third, isolated env for the pyRadPlan dose-evaluation downstream task (evaluation/dosimetry/). It MUST stay separate: pyRadPlan requires numpy>=2 (and drags in numba/numpydantic that force-upgrade numpy), which breaks mrct's numpy-1.26 world (scikit-image / tptbox). Never pip install pyRadPlan into mrct. See evaluation/dosimetry/README.md. (The old pyradplan/ scratch dir is archived at _archive/pyradplan/.)

xvr + diffdrr are installed into mrct (for evaluation/2d3d_registration/). Use the dev branch, not main: main depends on antspyx>=0.5.4, which would clash with TPTBox==0.3.0's antspyx==0.4.2 pin (TPTBox is required by CADS) — the dev branch dropped antspyx entirely (it reads transforms via SimpleITK instead), so there is no conflict and pip check stays clean. dev is also a much better API for us: Register is a keyword-arg attrs class with a pluggable pose initializer (FixedPose) and metric, a real multiscale pyramid, and a 9-line Pose module.

pip install -c <(echo numpy==1.26.4) "diffdrr>=0.6.0"          # pin numpy or pip upgrades 1.26.4 -> 2.x
pip install --no-deps "git+https://github.com/eigenvivek/xvr.git@dev"

--no-deps on xvr just skips pydicom's optional compressed-DICOM decoders (we render DRRs, never read DICOM). Two pre-existing pip check lines (anatomix's torch==2.8.0 pin, opencv-python-headless's numpy>=2 pin) are unrelated and benign. Rollback snapshot: ~/mrct_freeze_pre_xvr.txt.

Data lives on GPFS: /gpfs/accounts/jjparkcv_root/jjparkcv98/minsukc/MRI2CT/SynthRAD/ (= dataset/). The training root is 1.5mm_registered_flat_totalseg_masked/ (src/common/config.py:21) — CT/MR masked with the tight TotalSegmentator body mask. The older loose-masked 1.5mm_registered_flat_masked/ it was derived from is archived at dataset/_archive/. It is archived, not dead — do not delete it. No training data is read from it any more (the totalseg root owns its own files), but six live scripts still read it as a SOURCE: the retired-but-selectable cads35 / v2_35 teacher presets (src/unet_baseline/train.py:760, src/seg_baby_unet/train.py:108) and the four preprocessing builders that constructed the current root from it (src/preprocess/{make_totalseg_masked, make_12cls_masked_labels,run_cads_merge_all,apply_body_mask}.py).

Linting: ruff check src/ (E402 is globally ignored via pyproject.toml).

Running Training

# Anatomix translator -- ARCHIVED to src/_archive/amix/, no longer the main model (superseded by
# the plain U-Net below; kept for reference, not verified runnable from the archived location)

# UNet baseline (current default/main model)
python src/unet_baseline/train.py

# MAISI diffusion baseline -- ARCHIVED to src/_archive/maisi_baseline/ (see evaluation/generate/
# generate_maisi.py for the still-live MAISI inference path used by evaluation/)

# BabyUNet teacher (12-class CT organ segmentation)
python src/seg_baby_unet/train.py

# Via SLURM (A40 GPU)
sbatch sbatch/train_unet.sh   # ls sbatch/ — names churn. Scripts for archived trainers (amix,
                              # maisi) and finished one-off studies live in sbatch/_archive/.

KoalAI baseline (separate koalai env, per-region models, dual-fold)

Two splits live in one preprocessed dataset as separate folds — fold 0 = center-wise (OOD test), fold 1 = random (i.i.d.). Trained seg model is the perceptual extractor for synth's MAP loss; synth needs seg of the same fold to be done first. Full workflow in baselines/koalAI/TODO.md.

# Pilot: thorax fold 0 (centerwise) seg → synth → val inference
REGION=thorax                       sbatch sbatch/koalai_train_seg.sh
REGION=thorax                       sbatch sbatch/koalai_train_synth.sh   # after seg done
REGION=thorax SUBSET=val            sbatch sbatch/koalai_predict.sh       # SUBSET=val|test, default test

# Fold 1 (random) for any step: add FOLD=1
REGION=thorax FOLD=1                sbatch sbatch/koalai_train_seg.sh

Running Evaluation

src/evaluate/ (the old scratch-tier eval scripts) is archived at src/_archive/evaluate/. The current, curated pipeline lives in evaluation/: sCT volumes are produced once per model into a shared evaluation/volumes/<dataset>/<split>/<arm>/ tree (symlinked to GPFS), then every downstream task reads from there. See evaluation/README.md for the full produce-once/reuse-across-tasks design.

# 1. Generate an sCT arm from a checkpoint (mrct env)
python evaluation/generate/generate.py \
    --checkpoint /path/to/checkpoint.pt \
    --out_dir evaluation/volumes/synthrad/center_wise_val/<arm_name>
# koalAI / MAISI checkpoints use the family-specific generator (same contract):
#   evaluation/generate/generate_koalai.py (koalai env) · evaluation/generate/generate_maisi.py

# 2. Image metrics (MAE/PSNR/SSIM/Bone Dice/Dice) -- auto-creates the aux reference arms it needs
python evaluation/image_metrics/run.py --dataset synthrad --split center_wise_val
python evaluation/image_metrics/viz.py --dataset synthrad --split center_wise_val

# 3. Dosimetry: does the sCT deliver the same RT dose as the real planning CT? (pyradplan env)
micromamba run -n pyradplan python evaluation/dosimetry/run.py \
    --dataset goldatlas --split prostate_photon --arm real_ct <model_arm>
micromamba run -n pyradplan python evaluation/dosimetry/viz.py \
    --dataset goldatlas --split prostate_photon --sct_model <model_arm>

# 4. 2D/3D X-ray registration downstream task
#    NOTE: its old default split `center_wise_val_deprecated` (loose-masked volumes) is retired
#    and archived to evaluation/volumes/_archive/. Re-run against `center_wise_val`.
python evaluation/2d3d_registration/run.py --dataset synthrad --split center_wise_val

Metric/orientation landmines that have silently corrupted past analyses (body-MAE air dilution, RAS/LPS mismatch, hardcoded sigmoid head, two "bone" definitions, soft-vs-hard Dice): see docs/eval_gotchas.md before any cross-model or cross-region comparison.

Architecture

Source Structure

• src/common/ — shared modules used by all trainers:
  • config.py: DEFAULT_CONFIG dict with all hyperparameters; Config (SimpleNamespace wrapper)
  • data.py: subject-discovery helpers (get_subject_paths, get_split_subjects, get_region_key) + MONAI 3-stage pipeline (get_cached_transforms, get_random_crop, get_gpu_transforms, gpu_augment_batch, default_monai_cache_dir, build_data_dicts). Used by all trainers and eval scripts. The legacy torchio version is preserved at src/_archive/data_torchio.py for reference.
  • loss.py: CompositeLoss (L1 + SSIM + optional Anatomix perceptual + optional teacher Dice + optional bone Dice). When the perceptual loss is on (perceptual_w > 0), set ssim_w = 0 — perceptual replaces SSIM as the structural-similarity term on top of L1; using both double-counts structure. Dice loss is a per-class weighted macro-mean over all classes incl. background (bone class uses dice_bone_w, others dice_w).
  • trainer_base.py: BaseTrainer (seeding, WandB init, checkpoint resume, _log_monitoring() for RAM/VRAM/timings)
  • utils.py: anatomix_normalize(), compute_metrics(), cleanup_gpu(), send_notification(), get_ram_info() (parent + recursive children PSS via psutil)
• src/amix/ — anatomix translator trainer — ARCHIVED to src/_archive/amix/
• src/unet_baseline/ — plain U-Net baseline (current default/main model)
• src/maisi_baseline/ — MAISI diffusion ControlNet baseline — ARCHIVED to src/_archive/maisi_baseline/
• src/seg_baby_unet/ — teacher segmentation network
• src/evaluate/ — NIfTI export and metric evaluation scripts — ARCHIVED to src/_archive/evaluate/; current eval pipeline is evaluation/ (see Running Evaluation above)
• src/preprocess/ — MRI-CT registration, resampling, segmentation pipelines
• anatomix/ — local package with pre-trained U-Net feature extractor (versions v1, v1_2, v1_3)
• splits/ — train/val split files (SPLIT_NAME SUBJECT_ID per line); splits/koalai/ holds per-region derivations of both center_wise and random splits
• baselines/koalAI/ — cloned nnsyn fork + our mri2ct_scripts/ integration helpers (only the helpers + wandb hooks are our code; rest is upstream)
• baselines/koalAI-seg/ — git worktree of koalAI on the nnunetv2 branch (used by seg training; same wandb-hook patch applied)

Training Data Flow

MONAI 3-stage pipeline (src/common/data.py):

1. Cached CPU preprocessing (get_cached_transforms → monai.data.PersistentDataset): load NIfTIs → enforce RAS → CT clip-and-scale (e.g. -1024..1024 → 0..1) → MRI minmax or percentile (0–99.5) → snapshot pre-pad shape → pad-end to ≥patch_size and to multiple of res_mult → uint8 mask hygiene. Cached at /tmp/mri2ct_<USER>_monai_cache (single shared dir; PersistentDataset hashes data+transform spec, so different trainers get separate entries automatically).
2. CPU random crop (get_random_crop, wrapped via monai.data.Dataset(base, transform=crop)): RandWeightedCropd (body-mask-weighted) or RandSpatialCropSamplesd. Yields num_samples=patches_per_volume uniform 128³ patches per volume in DataLoader workers. Skipped entirely by MAISI (full-volume training).
3. GPU augmentation (get_gpu_transforms + gpu_augment_batch): per-item, on device — RandAffined → RandFlipd×3 → RandBiasFieldd (MRI) → RandAdjustContrastd (MRI) → RandGaussianNoised (MRI) → ScaleIntensityd. Operates on either 128³ patches (amix/unet) or full padded volumes (MAISI).

DataLoader uses default list_data_collate (uniform-shape patches stack cleanly; for MAISI batch_size=1 so no collation issue). Effective batch size for amix/unet = batch_size × patches_per_volume.

Validation: MONAI sliding_window_inference walks val_patch_size³ patches over the full padded volume → unpad to original_shape → MAE, PSNR, SSIM, Bone Dice logged to WandB. MAISI also runs the VAE encode/decode via sliding window (_encode_sliding_window, _decode_sliding_window).

Configuration Pattern

All trainers define an EXPERIMENT_CONFIG list at the top of train.py — each dict overrides DEFAULT_CONFIG. Multiple experiments run sequentially in one job. CLI args further override the first experiment's config.

Development Standards (from GEMINI.md)

• Minimalism: Do not refactor working code. Minimal, simple changes preferred.
• Propose before editing: For main code (training, production, shared src/ modules), propose changes in chat before editing. For throwaway analysis/testing scripts you use as tools, just write or edit directly.
• Verify paths: Many paths are hardcoded to /gpfs/accounts/jjparkcv_root/jjparkcv98/minsukc/. Verify before running.
• NEVER overwrite checkpoints, predictions, metrics or eval results. Scripts here write to fixed per-arm/per-model paths (ckpt_<arm>.pt, evaluation/volumes/<arm>/, results_*.json), so a short test run under a real arm/model name silently destroys a finished run — there is no GPFS snapshot to recover from. Before ANY run, check where it writes; smoke-test with a tag, throwaway name, or scratch dir. Same trap applies across evaluation/.
• Ad-hoc output goes in ~/MRI2CT/temp/: unless the user names a destination, write scratch outputs (debug figures, one-off exports, smoke-test dumps) to temp/, not into figs/, bone_study/, or other project dirs — keeps the repo top level uncluttered. Doesn't apply to a script's own documented --out_dir default unless asked to change it. temp/ is gitignored.