How wahib-elbessali/conditional-stylegan2-ada works
This repository trains class-conditional StyleGAN2-ADA models for brain MRI, chest X-ray, and dementia MRI, serves labeled synthetic images through a Gradio demo, and evaluates image quality and classifier utility. The main workflow is dataset preparation, per-domain training or fine-tuning, periodic cross-worker weight averaging, checkpoint use for generation, and evaluation against real images. Training, evaluation, and the demo are separate entry points. Training and evaluation excerpts are partial, so some links reflect the README’s explicit workflow rather than inspected implementation details.
Main entry points: distributed_training.py Training pipeline · evaluation.py Evaluation pipeline · app.py Gradio interface
Components 20
Distributed training
- Training pipeline
distributed_training.py - Dataset loaders
distributed_training.py - Stratified sampler
distributed_training.py - Conditional GAN
distributed_training.py - ADA augmentation
distributed_training.py - Domain fine-tuning
distributed_training.py - Weight averaging
distributed_training.py - Checkpoints
distributed_training.py
Image generation
- Gradio interface
app.py - Checkpoint loader
app.py - Image generator
app.py - Generated images
app.py
Model evaluation
- Evaluation pipeline
evaluation.py - Quality metrics
evaluation.py - Classifier utility
evaluation.py - Evaluation results
evaluation.py
Other components
- Demo user
- Kaggle datasets
- HuggingFace Hub
- GPU workers
Connections 26
- Kaggle datasets → Dataset loaders: provides data
- Training pipeline → Dataset loaders: loads samples
- Dataset loaders → Stratified sampler: labels batches
- Training pipeline → Stratified sampler: uses batches
- Training pipeline → Conditional GAN: trains models
- Training pipeline → ADA augmentation: updates augmentation
- Training pipeline → Domain fine-tuning: runs phases
- Domain fine-tuning → Checkpoints: loads base weights
- Training pipeline → Weight averaging: coordinates sync
- GPU workers → Weight averaging: participate in
- Weight averaging → HuggingFace Hub: exchanges weights
- Training pipeline → Checkpoints: saves weights
- Training pipeline → HuggingFace Hub: stores checkpoints
- Demo user → Gradio interface: selects condition
- Gradio interface → Checkpoint loader: requests model
- Checkpoint loader → HuggingFace Hub: downloads checkpoint
- Checkpoint loader → Image generator: loads weights
- Gradio interface → Image generator: requests images
- Image generator → Generated images: returns images
- Evaluation pipeline → Dataset loaders: loads real data
- Evaluation pipeline → HuggingFace Hub: downloads checkpoints
- Evaluation pipeline → Conditional GAN: generates samples
- Evaluation pipeline → Quality metrics: computes FID KID IS
- Evaluation pipeline → Classifier utility: tests utility
- Quality metrics → Evaluation results: records scores
- Classifier utility → Evaluation results: records scores