ESRGAN Ultrasound Image Denoising - Streamlit App

Interactive web application for medical ultrasound image denoising using ESRGAN

Demo • Installation • Usage • Features

---

🎯 Overview

This Streamlit application provides an interactive interface for enhancing noisy ultrasound images using a trained ESRGAN (Enhanced Super-Resolution Generative Adversarial Network) model. The app includes comprehensive quality metrics and visualization tools for evaluating denoising performance.

Key Features

• 🖼️ Interactive Image Enhancement - Upload and denoise ultrasound images in real-time
• 📊 Comprehensive Metrics - PSNR, SSIM, SNR, Variance, Edge Strength, Entropy, Contrast
• 🔍 Visual Comparison - Side-by-side comparison of noisy, enhanced, and ground truth images
• 📈 Histogram Analysis - Pixel intensity distribution visualization
• 💾 Export Results - Download enhanced images and comparison plots
• 🚀 GPU Acceleration - Optional CUDA support for faster processing
• ⚙️ Flexible Configuration - Adjustable model parameters and inference settings

---

📋 Requirements

Python Dependencies

streamlit>=1.28.0
torch>=2.0.0
torchvision>=0.15.0
numpy>=1.24.0
Pillow>=10.0.0
matplotlib>=3.7.0

Model Checkpoint

You'll need a trained ESRGAN model checkpoint (.pt file). The app expects:

• Default path: models/esrgan_final.pt
• Model configuration: nf=48, nb=11 (adjustable in sidebar)

---

🚀 Installation

1. Clone the Repository

git clone https://github.com/TechDaDev/esrgan_streamlit.git
cd esrgan_streamlit

2. Create Virtual Environment (Recommended)

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

3. Install Dependencies

pip install -r requirements.txt

4. Add Model Checkpoint

Place your trained ESRGAN checkpoint in the models/ directory:

mkdir -p models
# Copy your checkpoint file
cp /path/to/your/esrgan_final.pt models/

Don't have a trained model? See Training Your Own Model section.

---

🎮 Usage

Launch the App

streamlit run streamlit_esrgan.py

The app will open in your default browser at http://localhost:8501

Using the Interface

1. Configure Settings (Left Sidebar)

   • Set checkpoint path
   • Adjust model parameters (nf, nb)
   • Enable/disable GPU
   • Configure display options

2. Upload Images

   • Noisy Image: The ultrasound image to enhance (required)
   • Ground Truth: Clean reference image (optional, for metrics)

3. Enhance

   • Click "🚀 Enhance Image" button
   • Wait for processing (GPU: ~0.5s, CPU: ~2-5s for 512x512 images)

4. View Results

   • Visual comparison plots
   • Quality metrics dashboard
   • Histogram analysis (optional)

5. Download

   • Enhanced image (PNG)
   • Comparison plot (PNG)

---

📊 Quality Metrics Explained

Reference-Based Metrics (require ground truth)

Metric	Description	Good Range
PSNR	Peak Signal-to-Noise Ratio	>30 dB
SSIM	Structural Similarity Index	>0.90 (0-1 scale)
SNR	Signal-to-Noise Ratio	>25 dB

No-Reference Metrics

Metric	Description	Interpretation
Variance	Detail preservation	Too low = over-smoothing
Edge Strength	Sharpness measure	Higher = sharper
Entropy	Information content	Higher = more detail
Contrast	Dynamic range	Higher = better contrast
Noise Reduction	Variance decrease	Positive = noise removed

---

🏗️ Architecture

ESRGAN Model Structure

Input (1-channel grayscale)
    ↓
Conv (3×3) → nf filters
    ↓
[RRDB Block × nb]  ← Residual-in-Residual Dense Blocks
    ↓
Trunk Conv (3×3)
    ↓
Upsampling (if scale > 1)
    ↓
HR Conv + Output Conv
    ↓
Output (1-channel, Tanh activation)

Default Configuration:

• Number of filters (nf): 48
• Number of RRDB blocks (nb): 11
• Scale: 1 (denoising mode, no upsampling)
• Parameters: ~6.2M

---

🖼️ Demo

Example Enhancement

Noisy Input → ESRGAN Enhanced → Ground Truth

Input PSNR: 22.45 dB  →  Enhanced PSNR: 30.96 dB  →  Improvement: +8.51 dB
Input SSIM: 0.742     →  Enhanced SSIM: 0.816     →  Improvement: +0.074

---

🎓 Training Your Own Model

This app requires a pre-trained ESRGAN checkpoint. To train your own model:

1. Prepare Dataset

   • Paired noisy/clean ultrasound images
   • Recommended: 100+ image pairs
   • Format: PNG, JPG (grayscale)

2. Training Script

   • Use esrgan_phase1.py (not included in this repo)
   • Configuration: nf=48, nb=11, scale=1
   • Epochs: 15 pretrain + 60 GAN training

3. Checkpoint

   • Copy final checkpoint: runs/ckpts/esrgan_final.pt
   • Place in models/ directory

Training code available at: [Your training repository link]

---

⚙️ Configuration Options

Sidebar Settings

Model Settings:

• Checkpoint Path: Path to .pt model file
• Number of Filters (nf): Must match training (default: 48)
• Number of RRDB Blocks (nb): Must match training (default: 11)
• Use GPU: Enable CUDA acceleration

Inference Settings:

• Max Image Dimension: Downscale limit (default: 1024)
  • Prevents OOM errors with large images
  • Images auto-scaled preserving aspect ratio

Evaluation Settings:

• Show Quality Metrics: Display metric dashboard
• Show Comparison Plot: Side-by-side visualization
• Show Histogram Analysis: Pixel distribution plots

---

🚨 Troubleshooting

Common Issues

1. GPU Out of Memory (OOM)

Solution: Reduce "Max Image Dimension" or disable GPU

2. Model Loading Error

Solution: Verify nf/nb match your checkpoint's training config

3. Import Error: No module named 'streamlit'

Solution: pip install -r requirements.txt

4. Checkpoint Not Found

Solution: Update checkpoint path in sidebar or add model to models/

---

📁 Repository Structure

esrgan_streamlit/
├── streamlit_esrgan.py      # Main application
├── requirements.txt          # Python dependencies
├── README.md                 # This file
├── LICENSE                   # MIT License
├── .gitignore               # Git ignore rules
├── models/                  # Model checkpoints (create this)
│   └── esrgan_final.pt     # Trained model (29MB)
└── examples/                # Example images (optional)
    ├── noisy_sample.png
    └── clean_sample.png

---

🔧 Advanced Usage

Custom Model Configuration

If your model uses different parameters:

# In sidebar, adjust:
nf = 64  # Your model's number of filters
nb = 23  # Your model's number of blocks

Batch Processing

For processing multiple images:

# See examples/batch_process.py (coming soon)

---

📝 License

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

---

🙏 Acknowledgments

• ESRGAN Architecture: Original ESRGAN Paper
• Streamlit: For the amazing web framework
• PyTorch: For deep learning infrastructure

---

📧 Contact

Author: TechDaDev
GitHub: @TechDaDev
Repository: esrgan_streamlit

---

🌟 Star History

If you find this project useful, please consider giving it a star! ⭐

---

Built with ❤️ using Streamlit and PyTorch

Back to Top