This repository contains the source code for the master's thesis project, "Application of Deep Adversarial Networks for Synthesis and Augmentation of Medical Images." The primary goal is to generate ultra-realistic 256×256 H&E stained histopathology image patches by leveraging graph-based conditioning. The synthesized images are intended for data augmentation, particularly for addressing class imbalance in rare tissue types.
Disclaimer: This is research code and is not intended for clinical use.
- Graph-based Conditioning: Utilizes a graph representation of tissue structure (SLIC superpixels → Region Adjacency Graph → GATv2 encoder) to guide the image generation process with a structural latent vector,
z_graph. - Stable & High-Quality GAN Training: Implements relativistic losses (RpGAN), zero-center R1 regularization, and an Exponential Moving Average (EMA) of the generator weights to ensure stable training and produce high-fidelity images.
- Optimized for Small Datasets: Employs DiffAugment and optional Adaptive Discriminator Augmentation (ADA) to effectively prevent the discriminator from overfitting, a common challenge when training with limited data.
- Structure-Preserving Hook Loss: A novel GraphHook loss operates on a grid of Sobel edge weights, enforcing structural consistency between real and generated images and preserving the integrity of tissue boundaries.
- Domain-Specific Quality Control: Includes a comprehensive suite of metrics tailored for histopathology:
- Fréchet Inception Distance (FID) in a histology-specific feature space.
- Precision, Recall, Density, and Coverage (PRDC).
- H&E stain statistics validation.
- Optional texture analysis (GLCM, Moran's I) and nuclear segmentation metrics via HoVer-Net.
- Streamlined Workflow: Provides a simple command-line interface and shell scripts to manage the end-to-end workflow, from data preparation to training, generation, and quality control.
The core innovation of this project is to condition the GAN generator on both a random noise vector (z_noise) and a structural code (z_graph) derived from the morphology of a real tissue image. This dual-input approach allows the model to generate diverse yet structurally coherent images.
graph TD
subgraph "Input"
A[Real H&E Image<br/>256x256] --> B{SLIC Superpixels<br/>n~200 regions};
B --> C{Region Adjacency Graph<br/>(RAG)};
C --> D[Graph Encoder<br/>(GATv2 with mean RGB)];
D --> E[Graph Vector<br/>(z_graph)];
F[Random Noise<br/>(z_noise)];
end
subgraph "Generator"
G_IN(concat) --> G[Generator<br/>(StyleGAN-like)];
E --> G_IN;
F --> G_IN;
end
subgraph "Output & Losses"
G --> H[Generated (Fake) Image];
H --> L1[Discriminator<br/>(RpGAN)];
A --> L1;
H --> L2[GraphHook Loss<br/>(Structural Consistency)];
end
style G_IN fill:#fff,stroke:#333,stroke-width:2px
The entire process, from raw images to an augmented dataset, is managed by a series of scripts.
graph TD
A(STEP 1: Prepare Input Data<br/>e.g., stain normalization) --> B;
B(STEP 2: Cache SLIC/RAG Graphs<br/><code>scripts/prepare_graph_cache.sh</code>) --> C;
C(STEP 3: Train the GAN<br/><code>scripts/run_train.sh</code>) --> D;
D(STEP 4: Generate Patches<br/><code>scripts/run_generate.sh</code>) --> E;
E(STEP 5: Quality Control & Export<br/><code>scripts/run_qc_export.sh</code>) --> F;
F(STEP 6: Augmented Dataset<br/>Ready for use);
style A fill:#eee,stroke:#333,stroke-width:1px
style F fill:#dff,stroke:#333,stroke-width:2px
You can set up the required environment using Conda (recommended) or a standard Python virtual environment with pip.
Option A: Conda (Recommended)
# Create and activate the conda environment from the provided file
conda env create -f environment.yml
conda activate ultimate-ganOption B: pip + venv
# Create and activate a Python virtual environment
python -m venv .venv
source .venv/bin/activate
# Upgrade pip and install requirements
pip install --upgrade pip
pip install -r requirements.txt
# For advanced QC, WSI processing, etc., install optional extras
# pip install -r requirements-extras.txtNote: The
requirements.txtfile is configured for PyTorch 2.7 and CUDA 12.6. If your system has a different CUDA version, you may need to adjust the PyTorch and PyG package versions accordingly. Refer to the official PyTorch and PyG documentation for installation instructions.
Before training, you must pre-compute and cache the graph representations (SLIC superpixels and Region Adjacency Graphs) for your training images. This is a one-time process that significantly speeds up training.
bash scripts/prepare_graph_cache.sh \
-i /path/to/your/train_images \
-o /path/to/your/cache_dir \
-s 256 -n 200 -c 10 -w 8-i: Path to your input image directory.-o: Path to the output cache directory.
Configure your training run by editing config/train.yaml or by overriding parameters from the command line.
# Example training command using GPU 0
bash scripts/run_train.sh -c config/train.yaml -g 0 -- --opts batch=16 lambda_graph=0.5 use_ada=false-c: Path to the training configuration file.-g: GPU device ID to use for training.--opts: A list of key-value pairs to override settings in the YAML file.
Checkpoints and sample images will be saved to the outputs/ directory by default.
Use a trained checkpoint to generate new images. The EMA (Exponential Moving Average) checkpoint is recommended for the best results.
A) Unconditional Generation (from random z_graph)
bash scripts/run_generate.sh \
-k /path/to/your/checkpoint.pt \
-o /path/to/output_dir \
-n 1000 \
-t 0.2 \
-g 0B) Conditional Generation (from real image graphs) This method provides more control over the output morphology by using graph vectors from existing images.
bash scripts/run_generate.sh \
-k /path/to/your/checkpoint.pt \
-o /path/to/output_dir \
-n 1000 \
--from-graphs \
--data /path/to/your/train_images \
--cache /path/to/your/cache_dir-k: Path to the generator checkpoint (.ptfile).-o: Output directory for generated images.-n: Number of images to generate.-t: Discriminator score threshold to discard low-quality samples (e.g., 0.2 removes the bottom 20%).
Filter the generated images and package them for downstream use. The QC script can filter by stain statistics, de-duplicate based on similarity to real images, and apply other domain-specific checks.
bash scripts/run_qc_export.sh \
-g /path/to/generated_images \
-r /path/to/real_images \
-o /path/to/final_dataset_dir \
-e resnet_kather \
--stain-iqr \
-p 20 \
-s 256-g: Path to the generated images directory.-r: Path to the real images directory (used for de-duplication).-o: Output directory for the final, QC-passed dataset.
.
├── config/ # YAML configs for training, generation, etc.
├── scripts/ # Bash scripts to run main workflows.
├── src/ # Main Python source code.
│ ├── aug/ # Data augmentation (DiffAugment, ADA).
│ ├── cli/ # Command-line interfaces.
│ ├── data/ # Dataset loading and graph caching logic.
│ ├── hooks/ # GraphHook implementation.
│ ├── losses/ # Loss functions (RpGAN, R1).
│ ├── metrics/ # Evaluation metrics (FID, PRDC, etc.).
│ ├── models/ # G, D, and Encoder architectures.
│ └── train/ # Training loop and utilities.
├── tests/ # Smoke tests.
├── environment.yml # Conda environment file.
└── requirements.txt # Pip requirements file.
This project is licensed under the MIT License. See the LICENSE file for details.