How jonathangodwin/cvae works
This research application trains a conditional VAE to generate 3D prostate and lesion masks from clinical conditions, evaluates the resulting models, and provides an interactive patient explorer. A separate conditional 3D DDPM pipeline synthesizes T2-weighted MRI from segmentation masks, including masks produced by the CVAE. The main training orchestration and several documented capabilities are represented from the README; links for unsampled modules are limited to those documented workflows.
Main entry points: main.py CVAE orchestration · main_diffusion.py Diffusion orchestration · explorer_generation.py Mask generation · ProstateMasksPreprocessing.py Mask preprocessing
Components 24
Mask data
- Mask preprocessing
src/dataset/ProstateMasksPreprocessing.py - CVAE dataset
src/dataset/PICAIDataset.py
CVAE workflow
- Configuration
src/config/config_loader.py - Data paths
src/config/paths.py - CVAE orchestration
main.py - Conditional VAE
src/cvae/CVAE.py - CVAE trainer
src/cvae/CVAETrainer.py - Run artifacts
src/config/paths.py - Model evaluator
src/metrics/Evaluator.py - Evaluation plots
src/metrics/Plotter.py
Interactive exploration
- Patient explorer UI
gradio_app/explorer_app.py - Explorer callbacks
gradio_app/explorer_callbacks.py - Explorer model loader
gradio_app/explorer_model_loader.py - Mask generation
gradio_app/explorer_generation.py
MRI synthesis
- MRI preprocessing
src/diffusion/DDPMPreprocessing.py - NIfTI pair export
src/diffusion/NiftiExporter.py - Diffusion orchestration
main_diffusion.py - 3D DDPM trainer
src/diffusion/DDPMTrainer3D.py - Diffusion model
src/diffusion/med_ddpm/med_ddpm_trainer.py - Conditional 3D U-Net
src/diffusion/med_ddpm/med_ddpm_unet.py - MRI generation
src/diffusion/generate.py
Other components
- Researcher
- PI-CAI cases
- Synthetic T2 MRI
Connections 33
- Researcher → CVAE orchestration: starts training
- Researcher → Patient explorer UI: opens explorer
- PI-CAI cases → Mask preprocessing: provides masks
- Mask preprocessing → CVAE dataset: supplies processed masks
- CVAE orchestration → Configuration: loads settings
- CVAE orchestration → Data paths: resolves paths
- CVAE orchestration → CVAE dataset: loads cases
- CVAE orchestration → CVAE trainer: trains folds
- CVAE trainer → Conditional VAE: optimizes model
- CVAE trainer → Run artifacts: writes checkpoints
- CVAE orchestration → Model evaluator: evaluates folds
- Model evaluator → Run artifacts: reads latents
- CVAE orchestration → Evaluation plots: requests plots
- Model evaluator → Evaluation plots: provides metrics
- Patient explorer UI → Explorer callbacks: dispatches actions
- Explorer callbacks → Explorer model loader: uses loaded models
- Explorer callbacks → Mask generation: requests generation
- Explorer model loader → Conditional VAE: loads model
- Mask generation → Conditional VAE: encodes and decodes
- Mask generation → Run artifacts: reads latent statistics
- Patient explorer UI → PI-CAI cases: selects patient
- PI-CAI cases → MRI preprocessing: provides MRI cases
- Mask preprocessing → MRI preprocessing: aligns MRI crops
- Researcher → Diffusion orchestration: starts diffusion workflow
- Diffusion orchestration → MRI preprocessing: builds MRI pairs
- Diffusion orchestration → NIfTI pair export: exports pairs
- Diffusion orchestration → 3D DDPM trainer: starts training
- 3D DDPM trainer → Diffusion model: builds trainer
- 3D DDPM trainer → Conditional 3D U-Net: builds denoiser
- Diffusion model → Conditional 3D U-Net: denoises volumes
- Diffusion orchestration → MRI generation: starts sampling
- MRI generation → Synthetic T2 MRI: generates volumes
- Mask generation → MRI generation: supplies generated masks