U-Net for Polyp Segmentation on Kvasir-SEG

Dataset

• Dataset: Angelou0516/kvasir-seg (train/validation/test splits)
• Task: binary image segmentation

Preprocessing

• RGB image: resized to 256×256 (BILINEAR) + ToTensor() → [0, 1]
• Mask: resized to 256×256 (NEAREST) + binarized to 0/1

Training Augmentation

• Horizontal flip (p = 0.5)
• Light affine transformation: rotation ±10°, small translation, and scaling
• Elastic deformation (p = 0.2, alpha = 4, sigma = 4)

Architecture

• 4 encoder levels with MaxPool2d
• bottleneck
• decoder with ConvTranspose2d and skip connections (concatenation)
• DoubleConv blocks at each stage (Conv2d + BatchNorm2d + ReLU, repeated twice)

Loss Function and Rationale

• BCEWithLogitsLoss for stable pixel-wise supervision
• Dice loss to optimize region overlap and better handle class imbalance

• BCE / Dice = 0.5 / 0.5

Training Configuration

• Batch size: 4
• Learning rate: 1e-4
• Epochs: 25
• Best checkpoint selected using eval_dice (load_best_model_at_end=True)

Results

• train_dice: 0.8733, train_iou: 0.7995
• eval_dice: 0.7913, eval_iou: 0.7016
• test_dice: 0.8191, test_iou: 0.7294