PY-BRAIN

Brain tumor MRI segmentation pipeline with epistemic uncertainty quantification

---

⚠️ Research Use Only

PY-BRAIN is research software and is NOT a certified medical device. It has not been cleared or approved by the FDA, EMA, or any regulatory body. All AI-generated segmentations must be verified by a qualified radiologist before any clinical use. See SECURITY.md for full disclaimers.

---

Highlights

• Multi-model ensemble: SegResNet (60%) + TTA-4 (40%) with Platt-calibrated probabilities
• Epistemic uncertainty: MC-Dropout (15 samples) for confidence-aware segmentation
• BraTS-validated: Dice WT=0.933, TC=0.905, ET=0.925, HD95=10.3mm on BraTS2021
• Reproducible: Fixed random seeds, documented configuration, calibration coefficients versioned
• Modular: Stage-based pipeline (0–9), each stage independently executable
• Research-grade: Not for clinical diagnosis — see disclaimers throughout

---

Project Status

Component	Status	Notes
SegResNet + TTA-4 ensemble	✅ Active & Validated	Primary inference path
MC-Dropout uncertainty	✅ Active	15 samples, SegResNet only
Platt calibration	✅ Active	Per-subregion temperature scaling
STAPLE ensemble	✅ Implemented	Awaiting ensemble validation
SwinUNETR	⏸️ Disabled	Awaiting Platt calibration + channel verification

Validation: Validated on BraTS2021 cases (internal runs, 2026-04-13). Full external validation pending.

---

Repository Structure

pybrain/
├── pybrain/                   # Python package
│   ├── config/                # YAML configuration (thresholds, ensemble weights)
│   ├── core/                  # Metrics, normalization, brainmask utilities
│   ├── io/                    # NIfTI I/O, session, logging
│   ├── models/                # SegResNet, TTA, MC-Dropout, Platt, STAPLE
│   └── clinical/              # WHO grading, clinical consistency checks
├── scripts/                   # Pipeline stage scripts
│   ├── 0_preflight_check.py  # Environment validation
│   ├── 1_dicom_to_nifti.py   # DICOM → NIfTI conversion
│   ├── 2_ct_integration.py   # CT registration + solid tissue prior
│   ├── 3_brain_tumor_analysis.py  # AI segmentation (MAIN)
│   ├── 5_validate_segmentation.py # Dice, HD95, ASD metrics
│   ├── 6_tumour_location.py # Anatomical atlas location
│   ├── 7_tumour_morphology.py # Volume, surface area, compactness
│   ├── 8_radiomics_analysis.py # Radiomics + XGBoost IDH classifier
│   └── 9_generate_report.py  # PDF report
├── models/
│   └── calibration/           # Platt coefficients (version-controlled)
├── docs/
│   ├── DATA.md               # BraTS data acquisition guide
│   └── technical/
│       ├── VALIDATION.md      # Validation methodology & results
│       ├── MODEL_CARD.md      # Model documentation
│       ├── ARCHITECTURE.md    # System design
│       └── REPRODUCIBILITY.md # Reproduction guide
├── .github/workflows/ci.yml   # GitHub Actions (lint + smoke tests)
├── requirements.txt           # pip dependencies
├── environment.yml            # conda dependencies
├── pyproject.toml            # package metadata
├── LICENSE                   # MIT license
├── CITATION.cff              # Citation metadata
├── CONTRIBUTING.md           # Contribution guidelines
└── SECURITY.md               # Security policy & clinical disclaimers

---

Quick Start

1. Install Dependencies

Using pip:

git clone https://github.com/3ruiruirui-sketch/pybrain.git
cd pybrain
python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

Using conda:

git clone https://github.com/3ruiruirui-sketch/pybrain.git
cd pybrain
conda env create -f environment.yml
conda activate pybrain

2. Obtain BraTS Data

Download BraTS2021 from TCIA (requires registration, free for academic use). Place in data/datasets/BraTS2021/.

See docs/DATA.md for full instructions.

3. Run the Pipeline

# Set session ID for reproducibility
export PYBRAIN_SESSION="my_first_run"
export PYBRAIN_SEED=42

# Run main segmentation (Stage 3)
python scripts/3_brain_tumor_analysis.py \
    --input-dir data/datasets/BraTS2021/BraTS2021_00000 \
    --output-dir results/BraTS2021_00000 \
    --session-id BraTS2021_00000

4. Export DICOM-SEG (Optional)

For PACS integration and clinical workstation viewing (3D Slicer, OHIF, RadiAnt):

# Export segmentation as DICOM-SEG file
python scripts/4_export_dicom_seg.py \
    --segmentation results/BraTS2021_00000/segmentation_full.nii.gz \
    --source-dicom data/datasets/BraTS2021/BraTS2021_00000/T1c/ \
    --output results/BraTS2021_00000/segmentation.dcm

