Week04 Polyp Segmentation Unet
Binary segmentation model trained on the Kvasir-SEG dataset for polyp detection in gastrointestinal endoscopy images.
Model source
- Kind
- Other
- File size
- 119 MB
- Version
- vRevision 90ee276
- Updated
- Sep 27, 2026
- Downloads
- 0
- Tasks
- Image segmentation
- Source
- Hugging Face
Source description
U-Net for Gastrointestinal Polyp Segmentation
Binary segmentation model trained on the Kvasir-SEG dataset for polyp detection in gastrointestinal endoscopy images.
Also available: 3-class multi-class variant — same architecture with 3 output classes (val_dice=0.9036, test_dice=0.9169)
Architecture
Standard U-Net encoder-decoder with skip connections:
- Encoder: 4 stages of double 3×3 conv blocks (64 → 128 → 256 → 512) with MaxPool2d downsampling, followed by a 1024-channel bottleneck.
- Decoder: 4 stages of transposed convolution upsampling with skip concatenation and double conv blocks (512 → 256 → 128 → 64).
- Reduction head: extra Conv2d(64→32) + BatchNorm + ReLU before the final 1×1 convolution to reduce capacity and mitigate overfitting.
- Output: 1-channel logits (binary segmentation).
- All conv blocks use BatchNorm2d + ReLU activation.
- Input: 3-channel RGB, 256×256.
Loss Function
CombinedLoss = 0.5 × BCE + 0.5 × Dice Loss
- BCE (Binary Cross-Entropy with Logits): provides stable pixel-level gradients and handles class imbalance through its log-based formulation.
- Dice Loss: directly optimizes the overlap metric (Dice coefficient), which is the standard evaluation metric for medical segmentation. BCE alone tends to underperform on imbalanced masks where background dominates.
- The equal weighting balances pixel-wise accuracy (BCE) with region-level overlap (Dice).
Training Setup
| Parameter | Value |
|---|---|
| Epochs | 20 |
| Optimizer | Adam |
| Learning rate | 1e-4 |
| LR scheduler | ReduceLROnPlateau (factor=0.5, patience=3, min_lr=1e-6) |
| Batch size | 8 |
| Image size | 256×256 |
| GPU | NVIDIA GeForce RTX 3080 Ti |
| Dataset splits | train=800, val=100, test=100 |
Metrics
| Split | Dice Coefficient | Loss |
|---|---|---|
| Validation (best) | 0.7747 | — |
| Test | 0.8114 | 0.2661 |
Usage
import torch
from model import UNet
from safetensors.torch import load_file
model = UNet(in_channels=3, num_classes=1)
model.load_state_dict(load_file("model.safetensors"))
model.eval()
# input: (B, 3, 256, 256) float tensor in [0, 1]
# output: (B, 1, 256, 256) logits — apply sigmoid + threshold for mask
Or download from Hub:
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from model import UNet
weights = hf_hub_download(
repo_id="sebastiao-teixeira/week04-polyp-segmentation-unet",
filename="model.safetensors",
)
model = UNet(in_channels=3, num_classes=1)
model.load_state_dict(load_file(weights))
model.eval()
Sources
1 source
Sources
1 sourceVerified Sep 27
Model artifacts
1 artifact
Model artifacts
1 artifactSource excerpts
2 excerpts
Source excerpts
2 excerpts--- license: mit tags: - image-segmentation - medical-imaging - polyp-segmentation - unet - pytorch datasets: - Angelou0516/kvasir-seg pipeline_tag: image-segmentation --- # U-Net for Gastrointestinal Polyp Segmentation Binary segmentation model trained on the [Kvasir-SEG](https://huggingface.co/datasets/Angelou0516/kvasir-seg) dataset for polyp detection in gastrointestinal endoscopy images. > Also available: [3-class multi-class variant](https://huggingface.co/sebastiao-teixeira/week04-polyp-segmentation-unet-multiclass) — same architecture with 3 output classes (val_dice=0.9036, test_dice=0.9169) ## Architecture Standard U-Net encoder-decoder with skip connections: - **Encoder**: 4 stages of double 3×3 conv blocks (64 → 128 → 256 → 512) with MaxPool2d downsampling, followed by a 1024-channel bottleneck. - **Decoder**: 4 stages of transposed convolution upsampling with skip concatenation and double conv blocks (512 → 256 → 128 → 64). - **Reduction head**: extra Conv2d(64→32) + BatchNorm + ReLU before the final 1×1 convolution to reduce capacity and mitigate overfitting. - **Output**: 1-channel logits (binary segmentation). - All conv blocks use BatchNorm2d + ReLU activation. - Input: 3-channel RGB, 256×256. ## Loss Function **CombinedLoss** = 0.5 × BCE + 0.5 × Dice Loss - **BCE (Binary Cross-Entropy with Logits)**: provides stable pixel-level gradients and handles class imbalance through its log-based formulation. - **Dice Loss**: directly optimizes the overlap metric (Dice coefficient), which is the standard evaluation metric for medical segmentation. BCE alone tends to underperform on imbalanced masks where background dominates. - The equal weighting balances pixel-wise accuracy (BCE) with region-level overlap (Dice). ## Training Setup | Parameter | Value | |---|---| | Epochs | 20 | | Optimizer | Adam | | Learning rate | 1e-4 | | LR scheduler | ReduceLROnPlateau (factor=0.5, patience=3, min_lr=1e-6) | | Batch size | 8 | | Image size | 256×256 | | GPU | NVIDIA GeForce RTX 3080 Ti | | Dataset splits | train=800, val=100, test=100 | ## Metrics | Split | Dice Coefficient | Loss | |---|---|---| | Validation (best) | 0.7747 | — | | Test | 0.8114 | 0.2661 | ## Usage ```python import torch from model import UNet from safetensors.torch import load_file model = UNet(in_channels=3, num_classes=1) model.load_state_dict(load_file("model.safet...
Source context: 0 downloads · 0 likes · Pipeline image-segmentation · Repo sebastiao-teixeira/week04-polyp-segmentation-unet
