🫁 U-Net Lung Segmentation – GAN Fine-Tuning (Model 3B)

This model performs lung segmentation on grayscale chest X-ray images using a U-Net architecture enhanced with:

Dilated Residual Blocks
Squeeze-and-Excitation (SE) modules
Atrous Spatial Pyramid Pooling (ASPP)
Adversarial Fine-Tuning via PatchGAN (SegGAN)

It was specifically fine-tuned on structurally challenging cases (Hard Examples) using a GAN-based training strategy.

🧠 Model Overview

Base Model: U-Net with ASPP, SE Blocks, Dilated Convolutions
Framework: Keras / TensorFlow
Use Case: Semantic segmentation of lungs in chest X-rays
Focus: Robustness in difficult or degraded image regions

⚙️ Architecture Summary

Component	Description
Encoder	6 stages with Dilated Residual Blocks + SE modules
Bottleneck	ASPP module with dilation rates 6, 12, 18 + GAP
Decoder	Skip connections, UpSampling, Progressive channel reduction
Output Layer	1×1 Conv2D + Sigmoid for binary mask
GAN Discriminator	PatchGAN (B×H×W×1) for local realism checks

🎯 Training Strategy: SegGAN

Adversarial Setup:

Generator: U-Net (frozen encoder + ASPP)
Discriminator: PatchGAN
Loss Functions:
- combined_loss: Dice + Binary Crossentropy
- gan_loss: Binary Crossentropy with label smoothing
Optimizers:
- Generator: Adam (lr = 1e-4)
- Discriminator: Adam (lr = 1e-5, β1 = 0.5, β2 = 0.9)
EarlyStopping: Based on validation Dice stagnation

🧪 Fine-Tuning on Hard Examples (Model 3B)

To further increase robustness, the GAN-finetuned model was retrained on a curated subset of 300 low-performing masks (lowest Dice scores) out of 21,165 total samples.

These cases included:

Heavily rotated or flipped chest X-rays
Fragmented masks
Very low-contrast or noisy images

CLAHE preprocessing was applied to enhance contrast.

📊 Evaluation Results

✅ Final Test Set Performance (CLAHE Preprocessed)

Metric	Value
Average Dice Score	0.9768
Average IoU Score	0.9551
Combined Loss	0.0649

📈 Training Metrics (Epoch 50)

Epoch	Train Dice	Train IoU	Val Dice	Val IoU	Val Loss	LR
50	0.9843	0.9690	0.9603	0.9245	0.0644	5e-6

🧬 Usage

from huggingface_hub import hf_hub_download
import tensorflow as tf

# Download model
model_path = hf_hub_download(repo_id="maja011235/lung-segmentation-gan", filename="entwickeltes_masken_model.h5")

# Load with Keras
model = tf.keras.models.load_model(model_path, compile=False)

# Compile (optional)
model.compile(loss=..., metrics=...)