MRIComp4Flow

This is the official PyTorch implementation of the paper MRIComp4Flow: Compression of 3D Brain MRI for Training Multi-Modal Generative Models by Lisa K. Fischer, Mykhailo Riabets, Daniel Rueckert, Benedikt Wiestler, Anke Meyer-Baese and Sandeep Nagar.

If you find our work useful, please consider to star this repository and cite our paper.

MRIComp4Flow is a PyTorch-based workflow for studying how compression affects 3D brain MRI synthesis with Wavelet Flow Matching (WFM). The repository combines offline compression, model training, sampling, and quantitative evaluation in a single reproducible pipeline.

Project overview

The codebase supports three main stages:

1. Compression of BraTS volumes into JPEG2000 or JPEG-LS slices.
2. Training of a WFM model on compressed or uncompressed data.
3. Sampling of missing MRI contrasts and evaluation of synthesis quality.

Core components:

• Compression utilities: utils/compress_brats_j2k.py, utils/compress_brats_jpgLS.py
• Training entry point: scripts/train_wfm.py
• Sampling and evaluation: scripts/sample_wfm.py, scripts/evaluate.py
• Analysis utilities: scripts/quantify_compression.py, scripts/measure_slice_psnr.py, scripts/calculate_size.py
• Bash wrappers: bash_scripts/compress.sh, bash_scripts/train.sh, bash_scripts/sample.sh, bash_scripts/analyze.sh

Environment

Create and activate the conda environment:

conda env create -f environment.yml
conda activate MRIComp4Flow

Data layout

The expected dataset structure is:

MRIComp4Flow/
└── data/
    └── BRATS/
        ├── training/
        │   └── BraTS-GLI-00000-000/
        │       ├── BraTS-GLI-00000-000-t1n.nii.gz
        │       ├── BraTS-GLI-00000-000-t1c.nii.gz
        │       ├── BraTS-GLI-00000-000-t2w.nii.gz
        │       └── BraTS-GLI-00000-000-t2f.nii.gz
        └── validation/
            └── BraTS-GLI-00000-000/
                ├── BraTS-GLI-00000-000-t1n.nii.gz
                ├── BraTS-GLI-00000-000-t1c.nii.gz
                ├── BraTS-GLI-00000-000-t2w.nii.gz
                └── BraTS-GLI-00000-000-t2f.nii.gz

The repository is designed for BraTS 2024 glioma data from synapse (SynID: syn53708249) and assumes one patient directory per subject.

Compression

The recommended entry point is the wrapper script bash_scripts/compress.sh. It activates the conda environment and forwards the request to the appropriate compressor.

JPEG2000

bash bash_scripts/compress.sh \
  --codec jpeg2000 \
  --input ./data/BRATS/training \
  --output ./data/BRATS/compressed/j2k_training \
  --level 100

To generate multiple ratios in one call:

bash bash_scripts/compress.sh \
  --codec jpeg2000 \
  --input ./data/BRATS/training \
  --output ./data/BRATS/compressed/j2k_training \
  --level "50 100"

JPEG-LS

bash bash_scripts/compress.sh \
  --codec jpeg_ls \
  --input ./data/BRATS/training \
  --output ./data/BRATS/compressed/jls_training \
  --level 1

The wrapper writes outputs under folders named ratio_<r>/... for JPEG2000 and near_<n>/... for JPEG-LS.

Key flags (compress.sh)

• --codec: jpeg2000 (j2k) or jpeg_ls (jls)
• --input: input root containing BraTS patient subfolders
• --output: output root (the script will create ratio_<r>/... or near_<n>/... folders)
• --level: for jpeg2000 supply compression ratio(s); for jpeg_ls supply NEAR values (lossy_error). Multiple values are supported for jpeg2000.

Training

Training is handled by bash_scripts/train.sh, which forwards arguments to scripts/train_wfm.py.

Example:

