unet-kvasir-seg

ResNet34 U-Net for Kvasir-SEG

This project implements a U-Net for binary gastrointestinal polyp segmentation on Angelou0516/kvasir-seg. It uses the Hugging Face Trainer while keeping the architecture, losses, data pipeline, evaluation, and inference code in this repository.

Architecture

The encoder is an ImageNet-pretrained ResNet34 with features at five spatial scales. The custom decoder uses learned transposed convolutions, encoder-to-decoder skip connections, two-convolution blocks, and GroupNorm. A final 1×1 convolution produces one logit per pixel. The implementation accepts arbitrary spatial dimensions and restores the output to the exact input size.

The original vanilla U-Net remains in model.py only for compatibility with earlier checkpoints. ResNet34UNet is the final submitted architecture.

Data and preprocessing

The official 800/100/100 train/validation/test split is used. Images are resized to 320×320 with bilinear interpolation and normalized with ImageNet statistics. Masks use nearest-neighbor interpolation and are converted to binary tensors.

Training augmentation consists of horizontal and vertical flips plus probabilistic affine, elastic, color, and blur transformations. Image and mask geometry is transformed together through torchvision.transforms.v2.

Optimization

The decoder is warmed up for two epochs while the encoder is frozen. Full fine-tuning uses AdamW, cosine decay, gradient clipping, mixed precision, and discriminative learning rates:

• Encoder: 1e-5
• Decoder: 1e-4
• Batch size: 8
• Maximum fine-tuning epochs: 20
• Early-stopping patience: 4
• Checkpoint metric: validation Dice

BCE–Dice and BCE–Focal-Tversky were compared under an otherwise identical seed-42 experiment. BCE–Focal-Tversky was selected because it achieved higher validation Dice and foreground IoU.

Validation results

Across the three selected-loss seeds, validation Dice was 0.8732 ± 0.0219 and foreground IoU was 0.7987 ± 0.0347. Seed 42 was selected by validation Dice.

Final test result

The inference threshold was selected on validation data. Threshold 0.55 performed best. Horizontal-flip test-time augmentation was rejected because it did not improve validation Dice.

Performance by test foreground-size quartile shows the lowest Dice for the smallest quartile (0.8612) and largest quartile (0.8713). Large-polyps primarily lose recall, which indicates some under-segmentation of extensive regions.

Reproduction

Install dependencies and run tests:

uv sync
uv run python -m unittest discover -p "test_*.py" -v

Train one experiment:

uv run python experiment.py \
  --loss bce_focal_tversky \
  --seed 42 \
  --output-dir architecture-results/focal-tversky-seed42

Finalize a validation-selected checkpoint and evaluate the test split:

uv run python finalize.py \
  --checkpoint architecture-results/focal-tversky-seed42/best-model/model.safetensors \
  --output-dir architecture-results/final-model

The detailed rationale for each implementation decision is in reasoning.md.

Limitations

Kvasir-SEG contains only 1,000 images from a limited acquisition domain. Results vary across random seeds and may not generalize to other institutions, devices, or patient populations. The model is a research and teaching artifact and must not be used for clinical decisions without independent validation.