U-Net for Gastrointestinal Polyp Segmentation (Kvasir-SEG)
Model Description
This model is a U-Net convolutional neural network trained for binary segmentation of gastrointestinal polyps on the Kvasir-SEG dataset.
Architecture
The model follows the classic U-Net architecture with the following components:
- Encoder: 4 levels of DoubleConv blocks (Conv2d → BatchNorm → ReLU → Conv2d → BatchNorm → ReLU) followed by MaxPool2d, progressively increasing channels from 3 → 64 → 128 → 256 → 512 while halving spatial resolution.
- Bottleneck: A DoubleConv block at the lowest resolution (512 → 1024 channels, 16×16 spatial).
- Decoder: 4 levels of transposed convolutions (ConvTranspose2d) for upsampling, each followed by concatenation with the corresponding encoder skip connection and a DoubleConv block.
- Skip Connections: Feature maps from each encoder level are concatenated with the corresponding decoder level to preserve spatial information.
- Output: A 1×1 Conv2d reducing to 1 channel, producing a binary segmentation mask of shape [B, 1, H, W].
Loss Function
Training uses a combined BCE + Dice Loss:
- BCEWithLogitsLoss: Provides stable pixel-wise binary cross-entropy.
- Dice Loss: Directly optimizes overlap between prediction and ground truth, making it robust to the class imbalance present in this dataset (≈85% background, ≈15% polyp).
Training
- Dataset: Angelou0516/kvasir-seg
- Train/Val/Test split: 800 / 100 / 100
- Image size: 256×256
- Batch size: 16
- Epochs: 20
- Optimizer: AdamW (lr=1e-4)
Results
| Metric | Score |
|---|
| IoU | 0.6895 |
| Accuracy | 0.9374 |
