# jonathangodwin/cvae architecture

> Architecture diagram of the GitHub repository jonathangodwin/cvae, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-24). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/jonathangodwin/cvae
- Repository: https://github.com/jonathangodwin/cvae

## Overview

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`](https://github.com/jonathangodwin/cvae/blob/main/main.py) (CVAE orchestration), [`main_diffusion.py`](https://github.com/jonathangodwin/cvae/blob/main/main_diffusion.py) (Diffusion orchestration), [`gradio_app/explorer_generation.py`](https://github.com/jonathangodwin/cvae/blob/main/gradio_app/explorer_generation.py) (Mask generation), [`src/dataset/ProstateMasksPreprocessing.py`](https://github.com/jonathangodwin/cvae/blob/main/src/dataset/ProstateMasksPreprocessing.py) (Mask preprocessing)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_mask_data["Mask data"]
  node_mask_preprocessing["Mask preprocessing"]
  node_dataset["CVAE dataset<br/>[PICAIDataset.py]"]
end

subgraph group_cvae_workflow["CVAE workflow"]
  node_config["Configuration<br/>[config_loader.py]"]
  node_paths["Data paths<br/>[paths.py]"]
  node_main["CVAE orchestration<br/>[main.py]"]
  node_cvae["Conditional VAE<br/>[CVAE.py]"]
  node_trainer["CVAE trainer<br/>[CVAETrainer.py]"]
  node_run_artifacts["Run artifacts<br/>[paths.py]"]
  node_evaluator["Model evaluator<br/>[Evaluator.py]"]
  node_plotter["Evaluation plots<br/>[Plotter.py]"]
end

subgraph group_exploration["Interactive exploration"]
  node_explorer["Patient explorer UI<br/>[explorer_app.py]"]
  node_callbacks["Explorer callbacks"]
  node_model_loader["Explorer model loader"]
  node_explorer_generation["Mask generation"]
end

subgraph group_mri_synthesis["MRI synthesis"]
  node_mri_preprocessing["MRI preprocessing"]
  node_nifti_export["NIfTI pair export<br/>[NiftiExporter.py]"]
  node_diffusion_entry["Diffusion orchestration<br/>[main_diffusion.py]"]
  node_diffusion_trainer["3D DDPM trainer<br/>[DDPMTrainer3D.py]"]
  node_ddpm_core["Diffusion model"]
  node_unet["Conditional 3D U-Net<br/>[med_ddpm_unet.py]"]
  node_mri_generation["MRI generation<br/>[generate.py]"]
end

node_researcher(("Researcher"))
node_picai_data["PI-CAI cases"]
node_synthetic_mri["Synthetic T2 MRI"]

node_researcher -->|"starts training"| node_main
node_researcher -->|"opens explorer"| node_explorer
node_picai_data -->|"provides masks"| node_mask_preprocessing
node_mask_preprocessing -->|"supplies processed masks"| node_dataset
node_main -->|"loads settings"| node_config
node_main -->|"resolves paths"| node_paths
node_main -->|"loads cases"| node_dataset
node_main -->|"trains folds"| node_trainer
node_trainer -->|"optimizes model"| node_cvae
node_trainer -->|"writes checkpoints"| node_run_artifacts
node_main -->|"evaluates folds"| node_evaluator
node_evaluator -->|"reads latents"| node_run_artifacts
node_main -->|"requests plots"| node_plotter
node_evaluator -->|"provides metrics"| node_plotter
node_explorer -->|"dispatches actions"| node_callbacks
node_callbacks -->|"uses loaded models"| node_model_loader
node_callbacks -->|"requests generation"| node_explorer_generation
node_model_loader -->|"loads model"| node_cvae
node_explorer_generation -->|"encodes and decodes"| node_cvae
node_explorer_generation -->|"reads latent statistics"| node_run_artifacts
node_explorer -->|"selects patient"| node_picai_data
node_picai_data -->|"provides MRI cases"| node_mri_preprocessing
node_mask_preprocessing -->|"aligns MRI crops"| node_mri_preprocessing
node_researcher -->|"starts diffusion workflow"| node_diffusion_entry
node_diffusion_entry -->|"builds MRI pairs"| node_mri_preprocessing
node_diffusion_entry -->|"exports pairs"| node_nifti_export
node_diffusion_entry -->|"starts training"| node_diffusion_trainer
node_diffusion_trainer -->|"builds trainer"| node_ddpm_core
node_diffusion_trainer -->|"builds denoiser"| node_unet
node_ddpm_core -->|"denoises volumes"| node_unet
node_diffusion_entry -->|"starts sampling"| node_mri_generation
node_mri_generation -->|"generates volumes"| node_synthetic_mri
node_explorer_generation -.->|"supplies generated masks"| node_mri_generation

click node_config "https://github.com/jonathangodwin/cvae/blob/main/src/config/config_loader.py"
click node_paths "https://github.com/jonathangodwin/cvae/blob/main/src/config/paths.py"
click node_mask_preprocessing "https://github.com/jonathangodwin/cvae/blob/main/src/dataset/ProstateMasksPreprocessing.py"
click node_dataset "https://github.com/jonathangodwin/cvae/blob/main/src/dataset/PICAIDataset.py"
click node_main "https://github.com/jonathangodwin/cvae/blob/main/main.py"
click node_cvae "https://github.com/jonathangodwin/cvae/blob/main/src/cvae/CVAE.py"
click node_trainer "https://github.com/jonathangodwin/cvae/blob/main/src/cvae/CVAETrainer.py"
click node_run_artifacts "https://github.com/jonathangodwin/cvae/blob/main/src/config/paths.py"
click node_evaluator "https://github.com/jonathangodwin/cvae/blob/main/src/metrics/Evaluator.py"
click node_plotter "https://github.com/jonathangodwin/cvae/blob/main/src/metrics/Plotter.py"
click node_explorer "https://github.com/jonathangodwin/cvae/blob/main/gradio_app/explorer_app.py"
click node_callbacks "https://github.com/jonathangodwin/cvae/blob/main/gradio_app/explorer_callbacks.py"
click node_model_loader "https://github.com/jonathangodwin/cvae/blob/main/gradio_app/explorer_model_loader.py"
click node_explorer_generation "https://github.com/jonathangodwin/cvae/blob/main/gradio_app/explorer_generation.py"
click node_mri_preprocessing "https://github.com/jonathangodwin/cvae/blob/main/src/diffusion/DDPMPreprocessing.py"
click node_nifti_export "https://github.com/jonathangodwin/cvae/blob/main/src/diffusion/NiftiExporter.py"
click node_diffusion_entry "https://github.com/jonathangodwin/cvae/blob/main/main_diffusion.py"
click node_diffusion_trainer "https://github.com/jonathangodwin/cvae/blob/main/src/diffusion/DDPMTrainer3D.py"
click node_ddpm_core "https://github.com/jonathangodwin/cvae/blob/main/src/diffusion/med_ddpm/med_ddpm_trainer.py"
click node_unet "https://github.com/jonathangodwin/cvae/blob/main/src/diffusion/med_ddpm/med_ddpm_unet.py"
click node_mri_generation "https://github.com/jonathangodwin/cvae/blob/main/src/diffusion/generate.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_mask_preprocessing,node_dataset toneBlue
class node_config,node_paths,node_main,node_cvae,node_trainer,node_run_artifacts,node_evaluator,node_plotter toneAmber
class node_explorer,node_callbacks,node_model_loader,node_explorer_generation toneMint
class node_mri_preprocessing,node_nifti_export,node_diffusion_entry,node_diffusion_trainer,node_ddpm_core,node_unet,node_mri_generation toneRose
class node_researcher,node_picai_data,node_synthetic_mri toneIndigo
```