bash bash_scripts/train.sh \
  --codec jpeg2000 \
  --level 100 \
  --data_dir ./data/BRATS/compressed/j2k_training/ratio_100 \
  --samples_dir ./runs/j2k_ratio_100/samples \
  --checkpoint_dir ./runs/j2k_ratio_100/checkpoints \
  --max_iterations 50000

Key flags (train.sh)

• --codec: nifti, jpeg2000, or jpeg_ls
• --level: compression ratio for JPEG2000 or NEAR value for JPEG-LS
• --data_dir: path to the training data root
• --samples_dir: output location for sample images
• --checkpoint_dir: output location for checkpoints
• --batch_size: training batch size
• --max_iterations: maximum number of training iterations
• --wandb_entity / --wandb_project: optional Weights & Biases tracking
• --disable_wandb: disable wandb logging

Sampling and evaluation

After training, generate missing contrasts with the sampling wrapper:

bash bash_scripts/sample.sh \
  --mode wfm \
  --data_dir ./data/BRATS/validation \
  --model_path ./runs/j2k_ratio_100/checkpoints/wfm_unified_050000.pt \
  --output_dir ./results/j2k_ratio_100 \
  --contr all \
  --sampling_steps 1 \
  --evaluate

The wrapper runs scripts/sample_wfm.py and, if enabled, evaluates the outputs with scripts/evaluate.py.

For direct sampling without the wrapper:

python scripts/sample_wfm.py \
  --data_dir ./data/BRATS/validation \
  --model_path ./runs/j2k_ratio_100/checkpoints/wfm_unified_050000.pt \
  --output_dir ./results/j2k_ratio_100 \
  --sampling_steps 1 \
  --contr all

For evaluation only:

python scripts/evaluate.py \
  --pred_dir ./results/j2k_ratio_100 \
  --contr t1n \
  --output metrics_t1n.txt

Key flags (sample.sh)

• --mode: wfm for model-based sampling or prior for the source-prior baseline
• --data_dir: path to the validation data root
• --model_path: path to a trained model checkpoint (required for --mode wfm)
• --output_dir: directory for sampled outputs
• --contr: modality to synthesize (t1n, t1c, t2w, t2f, or all)
• --sampling_steps: number of sampling steps
• --evaluate: run evaluation after sampling
• --use_mask: restrict evaluation metrics to brain tissue regions

Analysis utilities

The analysis wrapper supports four tasks:

# Evaluate predictions
bash bash_scripts/analyze.sh \
  --task evaluate \
  --pred_dir ./results/j2k_ratio_100 \
  --output_prefix metrics_j2k_ratio_100

# Quantify compression fidelity
bash bash_scripts/analyze.sh \
  --task quantify \
  --original_dir ./data/BRATS/training \
  --compressed_dir ./data/BRATS/compressed/j2k_training \
  --codec j2k \
  --levels "50 100" \
  --output compression_fidelity_j2k.csv

# Measure per-slice PSNR/SSIM across methods
bash bash_scripts/analyze.sh \
  --task psnr \
  --method_dirs "./data/BRATS/compressed/jls_training ./data/BRATS/compressed/j2k_training" \
  --original ./data/BRATS/training \
  --output slice_psnr.csv

# Report disk usage per patient
bash bash_scripts/analyze.sh \
  --task size \
  --data_dir ./data/BRATS/training

Key flags (analyze.sh)

• --task: evaluate, psnr, quantify, or size
• --pred_dir: path to prediction outputs for evaluate
• --method_dirs: space-separated directories for psnr
• --original / --original_dir: original data path for psnr and quantify
• --compressed_dir: compressed data path for quantify
• --codec: compression codec for quantify (j2k or jls)
• --levels: compression levels for quantify
• --data_dir: data root for size

Notes

• The bash wrappers assume a POSIX shell and a working conda installation.
• The example commands above follow the current wrapper interfaces in bash_scripts/.
• You can also run the underlying Python scripts directly if you prefer not to use the wrappers.
• You can find pretrained weights in ./pretrained_weights folder

When using this codebase please cite our paper.