Brain Tumor Segmentation — Project

This repository is a project for volumetric brain tumor segmentation using MONAI and PyTorch. It provides a reproducible preprocessing, training, and evaluation pipeline (delivered as Jupyter notebooks and supporting instructions) targeting the BraTS-style MRI dataset.

Goals

• Provide a reproducible preprocessing pipeline for BraTS-style MRI volumes (4 modalities + segmentation).
• Train and evaluate 3D segmentation models (UNet / VNet style) with MONAI transforms and patch-based training.
• Handle real-world dataset issues (missing or all-zero modalities, corrupted files) automatically.
• Offer clear instructions to reproduce training runs, experiments, and inference.

Repository layout

• brain-tumor-unet.ipynb — Main project notebook implementing the full preprocessing pipeline, dataset cleaning/imputation, transforms, patch sampling, training loop, and basic visualization. This notebook is the canonical pipeline used in the project.
• brain-tumor-vnet.ipynb — Alternative architecture experiments (VNet-style). Use for model comparison or as a starting point for different network designs.
• README.md — (this file) project-level documentation and instructions.

Project contract (inputs / outputs)

• Inputs:
  • BraTS-style training folder containing subject subfolders with the 4 MRI modalities and a segmentation: *_t1.nii, *_t1ce.nii, *_t2.nii, *_flair.nii, *_seg.nii.
  • Optional: pre-trained model checkpoints for inference.
• Outputs:
  • Training logs and saved model checkpoints (user-configured paths in the notebooks).
  • Example visualizations of slices or patch predictions.
  • (Optional) evaluation metrics (Dice, Hausdorff) when evaluation cells are enabled.

Error modes / behavior:

• Subjects missing segmentation are skipped for training and logged.
• Missing or all-zero image modalities are imputed with a blank NIfTI (keeps a consistent 4-channel input to the model).
• Corrupted files that fail to load are skipped and printed in the preprocessing stage.

Requirements & Reproducibility

Recommended environment:

• Python 3.8+
• PyTorch (compatible with your CUDA version or CPU)
• MONAI (tested with monai==1.5.0 in the notebooks)
• nibabel, numpy, matplotlib

I can add a requirements.txt or Conda environment.yml pinned to the exact versions tested. Tell me whether you prefer Pip or Conda and whether you target CUDA or CPU, and I will add it.

Quick install (Windows PowerShell, CPU example):

python -m venv .venv; .\.venv\Scripts\Activate.ps1
pip install --upgrade pip
pip install monai==1.5.0 nibabel matplotlib numpy
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu

If you have CUDA, replace the PyTorch install with the appropriate command from https://pytorch.org.

How to run (project workflow)

1. Prepare the dataset in a BraTS-style layout. Point data_root_dir (first cells) in brain-tumor-unet.ipynb to your dataset root.
2. Open brain-tumor-unet.ipynb in Jupyter Lab / Notebook or VS Code and run cells top-to-bottom. The notebook contains cells for:
   • Installing dependencies (optional for clean environments)
   • Building a reusable blank NIfTI for imputation
   • Building data dictionaries for MONAI Dataset
   • Preprocessing/transforms (resampling, orientation, normalization, label mapping)
   • Patch sampling and dataloaders
   • Model definition, training loop, and checkpoint saving
   • Basic visualization and sanity checks
3. For evaluation or inference, run the evaluation cells (or ask me to add a dedicated inference.py script and an evaluation notebook).

Notes:

• The notebook is intended to be the canonical runnable pipeline for the project. If you want a script-based runner (for HPC or reproducible CI training), I can extract the notebook logic into train.py + eval.py with a CLI.

Training & Evaluation

• The notebooks show a minimal training loop. For production-quality experiments you will want to:
  • Add configuration via YAML / argparse (learning rate, batch size, patch size, epochs)
  • Integrate a logging backend (TensorBoard, WandB)
  • Add deterministic seeding and checkpoint resume behavior
  • Add a validation loop that computes Dice per class and saves best-checkpoint by metric

If you'd like, I can implement these improvements and add a sample train.sh / train.ps1 runner.

Model, labels and CUDA safety

This project maps BRATS label values (1, 2, 4) to contiguous class indices (1, 2, 3) to avoid CUDA assert errors when using cross-entropy or similar losses. Labels are cast to torch.long before being passed to the loss.

Contribution & development

If you plan to continue development, consider the following additions I can make:

• requirements.txt or environment.yml (you pick Pip vs Conda)
• Extract notebook code into src/ with train.py and predict.py for scriptable runs
• Add a tests/ folder with small unit tests verifying preprocessing steps (header/affine preservation, label mapping, imputation behavior)
• Add a small Makefile / tasks.json to start Jupyter, run training, and run evaluation

Tell me which of these you'd like and I will implement it.

License

Add a license file when you're ready to publish or share. Common choices:

• MIT — permissive, simple for code sharing
• Apache-2.0 — permissive with patent grant

If you tell me which license you prefer, I can add LICENSE now.

Contact / authoring

This repository holds the project notebooks and documentation. If you want me to take the next step, pick one or more of:

• Add requirements.txt for a Pip environment
• Add environment.yml for Conda
• Extract notebook into src/train.py and src/predict.py (script mode)
• Add an inference.ipynb that loads a checkpoint and runs full-volume predictions

I'll implement the chosen tasks and run quick smoke tests to verify environment and basic flow.