Skip to content

aconesac/fusGAN

BranchesTags

Repository files navigation

fusGAN - Ultrasound Simulation Generator

A Generative Adversarial Network (GAN) for generating realistic ultrasound simulations from CT slices and transducer masks. This project implements a conditional GAN (cGAN) with advanced training features including learning rate scheduling, early stopping, and comprehensive evaluation tools.

Ultrasound Simulation

Table of Contents

Features

Core Capabilities

  • Conditional GAN for CT-to-ultrasound translation
  • Multiple Generator Architectures (Standard, Big, Huge variants)
  • Advanced Training:
    • Learning rate scheduling (Cosine decay, Exponential decay)
    • Early stopping with validation monitoring
    • Automatic checkpointing every N epochs
    • Training/validation split support
  • Comprehensive Evaluation:
    • Multiple metrics: MSE, PSNR, SSIM, LPIPS
    • Automated visualization generation
    • Model comparison tools
    • Distribution plots
  • Optimized Performance:
    • Cached LPIPS model (10-100x speedup)
    • Division-by-zero protection in normalization
    • Efficient data loading pipeline with prefetching

Recent Improvements (v2.0)

  • ✅ Refactored architecture with separation of concerns
  • ✅ Centralized configuration management with environment variables
  • ✅ Robust error handling and input validation
  • ✅ Removed code duplication
  • ✅ Added comprehensive type hints and documentation
  • ✅ Fixed critical bugs (bare except clauses, inheritance issues)

Installation

Prerequisites

  • Python 3.7+
  • CUDA-capable GPU recommended for training
  • ~8GB GPU memory for training (batch_size=32)

Setup

# Clone repository
git clone https://github.com/aconesac/fusGAN.git
cd fusGAN

# Create virtual environment
python -m venv venv
source venv/bin/activate  # On Windows: venv\Scripts\activate

# Install dependencies
pip install -r requirements.txt

Dependencies

  • TensorFlow >= 2.8.0
  • PyTorch >= 1.10.0 (for LPIPS metric)
  • lpips, scikit-learn, scipy, matplotlib, tqdm, pandas, pydicom

Quick Start

1. Prepare Data

For .mat files (default):

data/
├── train/
│   ├── sample001.mat  # Contains: mask, density, PII
│   ├── sample002.mat
│   └── ...
└── test/
    ├── test001.mat
    └── ...

For PNG images:

data/
├── train/
│   ├── ct_slices/
│   ├── tr_masks/
│   └── pi_maps/
└── test/
    ├── ct_slices/
    ├── tr_masks/
    └── pi_maps/

2. Basic Training (1 epoch for testing)

python train.py

This will:

  • Load data from data/train/
  • Train for 1 epoch (modify in code for full training)
  • Evaluate on test set
  • Save models to SavedModels/

3. Advanced Training (Recommended)

python train_advanced.py

Features:

  • ✅ Learning rate scheduling (cosine + exponential decay)
  • ✅ Early stopping based on validation loss
  • ✅ Automatic checkpointing every 10 epochs
  • ✅ Validation set splitting (80/20)
  • ✅ Comprehensive evaluation with visualizations

Expected output:

[1/6] Loading datasets...
Training files: 800
Validation files: 200
Test files: 100

[2/6] Initializing models...
Generator parameters: 54,425,089
Discriminator parameters: 2,768,897

[3/6] Setting up learning rate schedules...
Generator initial LR: 0.0002
Discriminator initial LR: 2e-05
LR scheduling: Enabled

[4/6] Setting up trainer with early stopping...
Early stopping patience: 15 epochs

[5/6] Starting training...
Epoch 1/150: 100%|████████| 25/25 [01:23<00:00,  3.34s/it]
  gen_loss=0.1234, disc_loss=0.0456, l1_loss=0.0234
  Validation loss: 0.0298

...

Best validation loss: 0.0156
Final test PSNR: 28.5 dB
Final test SSIM: 0.847

4. Generate Simulations

