kvasir-seg-unet

ResNet-50 U-Net — Kvasir-SEG Polyp Segmentation

Binary segmentation model for gastrointestinal polyp detection, trained on the Kvasir-SEG dataset.

Architecture

A U-Net with a ResNet-50 encoder pretrained on ImageNet.

Feature map sizes for a 256×256 input:

stride 2 ->  64ch, 128×128
stride 4 -> 256ch, 64×64
stride 8 -> 512ch, 32×32
stride 16 -> 1024ch, 16×16
stride 32 -> 2048ch, 8×8 (bottleneck)

Loss Function

Medical image segmentation suffers from severe class imbalance (background pixels vastly outnumber polyp pixels). To address this, the model is trained with a combined loss:

Loss = 0.5 × Focal + 0.5 × Tversky

Training

• Dataset: 880 train / 60 val / 60 test (Kvasir-SEG split)
• Input size: 256×256, normalised with ImageNet mean/std
• Batch size: 8
• Epochs: 20
• LR schedule: Cosine decay with 10% linear warm-up
• Hyperparameter search: Optuna over learning_rate {5e-4, 1e-3} and weight_decay {1e-4, 1e-3}
• Best weights: snapshot of the epoch with highest validation Dice across all trials

Data Augmentation (train only)

Spatial (applied identically to image and mask): Random horizontal/vertical flip, ±20° rotation, affine (translate ±5%, scale 0.9–1.1), elastic deformation (α=60, σ=6)

Photometric (image only, p=0.8): Color jitter (brightness/contrast ±0.3, saturation ±0.2, hue ±0.05), Gaussian blur (3×3, σ∈[0.1, 2.0])

Usage

from model import ResNet50UNet
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download

model = ResNet50UNet(num_classes=1, pretrained=False)
weights_path = hf_hub_download(repo_id="VioletaR/kvasir-seg-unet", filename="model.safetensors")
model.load_state_dict(load_file(weights_path))
model.eval()

Inputs must be 256×256 RGB tensors normalised with ImageNet statistics (mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]). The model outputs raw logits; apply torch.sigmoid and threshold at 0.5 for the binary mask.