Modified Double U-Net for Multi-Class Medical Image Segmentation

This repository contains a PyTorch implementation of the Modified Double U-Net architecture, specifically optimized for multi-class segmentation of medical images. This project builds upon the Double U-Net framework by introducing ensemble-based feature extraction and multi-class spatial attention gates, tailored specifically for datasets like BUSI and CBIS-DDSM.

🚀 Key Features

• Architecture: Dual-stacked U-Net design with an advanced ensemble encoder.
• Backbones: Integration of pretrained VGG-19 (with architectural support for Xception and DenseNet feature fusion) to leverage ImageNet-learned spatial hierarchies.
• Multi-Class Support: Optimized for 3-class segmentation:
  • Class 0: Background (Normal)
  • Class 1: Benign
  • Class 2: Malignant
• Performance Optimizations:
  • Automatic Mixed Precision (AMP): High-speed training with a reduced VRAM footprint using torch.amp.
  • Persistent Workers & Memory Pinning: Optimized DataLoader settings to eliminate CPU-GPU bottlenecks and keep GPU utilization at 100%.
  • Combined Loss: A robust fusion of Cross-Entropy and Multi-Class Dice Loss (with One-Hot Encoding) to effectively handle class imbalance.

---

🏗 Architecture Overview

The Modified Double U-Net consists of two primary subnetworks:

1. Network 1: Extracts high-level features using a pretrained encoder and an ASPP (Atrous Spatial Pyramid Pooling) block to capture multi-scale contextual information.
2. Network 2: Refines the segmentation by combining the original input image with the class-aware attention maps generated by Network 1.

Crucial Modification: The bridge between networks uses a Softmax-based attention gate (summing the probabilities of the benign and malignant classes) rather than the binary Sigmoid gate used in the original architecture. This allows the second U-Net to correctly focus on multi-class foreground pixels.

---

📁 Project Structure

Our pipeline is functionally split to handle the specific preprocessing and architectural nuances of the BUSI and CBIS datasets independently.

├── dataset_busi/            # Formatted BUSI dataset (train/val splits, images, masks)
├── dataset_seg/             # Formatted CBIS-DDSM dataset (train, valid, test)
├── files/                   # Saved model checkpoints (.pth) and train logs
├── prepare_busi_dataset.py  # Preprocessing logic and mask merging for BUSI
├── prepare_cbis_dataset.py  # Preprocessing and label extraction logic for CBIS
├── train_BUSI.py            # Main training loop for BUSI (AMP & Auto-Resume)
├── train_CBIS.py            # Main training loop for CBIS (AMP & Auto-Resume)
├── test_BUSI.py             # Quantitative evaluation and benchmarking for BUSI
├── test_CBIS.py             # Quantitative evaluation and benchmarking for CBIS
├── predict_BUSI.py          # Inference & RGB overlay generation for BUSI
├── predict_CBIS.py          # Inference & RGB overlay generation for CBIS
├── BUSI_model.py            # Modified DoubleU-Net architecture mapped for BUSI
├── CBIS_model.py            # Modified DoubleU-Net architecture mapped for CBIS
├── metrics.py               # CombinedLoss and MultiClassDiceLoss functions
└── utils.py                 # Class-wise metric calculators & tensor safety checks

---

🛠 Prerequisites

Hardware

• Minimum: NVIDIA GPU with 8GB VRAM (e.g., RTX 3060/4060).
• Recommended: 12GB+ VRAM for larger batch sizes and optimal worker counts.

Environment

Ensure you have Python 3.8+ installed and your virtual environment activated.

# If using requirements.txt
pip install -r requirements.txt

(Note: Ensure you have a compatible PyTorch build for your specific CUDA architecture).

---

🏋️ Training

Both training scripts include an Auto-Resume Capability—if stopped, they will automatically detect the respective checkpoint in the files/ folder and seamlessly continue training from the last saved epoch without losing progress.

To train on the BUSI Dataset:

python train_BUSI.py

To train on the CBIS Dataset:

python train_CBIS.py

Current Training Configuration:

• Image Size: 256x256
• Batch Size: 8 (Adjustable based on VRAM)
• Optimizer: Adam (Learning Rate: 1e-4)
• Scheduler: ReduceLROnPlateau
• Early Stopping Patience: 50 epochs
• Loss Function: CrossEntropy + Multi-Class Dice Loss

---

📊 Evaluation & Inference

Quantitative Testing

Evaluate the model on the unseen test sets to generate research-grade, class-wise metrics (F1-Score, Jaccard Index/IoU, Recall, Precision). The evaluation includes tensor safety nets to prevent false metric inflation on empty classes.

# For BUSI metrics
python test_BUSI.py

# For CBIS metrics
python test_CBIS.py

Visual Prediction

Generate segmentation masks with discrete colored overlays to visually verify network performance:

• Green Overlay: Benign Tissue
• Red Overlay: Malignant Tissue

# Generate visual overlays for BUSI
python predict_BUSI.py

# Generate visual overlays for CBIS
python predict_CBIS.py

---

📜 References

• Modified Double U-Net: Sagar Deep Deb & Rajib Kumar Jha, "Modified Double U-Net Architecture for Medical Image Segmentation".
• Original Double U-Net: Jha et al., "DoubleU-Net: A Deep Convolutional Neural Network for Medical Image Segmentation".

---

🎓 Acknowledgments

Developed as part of the BTP Semester 6 project. This implementation focuses on translating theoretical architectural improvements into a high-performance, robust PyTorch pipeline tailored for clinical decision support systems.