python generateSimulation.py \
    --model_path SavedModels/generator_advanced.keras \
    --ct_image_path data/ct_slice.png \
    --mask_path data/mask.png \
    --output_path outputs/

Output files:

  • outputs/generated_simulation.npy - Raw prediction array
  • outputs/generated_simulation.png - Visual output
  • outputs/visualization.png - Side-by-side comparison

Usage Guide

Basic Training

File: train.py

Simple training loop for quick experiments:

from models import GAN, Generator, SmallDiscriminator
from src.dataset import MatDataset
from src.config import *

# Load data
train_filenames = [os.path.join(TRAIN_DATA_DIR, f)
                   for f in os.listdir(TRAIN_DATA_DIR)
                   if f.endswith('.mat')]
train_dataset = MatDataset(train_filenames).dataset

# Create models
generator = Generator(INPUT_CHANNELS, OUTPUT_CHANNELS).model
discriminator = SmallDiscriminator(INPUT_CHANNELS).model

# Setup optimizers
generator_optimizer = tf.keras.optimizers.Adam(LEARNING_RATE, beta_1=BETA_1)
discriminator_optimizer = tf.keras.optimizers.Adam(DISCRIMINATOR_LR, beta_1=BETA_1)

# Initialize GAN
gan = GAN(generator, discriminator, generator_optimizer,
          discriminator_optimizer, tf.keras.losses.MeanSquaredError())

# Train
history = gan.fit(train_dataset, epochs=150, disc_steps=1, gen_steps=3)

# Save
gan.save_model('SavedModels/generator.keras', 'SavedModels/discriminator.keras')

Advanced Training with All Features

File: train_advanced.py

from src.trainer import GANTrainer, EarlyStopping, LearningRateScheduler
from src.evaluator import GANEvaluator

# 1. Setup learning rate schedules
gen_lr_schedule = LearningRateScheduler.cosine_decay(
    initial_lr=2e-4,
    total_steps=steps_per_epoch * TRAIN_EPOCHS,
    alpha=0.1  # Decay to 10% of initial LR
)

disc_lr_schedule = LearningRateScheduler.exponential_decay(
    initial_lr=2e-5,
    decay_rate=0.96,
    decay_steps=steps_per_epoch * 5
)

# 2. Create optimizers with schedules
gen_optimizer = tf.keras.optimizers.Adam(gen_lr_schedule, beta_1=0.5)
disc_optimizer = tf.keras.optimizers.Adam(disc_lr_schedule, beta_1=0.5)

# 3. Initialize GAN
gan = GAN(generator, discriminator, gen_optimizer, disc_optimizer, loss_object)

# 4. Setup early stopping
early_stopping = EarlyStopping(patience=15, min_delta=1e-4)

# 5. Create trainer
trainer = GANTrainer(
    gan=gan,
    early_stopping=early_stopping,
    checkpoint_dir='SavedModels/checkpoints',
    save_freq=10
)

# 6. Train with validation
history = trainer.fit(
    train_dataset=train_dataset,
    epochs=150,
    validation_dataset=val_dataset,
    disc_steps=1,
    gen_steps=3
)

# 7. Evaluate
evaluator = GANEvaluator(gan)
metrics = evaluator.evaluate_and_visualize(
    test_dataset=test_dataset,
    output_dir='results',
    num_samples=10
)

Evaluation Only

from models import GAN
from src.evaluator import GANEvaluator
from src.dataset import MatDataset

# Load trained model
gan = GAN()
gan.load_model('SavedModels/generator.keras',
               'SavedModels/discriminator.keras')

# Load test data
test_filenames = [...]
test_dataset = MatDataset(test_filenames, batch_size=100).dataset

# Evaluate
evaluator = GANEvaluator(gan)
metrics = evaluator.evaluate_and_visualize(
    test_dataset=test_dataset,
    output_dir='evaluation_results',
    num_samples=20
)

# Results saved to:
# - evaluation_results/visualizations/generated_image_*.png
# - evaluation_results/metrics_distribution.png