The DICOM-SEG file contains three segments (WT, TC, ET) with proper referencing to the source DICOM series.

5. Export DICOM-SR (Optional)

For quantitative measurements in PACS using TID 1500 (Measurement Report) template:

# Export measurements as DICOM-SR file
python scripts/5_export_dicom_sr.py \
    --measurements results/BraTS2021_00000/morphology.json \
    --source-dicom data/datasets/BraTS2021/BraTS2021_00000/T1c/ \
    --segmentation-dicom results/BraTS2021_00000/segmentation.dcm \
    --output results/BraTS2021_00000/measurements.dcm

The DICOM-SR file contains volume measurements (WT, TC, ET, NC) in a machine-readable format for clinical workstations.

6. Longitudinal Analysis (Optional)

For comparing current scan with a prior scan to assess tumor response using RANO criteria:

# Compare current session with prior session
python scripts/10_longitudinal_analysis.py \
    --current results/BraTS2021_00000/ \
    --prior results/BraTS2021_00000_prior/ \
    --output results/BraTS2021_00000/longitudinal/

This performs rigid registration of the prior scan to the current scan, computes volume changes per subregion (WT, TC, ET), and applies RANO criteria for response assessment (CR, PR, SD, PD, NE). Results include:

• Registration quality (NMI score)
• Volume change table with percentage changes
• RANO response classification
• Side-by-side overlay images (axial, coronal, sagittal)

Note: This is a research implementation of RANO criteria and should not be used for clinical decision-making.

Results are written to results/{session-id}/:

results/BraTS2021_00000/
├── segmentation_full.nii.gz     # Final segmentation (WT/TC/ET labels)
├── segmentation.dcm             # DICOM-SEG file (if exported)
├── measurements.dcm            # DICOM-SR file (if exported)
├── longitudinal/               # Longitudinal analysis (if run)
│   ├── longitudinal_results.json
│   ├── prior_registered.nii.gz
│   ├── prior_seg_registered.nii.gz
│   ├── overlay_axial.png
│   ├── overlay_coronal.png
│   └── overlay_sagittal.png
├── probabilities_wt.nii.gz     # WT probability map
├── probabilities_tc.nii.gz     # TC probability map
├── probabilities_et.nii.gz     # ET probability map
├── uncertainty_mc.nii.gz       # MC-Dropout std (uncertainty map)
├── platt_calibration.json      # Calibration metadata
└── report_BraTS2021_00000.pdf  # PDF report

7. PACS Integration (Optional)

For automatic DICOM communication with PACS using Orthanc:

# Start Orthanc (development PACS)
cd docker/orthanc
docker compose up -d

# Start PACS listener daemon
python scripts/11_pacs_listener.py --config pybrain/config/defaults.yaml

The listener will:

• Receive DICOM studies from PACS via C-STORE
• Filter studies based on triage rules (modalities, series descriptions)
• Automatically trigger the pipeline for matching studies
• Auto-route results (DICOM-SEG, DICOM-SR) back to PACS

Configuration (in pybrain/config/defaults.yaml):

pacs:
  enabled: true
  local_ae_title: "PYBRAIN"
  local_port: 11112
  peer:
    ae_title: "ORTHANC"
    host: "localhost"
    port: 4242
  triage:
    require_modalities: ["MR"]
    require_series_descriptions: ["T1", "T1c", "T2", "FLAIR"]
    auto_route_back: true

Note: This requires pynetdicom>=2.0. See docs/technical/PACS_INTEGRATION.md for full setup guide.

8. API Quick Start (Optional)

For running PY-BRAIN as a REST API service with async workers:

# Start the full stack with Docker Compose
docker compose up -d

# Run database migrations
docker compose exec api alembic upgrade head

# Verify deployment
curl http://localhost:8000/health
curl http://localhost:8000/ready

# Open API documentation
open http://localhost:8000/docs

The API service includes:

• FastAPI application on port 8000
• Celery workers for async job processing
• PostgreSQL database for persistent storage
• Redis for Celery broker and backend
• Orthanc PACS (optional)
• Flower for Celery monitoring on port 5555

Example usage:

# Upload a case
curl -X POST http://localhost:8000/cases \
  -H "Authorization: Bearer dev-api-key-123" \
  -F "files=@case.zip" \
  -F "analysis_mode=auto"

# Trigger segmentation
curl -X POST http://localhost:8000/cases/{case_id}/segment \
  -H "Authorization: Bearer dev-api-key-123"

# Get results
curl http://localhost:8000/cases/{case_id} \
  -H "Authorization: Bearer dev-api-key-123"

See docs/API.md for full API reference and docs/DEPLOYMENT.md for deployment details.

9. Validate Results

python scripts/5_validate_segmentation.py \
    --prediction results/BraTS2021_00000/segmentation_full.nii.gz \
    --ground-truth data/datasets/BraTS2021/BraTS2021_00000/BraTS2021_00000_seg.nii.gz

