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.

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.

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.

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:

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.

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.

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.

Install dependencies and run tests:

Train one experiment:

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

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

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.