## Components

### Mask data

- [Mask preprocessing](https://github.com/jonathangodwin/cvae/blob/main/src/dataset/ProstateMasksPreprocessing.py): `src/dataset/ProstateMasksPreprocessing.py`
- [CVAE dataset](https://github.com/jonathangodwin/cvae/blob/main/src/dataset/PICAIDataset.py): `src/dataset/PICAIDataset.py`

### CVAE workflow

- [Configuration](https://github.com/jonathangodwin/cvae/blob/main/src/config/config_loader.py): `src/config/config_loader.py`
- [Data paths](https://github.com/jonathangodwin/cvae/blob/main/src/config/paths.py): `src/config/paths.py`
- [CVAE orchestration](https://github.com/jonathangodwin/cvae/blob/main/main.py): `main.py`
- [Conditional VAE](https://github.com/jonathangodwin/cvae/blob/main/src/cvae/CVAE.py): `src/cvae/CVAE.py`
- [CVAE trainer](https://github.com/jonathangodwin/cvae/blob/main/src/cvae/CVAETrainer.py): `src/cvae/CVAETrainer.py`
- [Run artifacts](https://github.com/jonathangodwin/cvae/blob/main/src/config/paths.py): `src/config/paths.py`
- [Model evaluator](https://github.com/jonathangodwin/cvae/blob/main/src/metrics/Evaluator.py): `src/metrics/Evaluator.py`
- [Evaluation plots](https://github.com/jonathangodwin/cvae/blob/main/src/metrics/Plotter.py): `src/metrics/Plotter.py`

### Interactive exploration

- [Patient explorer UI](https://github.com/jonathangodwin/cvae/blob/main/gradio_app/explorer_app.py): `gradio_app/explorer_app.py`
- [Explorer callbacks](https://github.com/jonathangodwin/cvae/blob/main/gradio_app/explorer_callbacks.py): `gradio_app/explorer_callbacks.py`
- [Explorer model loader](https://github.com/jonathangodwin/cvae/blob/main/gradio_app/explorer_model_loader.py): `gradio_app/explorer_model_loader.py`
- [Mask generation](https://github.com/jonathangodwin/cvae/blob/main/gradio_app/explorer_generation.py): `gradio_app/explorer_generation.py`

### MRI synthesis

- [MRI preprocessing](https://github.com/jonathangodwin/cvae/blob/main/src/diffusion/DDPMPreprocessing.py): `src/diffusion/DDPMPreprocessing.py`
- [NIfTI pair export](https://github.com/jonathangodwin/cvae/blob/main/src/diffusion/NiftiExporter.py): `src/diffusion/NiftiExporter.py`
- [Diffusion orchestration](https://github.com/jonathangodwin/cvae/blob/main/main_diffusion.py): `main_diffusion.py`
- [3D DDPM trainer](https://github.com/jonathangodwin/cvae/blob/main/src/diffusion/DDPMTrainer3D.py): `src/diffusion/DDPMTrainer3D.py`
- [Diffusion model](https://github.com/jonathangodwin/cvae/blob/main/src/diffusion/med_ddpm/med_ddpm_trainer.py): `src/diffusion/med_ddpm/med_ddpm_trainer.py`
- [Conditional 3D U-Net](https://github.com/jonathangodwin/cvae/blob/main/src/diffusion/med_ddpm/med_ddpm_unet.py): `src/diffusion/med_ddpm/med_ddpm_unet.py`
- [MRI generation](https://github.com/jonathangodwin/cvae/blob/main/src/diffusion/generate.py): `src/diffusion/generate.py`

### Other components

- **Researcher**
- **PI-CAI cases**
- **Synthetic T2 MRI**

## Connections

- 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