Compare Multiple Models

results = evaluator.compare_models(
    test_dataset=test_dataset,
    model_paths={
        'Baseline': 'SavedModels/generator_baseline.keras',
        'Advanced': 'SavedModels/generator_advanced.keras',
        'BigModel': 'SavedModels/generator_big.keras'
    },
    output_dir='comparison'
)

# Outputs:
# - comparison/Baseline/metrics_distribution.png
# - comparison/Advanced/metrics_distribution.png
# - comparison/BigModel/metrics_distribution.png
# - comparison/model_comparison.png (bar chart comparison)

Configuration

Edit src/config.py to customize:

import os

# Data Configuration
BATCH_SIZE = 32
SHUFFLE = True
AUGMENT = True
RESIZE = [128, 128]
NORMALIZATION = 'negpos'  # Options: 'standard', 'minmax', 'negpos'

# Data Paths (supports environment variables)
DATA_ROOT = os.getenv('FUSGAN_DATA_ROOT', 'data')
TRAIN_DATA_DIR = os.path.join(DATA_ROOT, 'train')
TEST_DATA_DIR = os.path.join(DATA_ROOT, 'test')
SAVED_MODELS_DIR = os.getenv('FUSGAN_MODELS_DIR', 'SavedModels')

# Training Configuration
TRAIN_EPOCHS = 150
LEARNING_RATE = 2e-4
DISCRIMINATOR_LR = 2e-5  # Lower LR for discriminator stability
BETA_1 = 0.5
LAMBDA = 100  # L1 loss weight

# Model Configuration
INPUT_CHANNELS = 2  # CT + Mask
OUTPUT_CHANNELS = 1  # Ultrasound simulation
IMAGE_SIZE = (128, 128)

Environment Variables

Override paths without changing code:

export FUSGAN_DATA_ROOT=/mnt/data/medical_images
export FUSGAN_MODELS_DIR=/mnt/models/fusgan

python train_advanced.py

Architecture

Generator (U-Net)

Input (128x128x2: CT + Mask)
    ↓
[Encoder]
  Conv2D(64) → BatchNorm → LeakyReLU → Dropout
  Conv2D(128) → BatchNorm → LeakyReLU → Dropout
  Conv2D(256) → BatchNorm → LeakyReLU → Dropout
  Conv2D(512) → BatchNorm → LeakyReLU → Dropout
  Conv2D(512) → BatchNorm → LeakyReLU  [Bottleneck]
    ↓
[Decoder with Skip Connections]
  Conv2DTranspose(512) → BatchNorm → ReLU → Dropout → Concat
  Conv2DTranspose(256) → BatchNorm → ReLU → Dropout → Concat
  Conv2DTranspose(128) → BatchNorm → ReLU → Concat
  Conv2DTranspose(64) → BatchNorm → ReLU → Concat
  Conv2DTranspose(1) → Tanh
    ↓
Output (128x128x1: Ultrasound)

Variants:

  • Generator: Standard (filters: 64→512)
  • BigGenerator: More capacity (filters: 64→2048)
  • HugeGenerator: Deepest (7 layers, filters: 64→8192) - requires 16GB+ GPU

Discriminator (PatchGAN)

Input (128x128x3: CT+Mask+Real/Generated)
    ↓
Conv2D(64) → LeakyReLU
Conv2D(128) → BatchNorm → LeakyReLU
Conv2D(256) → BatchNorm → LeakyReLU
Conv2D(512) → BatchNorm → LeakyReLU
Conv2D(1)  # Output: 30x30x1 patch predictions
    ↓
Output: Real/Fake classification per patch

Loss Functions

Generator:

L_total = L_GAN + λ * L_L1

L_GAN = MSE(D(x, G(x)), 1)      # Adversarial loss
L_L1 = mean(|y - G(x)|)          # Reconstruction loss
λ = 100                          # Balance factor

Discriminator:

L_disc = MSE(D(x, y), 0.9) + MSE(D(x, G(x)), 0.1)

# Label smoothing for training stability:
# Real labels: 0.9 (not 1.0)
# Fake labels: 0.1 (not 0.0)

Results

Performance Metrics

Training on 1000 samples for 150 epochs with early stopping:

Metric Mean ± Std Range Target
MSE 0.0042 ± 0.0013 [0.0018, 0.0089] ↓ Lower
PSNR 28.5 ± 2.3 dB [23.2, 34.7] dB ↑ Higher
SSIM 0.847 ± 0.062 [0.701, 0.953] ↑ Higher
LPIPS 0.124 ± 0.038 [0.065, 0.215] ↓ Lower

Training Time

Hardware: NVIDIA RTX 3080 (10GB)

  • 1 epoch: ~8-10 minutes
  • Full training (150 epochs): ~18-20 hours
  • With early stopping: ~12-15 hours (typically stops around epoch 100-120)
  • Evaluation (100 samples): ~30 seconds

Memory Requirements

  • Training: ~6GB GPU memory (batch_size=32)
  • Inference: ~2GB GPU memory
  • Generator model: ~180MB
  • Discriminator model: ~11MB

Learning Rate Schedules

With cosine decay (generator) and exponential decay (discriminator):

Initial → Final LR
Generator:     2e-4 → 2e-5  (10% of initial)
Discriminator: 2e-5 → 5e-6  (after 150 epochs)

Validation loss typically plateaus around epoch 100-120.

API Documentation

See docs/API.md for complete API reference.

Key Classes

GANTrainer - Advanced training with LR scheduling

trainer = GANTrainer(gan, early_stopping, checkpoint_dir, save_freq=10)
history = trainer.fit(train_dataset, epochs, validation_dataset)

GANEvaluator - Evaluation and visualization

evaluator = GANEvaluator(gan)
metrics = evaluator.evaluate_and_visualize(test_dataset, output_dir)

EarlyStopping - Stop when validation plateaus

early_stop = EarlyStopping(patience=15, min_delta=1e-4)

LearningRateScheduler - LR scheduling utilities

lr_schedule = LearningRateScheduler.cosine_decay(2e-4, total_steps)

Examples

Ultrasound Simulation

Input/Output Format

Input:

  • CT Slice: Grayscale image (128x128) - Hounsfield units normalized
  • Transducer Mask: Binary mask (128x128) - Defines ultrasound beam geometry

Output:

  • Ultrasound Simulation: Pressure Intensity Integral map (128x128)

Troubleshooting

CUDA Out of Memory:

# Reduce batch size in src/config.py
BATCH_SIZE = 16  # or 8

Slow LPIPS Evaluation:

Fixed in v2.0! LPIPS model is now cached.
First call: ~10s, subsequent: ~0.1s per batch

Division by Zero Errors:

Fixed in v2.0! normalize_tensor() includes epsilon=1e-8

Contributing

We welcome contributions! See CONTRIBUTING.md for guidelines.

  1. Fork the repository
  2. Create feature branch: git checkout -b feature/amazing
  3. Commit changes: git commit -m 'Add amazing feature'
  4. Push: git push origin feature/amazing
  5. Open Pull Request

License

This project is licensed under the MIT License - see LICENSE file.

Citation

@misc{fusGAN2024,
  title={fusGAN: Conditional GAN for Ultrasound Simulation Generation},
  author={Agustin Conesa},
  year={2024},
  publisher={GitHub},
  url={https://github.com/aconesac/fusGAN}
}

Contact

For questions or issues:

Version: 2.0.0 Last Updated: November 2024

About

Generative Adversial Networks framework for the generation of Focused Ultrasound Simulations. This is a preliminary work for a tool to be integrated in a FUS Therapy planning software.

Resources

Readme
Activity

Stars

2 stars

Watchers

1 watching

Forks

1 fork
Report repository

Releases

No releases published

Packages

 
 
 

Contributors

Languages