Skip to content

manishraut77/DoseNet

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

2 Commits
 
 
 
 
 
 
 
 
 
 

Repository files navigation

DoseNet

AI-assisted radiotherapy dose planning — predict, visualize, and evaluate treatment plans in seconds.

Radiotherapy treatment planning is one of the most time-intensive steps in cancer care. A dosimetrist manually iterates a 3D dose distribution across CT scans and anatomical structures, trying to cover tumors while protecting surrounding organs. This process can take hours per patient and requires significant expertise.

DoseNet is an end-to-end system that takes a patient's CT scan and organ/tumor structure masks, runs a deep 3D neural network to predict the optimal dose distribution, and lets clinicians explore the result interactively — all in a browser.


What It Does

  1. Upload a patient folder (CT + structure masks as sparse CSV files)
  2. Predict a 3D dose distribution instantly via a trained 3D U-Net
  3. Visualize the result across axial, coronal, and sagittal planes with dose overlay
  4. Evaluate plan quality with Dose-Volume Histograms (DVH) and a clinical scorecard
  5. Export the predicted dose as a downloadable CSV or NIfTI volume

Demo

Upload patient folder → instant dose prediction → 3-plane CT viewer + DVH chart

Minimum required files per patient:

pt_001/
├── ct.csv                  ← CT scan (sparse HU values)
├── possible_dose_mask.csv  ← clinically reachable voxels
├── PTV56.csv / PTV63.csv / PTV70.csv   ← tumor targets (optional)
└── Brainstem.csv / SpinalCord.csv / ...  ← organs at risk (optional)

The app auto-predicts as soon as valid files are detected. No button clicking needed.


Architecture

┌──────────────────────────────────────────────────────────────────┐
│                        Browser (React)                           │
│                                                                  │
│  Upload CSVs → parse CT locally → POST /predict-folder-download  │
│                                                                  │
│  ┌────────────────┐   ┌──────────────────┐   ┌───────────────┐  │
│  │  3-Plane Viewer│   │   DVH Chart (SVG)│   │   Scorecard   │  │
│  │  axial/cor/sag │   │  per-ROI curves  │   │  pass/fail    │  │
│  │  Canvas 2D     │   │  Web Worker      │   │  PTV+OAR      │  │
│  └────────────────┘   └──────────────────┘   └───────────────┘  │
└──────────────────────────────────────────────────────────────────┘
                              │ HTTP
┌──────────────────────────────────────────────────────────────────┐
│                       FastAPI Backend                            │
│                                                                  │
│  sparse CSV → dense tensors → DosePredictorUNet3D → sparse CSV   │
│                                                                  │
│  also serves: NIfTI bundles, CT slice enhancement                │
└──────────────────────────────────────────────────────────────────┘
                              │ PyTorch
┌──────────────────────────────────────────────────────────────────┐
│           DosePredictorUNet3D  (Residual 3D U-Net)               │
│                                                                  │
│  Input: CT [1ch] + ROI masks [10ch] → 128×128×128 volume         │
│  FiLM conditioning on plan variant (tumor / balanced / OAR)      │
│  Output: predicted dose [1ch] in Gy                              │
└──────────────────────────────────────────────────────────────────┘

Tech Stack

Layer Technology
Model PyTorch, Residual 3D U-Net, FiLM conditioning
Backend FastAPI, Uvicorn, nibabel, OpenCV
Frontend React, Vite, Canvas 2D, SVG, Web Workers
Dataset OpenKBP Challenge (340 head-and-neck patients)
Inference device CUDA GPU or CPU (auto-detected)

Quickstart

You need two terminals — one for the backend, one for the frontend.

Backend

cd backend
python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt
./run.sh

Starts at http://127.0.0.1:8000. Verify with:

curl http://127.0.0.1:8000/health

Frontend

cd frontend
npm install
npm run dev

Opens at http://127.0.0.1:5173. Upload a patient folder and the app handles the rest.

Environment (optional)

# backend/.env
MODEL_DEVICE=cuda        # or cpu
BACKEND_PORT=8000
CORS_ORIGINS=http://localhost:5173

# frontend/.env
VITE_API_PROXY_TARGET=http://127.0.0.1:8000

The Model

Architecture: DosePredictorUNet3D — a Residual 3D U-Net trained on 128×128×128 voxel CT volumes.

Input [B, 11, 128, 128, 128]
  └─ 1 CT channel (normalized HU)
  └─ 10 ROI channels (7 OARs + 3 PTVs, binary masks)

Encoder:
  Stem  → 28 ch
  Enc1  → 56 ch   + 2× downsample
  Enc2  → 112 ch  + 2× downsample
  Enc3  → 224 ch  + 2× downsample

Bottleneck:
  2× ResidualBlock3D
  FiLMConditioner (injects 2D plan tradeoff weights → γ, β)

Decoder (with skip connections):
  Up2 → 112 ch
  Up1 → 56 ch
  Up0 → 28 ch