---

Inference Workflow

Input MRI (T1, T1c, T2, FLAIR)
    │
    ▼
ROI Localization (whole-brain crop → focused ROI)
    │
    ▼
┌─────────────────────────────────────────────────────────┐
│          Ensemble Inference (SegResNet + TTA-4)          │
│                                                          │
│  SegResNet ─────────────────────► P_std (60%)          │
│  TTA-4 (4-axis flips) ──────────► P_tta  (40%)          │
│                                                          │
│  P_ensemble = 0.60 × P_std + 0.40 × P_tta               │
└─────────────────────────────────────────────────────────┘
    │
    ▼
┌─────────────────────────────────────────────────────────┐
│            MC-Dropout (if enabled, 15 samples)            │
│  SegResNet forward × 15 with dropout                     │
│  → mean probability + std uncertainty map                │
│  Platt calibration applied to probabilities              │
└─────────────────────────────────────────────────────────┘
    │
    ▼
CT Boost (optional — solid tissue prior from HU[35-75])
    │
    ▼
Per-subregion thresholding: WT / TC / ET
    │
    ▼
Hierarchical consistency (WT ⊇ TC ⊇ ET)
    │
    ▼
Output: segmentation + probability maps + uncertainty

---

Validation Summary

Validated on BraTS2021 cases (2026-04-13):

Case	Dice WT	Dice TC	Dice ET	HD95	Quality
BraTS2021_00000	0.933	0.905	0.925	10.3 mm	EXCELLENT
BraTS2021_00002	0.702	0.827	0.683	19.1 mm	CHALLENGING
BraTS2021_00003	0.967	0.937	0.912	1.4 mm	OUTSTANDING

Case 00002 is a clinically challenging infiltrative tumor — low WT Dice reflects inherent difficulty, not a pipeline failure.

Full validation details: docs/technical/VALIDATION.md

---

Current Model Configuration

Active and disabled models:

Model	Ensemble Weight	Status	Notes
SegResNet	0.60	✅ Active	Primary model, MONAI BraTS pretrained
TTA-4	0.40	✅ Active	4-axis flip averaging
MC-Dropout	—	✅ Active	15 samples, SegResNet only
Platt Calibration	—	✅ Active	Per-subregion temperature scaling
STAPLE	—	✅ Implemented	Awaiting ensemble validation
SwinUNETR	0.00	⏸️ Disabled	Awaiting Platt + channel-order verification

Configuration file: pybrain/config/defaults.yaml

---

Reproducibility Notes

• Set PYBRAIN_SEED=42 before running for deterministic results
• MC-Dropout is stochastic: disable mc_dropout.enabled in config for strict determinism
• Platt calibration coefficients are stored in models/calibration/
• MONAI model weights are downloaded on first run (~500 MB)

Full guide: docs/technical/REPRODUCIBILITY.md

---

Limitations

1. Research only — not a certified medical device; not for clinical diagnosis
2. BraTS-trained — validated only on BraTS-style 4-channel MRI (T1/T1ce/T2/FLAIR)
3. Infiltrative tumors — Case 00002-type tumors with ill-defined edema margins show lower WT Dice (see validation report)
4. Post-operative scans — not validated on post-surgical cases
5. Apple MPS — some PyTorch MPS operations are non-deterministic; CUDA on Linux recommended for strict reproducibility
6. MC-Dropout scope — uncertainty only available for SegResNet; TTA-4 does not produce uncertainty maps

---

Roadmap

Priority	Item	Status
1	SwinUNETR validation and activation	⏸️ Pending
2	Full regression test suite	🔄 In progress
3	STAPLE ensemble validation	🔄 In progress
4	External multi-site validation	🔜 Planned
5	Zenodo DOI registration	🔜 Planned

---

Citation

If you use PY-BRAIN in your research, please cite:

# BibTeX (update with Zenodo DOI when assigned)
@misc{pybrain2026,
  title = {{PY-BRAIN}: Brain Tumor MRI Segmentation with Uncertainty Quantification},
  author = {{PY-BRAIN Contributors}},
  year = {2026},
  publisher = {GitHub},
  journal = {GitHub repository},
  url = {https://github.com/3ruiruirui-sketch/pybrain}
}

Or use the Cite this repository button on GitHub (uses CITATION.cff).

---

License

MIT License — see LICENSE.

Note on model licenses: The MONAI SwinUNETR BraTS bundle is licensed under Apache 2.0. The BrainIAC model has a separate non-commercial academic license from Mass General Brigham. Model weights are downloaded separately and are not included in this repository.

---

Contributing

Contributions are welcome! See CONTRIBUTING.md for:

• Development setup
• Coding standards
• Testing requirements
• Pull request process

Please read our Code of Conduct and Security Policy before contributing.

---

Links

• Documentation
• Model Card
• Validation Report
• Reproducibility Guide
• Data Guide
• GitHub Issues