U-Net model for Polyp Segmentation (Kvasir-SEG)
Model description
This repository contains a U-Net model trained for gastrointestinal polyp segmentation using the Kvasir-SEG dataset.
The model performs binary semantic segmentation, predicting a segmentation mask that identifies the pixels corresponding to polyps in a gastrointestinal image.
The architecture follows the standard encoder-decoder structure with skip connections:
- Input shape: 3 x H x W (RGB)
- Backbone: convolutional blocks with Batch normalization and ReLU activation
- Downsampling: MaxPooling
- Upsampling: transposed convolutions
- Output shape: 1 x H x W
The output mask represents the predicted probability oh each pixel belonging to a polyp class.
Binarization uses a default threshold of 0.5.
Intended uses and limitations
This architecture is intended for educational purposes and for experimenting with the U-Net model for medical image segmentation.
Although the model performed well, there are some limitations: it was trained with only one dataset (Kvasir-SEG) and it may have a worse performance with images from different gastrointestinal datasets.
It's not for clinical use.
How to use
The model weights are provided in the file model.safetensors.
Example usage in PyTorch:
1from safetensors.torch import load_file
2from model import UNet
3
4model = UNet(in_channels=3, n_classes=1)
5
6weights = load_file("model.safetensors")
7model.load_state_dict(weights)
8
9model.eval()
Expected input shape: (B,3,H,W)
Expected output shape: (B,1,H,W)
Training data
The model was trained using the Kvasir-SEG dataset for polyp segmentation.
The dataset is split into: train (model fitting), validation (model selection) and test (reserved for final reporting, not executed yet).
The train dataset has 800 images, each one with a corresponding binary segmentation mask.
Training procedure
The training procedure was divided in:
-
Preprocessing:
- images were resized to a fixed resolution (256)
- convert images to RGB and masks to grayscale
- images were converted to tensors
- masks where binarized using a threshold greater than 0.5
- data augmentation applied to training set only (random horizontal flip)
-
Training configuration (using PyTorch + HuggingFace Trainer):
- U-Net architecture (in_channels=3, n_classes=1)
- Loss Function: DiceLoss * 0.5 + BCEWithLogitsLoss * 0.5
- Number of epochs: 20
- Learning rate: 3e-4
- Batch size: 8
- Model selection: metric_for_best_model = "mean_dice"
Loss function rationale
The model was trained using a combination of Dice Loss and Binary Cross-Entropy with Logits.
Dice Loss can handle segmentation tasks with class imbalance as it focuses on maximazing the overlap between the predicted and the ground truth mask, but it can produce unstable gradients when the predicted regions are very small.
BCE helps stabilize the training process by providing pixel-wise supervision and improving convergence, but it can't handle class imbalance very well.
By combining both loss functions, there is a balance between region overlap optimization (Dice) with pixel-level classification accuracy (BCE), which may lead to better segmentation performances.
Variable and metrics
The model performance was evaluated using the the Dice coefficient and the IoU.