Output Head → [B, 1, 128, 128, 128] dose in Gy

Building blocks: ConvNormAct3D (Conv → GroupNorm → SiLU), ResidualBlock3D, UpBlock3D (trilinear upsample + skip concat).

FiLM conditioning lets you steer the prediction toward three clinical goals at inference time by passing a 2D tradeoff vector:

Plan Variant Tumor emphasis OAR emphasis
tumor_focused 0.85 0.15
balanced 0.50 0.50
oar_sparing 0.30 0.70

Training

Trained on the OpenKBP head-and-neck dataset — 340 patients with CT, structure masks, and clinical dose plans.

Run training

python3 data/training/train.py \
  --data-root /path/to/openkbp \
  --device cuda \
  --epochs 80 \
  --batch-size 1 \
  --plan-variant balanced

Loss function

Plain L1 on 128³ volumes makes the model predict near-zero everywhere (most voxels receive no dose). We use a three-term composite loss restricted to the clinically reachable region:

L = w_dose × masked_L1
  + w_ptv  × PTV_underdose_penalty    (coverage shortfall per target)
  + w_oar  × OAR_overdose_penalty     (squared excess above 30 Gy)

Weights vary by plan variant. This pushes the model toward meaningful clinical behavior rather than the trivial near-zero solution.

Training setup

Setting Value
Optimizer AdamW
LR schedule OneCycleLR
Batch size 1 (128³ volumes are large)
Mixed precision AMP enabled by default
Gradient clipping norm ≤ 1.0
EMA decay 0.999
Early stopping patience 16 epochs

Validation results (example run)

Metric Value What it means
train_loss 2.77 Composite training objective on the training split (masked_l1 + PTV underdose penalty + OAR overdose penalty, with variant weights). Lower is better.
val_loss 4.16 Same composite objective on held-out validation patients. Higher than train loss indicates the expected train/validation generalization gap.
dose_score 2.94 Gy OpenKBP dose score: mean absolute voxel-wise dose error (Gy) inside the possible-dose region. Lower is better.
dvh_score 1.74 Gy OpenKBP DVH score: mean absolute error (Gy) over ROI DVH metrics (D_0.1_cc, mean, D_99, D_95, D_1). Lower is better.
Prescription deviation ~2-3% Predicted target coverage is typically within about 2-3% of prescribed dose levels on validation cases.
Clinical interpretation Clinically aligned surrogate model The model tracks clinically meaningful target-vs-OAR tradeoffs and is useful as a planning aid (not a final clinical plan without clinician review).

Deploy a newly trained model

cp data/training/runs/<run_name>/checkpoints/best.pt backend/models/best.pt
# restart backend

Visualization

The viewer is built entirely on Canvas 2D — no WebGL dependency.

Three-plane MPR viewer

Axial, coronal, and sagittal slices rendered simultaneously. Each plane extracts a 2D slice from the 128³ volume and renders it at 2× resolution for sharpness.

CT windowing (DICOM window/level):

Preset Level Width
Soft Tissue 40 HU 400 HU
Bone 400 HU 2000 HU
Lung −600 HU 1500 HU

Dose overlay: 7-stop perceptual colormap (blue → cyan → green → yellow → orange → red), alpha-blended over grayscale CT with gamma correction so anatomy stays visible under the dose wash.

Dose contours: Marching-squares algorithm generates iso-dose lines, linked into continuous polylines, smoothed with a weighted averaging pass, and rendered with a two-pass stroke (dark underlay + colored overlay).

Optional CT enhancement: Backend OpenCV pipeline (CLAHE + denoise + detail boost + tone shaping) applied per slice with a 140ms debounce on drag.

DVH chart

Dose-Volume Histograms computed in a Web Worker (non-blocking):

  • 121-bin cumulative histogram per ROI
  • Interactive crosshair: hover to read volume% at any dose
  • Toggle individual curves on/off via legend
  • Reference lines: prescription doses (56/63/70 Gy) and OAR limits

Plan quality scorecard

Automated clinical heuristics:

Check Metric Goal
PTV56 coverage D95 ≥ 53.2 Gy
PTV63 coverage D95 ≥ 59.9 Gy
PTV70 coverage D95 ≥ 66.5 Gy
Parotid mean mean dose ≤ 26 Gy
Spinal Cord Dmax ≤ 50 Gy
Brainstem Dmax ≤ 54 Gy

Overall status: Clinically Acceptable or Needs Review.


Data Format

Patient data follows OpenKBP sparse CSV conventions:

Flat index format (CT, dose, masks):

index,data
18432,1024
18433,890

XYZ format (alternative, also supported):

x,y,z,value
10,20,30,7.5

Mask-only format (just indices, value defaults to 1):

index
18432
18433

All volumes are 128×128×128 voxels. CT HU values are normalized to [0, 1] by clipping to [0, 4095] and dividing.

Supported ROI names:

  • OARs: Brainstem, SpinalCord, RightParotid, LeftParotid, Esophagus, Larynx, Mandible
  • Targets: PTV56, PTV63, PTV70

API Reference

Base URL: http://127.0.0.1:8000

Method Endpoint Description
GET /health Model load status, device, errors
POST /predict-folder Predict dose, return stats JSON
POST /predict-folder-download Predict dose, return sparse CSV stream
POST /volumes/prepare Generate NIfTI bundle from uploaded files
GET /volumes/bundles/{id}/ct.nii.gz Download CT NIfTI
GET /volumes/bundles/{id}/dose.nii.gz Download dose NIfTI
GET /volumes/{patient_id}/ct.nii.gz Latest CT bundle for patient
GET /volumes/{patient_id}/dose.nii.gz Latest dose bundle for patient
POST /enhance_slice Enhance a single CT slice (OpenCV)

POST /predict-folder-download form fields:

  • patient_id — string identifier
  • plan_varianttumor_focused | balanced | oar_sparing
  • files — one or more CSV files

Response headers include:

  • X-Predicted-Dose-Mean/Min/Max — dose statistics
  • X-Model-Devicecuda or cpu
  • X-Volume-Bundle-Id — NIfTI bundle identifier
  • X-Ct-Nifti-Url / X-Dose-Nifti-Url — 3D volume download links

Project Structure

DoseNet/
├── backend/
│   ├── app.py                        # FastAPI app + all endpoints
│   ├── run.sh                        # Startup script
│   ├── requirements.txt
│   ├── models/
│   │   └── best.pt                   # Trained model checkpoint
│   ├── services/
│   │   ├── volume_conversion.py      # sparse CSV ↔ dense volume ↔ NIfTI
│   │   └── enhancer.py               # OpenCV CT slice enhancement
│   └── tests/
│       └── test_volume_conversion.py
│
├── frontend/
│   └── src/
│       ├── pages/HomePage.jsx        # Upload, predict, orchestration
│       ├── components/VolumeViewer.jsx  # 3-plane viewer + DVH + scorecard
│       ├── services/api.js           # Backend API client
│       ├── utils/sparseVolume.js     # CSV → dense volume utilities
│       └── workers/dvhWorker.js      # DVH computation (Web Worker)
│
└── data/
    └── training/
        ├── train.py                  # Training entry point
        ├── network_architectures.py  # DosePredictorUNet3D
        ├── network_functions.py      # Training loop, losses, EMA, checkpointing
        ├── data_loader.py            # OpenKBP dataset class
        ├── dose_evaluation_class.py  # Dose score + DVH score metrics
        └── best.pt                   # Training checkpoint copy

Tests

python3 -m unittest backend.tests.test_volume_conversion

Covers: sparse CSV parsing (flat and XYZ), bounds validation, NIfTI shape/dtype/affine, orientation debug volume.


Dependencies

Backend:

fastapi  uvicorn  pydantic  python-multipart
numpy  pandas  torch  nibabel  pynrrd  opencv-python

Frontend:

react  react-dom  vite  @vitejs/plugin-react

Challenges We Ran Into

Memory with 3D volumes. A single 128³ volume already stresses GPU memory. We fixed batch size at 1, enabled mixed-precision training (AMP), clipped gradients, and stabilized weights with EMA.

Silent misalignment from sparse indexing. A single off-by-one in flat-to-3D index reconstruction misaligns CT, masks, and dose with no obvious error. We built a centralized conversion utility with explicit bounds checking and validated voxel counts at each load step.

Model predicting near-zero everywhere. Standard L1 over a volume where 80%+ of voxels have zero dose produces a trivially good loss for near-zero predictions. We masked the loss to the possible dose region and added explicit penalties for under-covering tumors and over-irradiating organs.

Dose overlay hiding anatomy. Opaque dose color makes it impossible to see where you are in the CT. We tuned gamma-corrected alpha blending so low-dose regions stay semi-transparent, and added contour lines so dose iso-lines remain readable over both dark and bright CT regions.


What's Next

  • Uncertainty estimates via MC dropout (decoder dropout is already set up for this)
  • Comparison mode: predicted vs. ground truth dose side-by-side
  • DICOM import support (currently requires OpenKBP-format CSVs)
  • Dose editing: let users adjust and re-evaluate modified plans
  • Faster inference via TorchScript export

Disclaimer

DoseNet is a research prototype built to explore AI-assisted treatment planning. It is not a cleared medical device and must not be used for autonomous clinical treatment decisions. Any clinical use requires formal validation, regulatory clearance, and physician oversight.


Dataset

Trained on the OpenKBP Grand Challenge dataset — 340 head-and-neck cancer patients with paired CT scans, organ/tumor structure masks, and clinical dose plans from multiple institutions. Training data lives under data/training/ alongside the model definition and training scripts.

About

No description, website, or topics provided.

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published