Kvasir Polyp Segmentation (U-Net & DeepLabV3+)

Two TensorFlow/Keras models for gastrointestinal polyp segmentation on endoscopy images, trained on the Kvasir-SEG dataset augmented with normal (no-polyp) images from Kvasir v2.

Models

File Architecture Backbone Size
unet_best.h5 U-Net from scratch ~373 MB
deeplabv3plus_best.h5 DeepLabV3+ EfficientNet-B0 (ImageNet) ~40 MB

Both are saved as full Keras models (architecture + weights) in legacy HDF5 format.

Intended use

  • Input: RGB image resized to 256x256, pixel values normalized to [0, 1].
  • Output: single-channel sigmoid mask of shape (256, 256, 1). Threshold at 0.5 to obtain a binary polyp mask.

Training data

  • Positives: Kvasir-SEG — 1,000 polyp images with manually annotated segmentation masks.
  • Negatives: ~500 normal images from Kvasir v2 (normal-cecum, normal-pylorus, normal-z-line), paired with all-zero masks so the model learns that "no polyp" is a valid output and avoids false positives on healthy tissue.

Training setup

  • Loss: Dice + Binary Cross-Entropy (BCE). The BCE term keeps gradients well-behaved on empty (normal) masks, where pure Dice is unstable.
  • Optimizer: Adam, learning rate 1e-4.
  • Batch size: 8. Epochs: 15 (with ReduceLROnPlateau and EarlyStopping).
  • Metrics: Dice coefficient, IoU (on polyp images); sensitivity / specificity / false-positive rate (across polyp vs. normal images).

How to load

These models use custom loss/metric functions. For inference only, skip them with compile=False:

from tensorflow.keras.models import load_model
from huggingface_hub import hf_hub_download

path = hf_hub_download("bekmon/kvasir-polyp-segmentation", "unet_best.h5")
model = load_model(path, compile=False)

# preprocess: RGB -> 256x256 -> /255.0, shape (1, 256, 256, 3)
pred = model.predict(image_batch)        # (1, 256, 256, 1), values in [0, 1]
mask = (pred > 0.5).astype("float32")    # binary polyp mask

To resume training (or use the custom metrics), pass custom_objects with dice_bce_loss, dice_loss, bce_loss, dice_coef, and iou_coef.

Limitations

  • Trained on single-center data (Kvasir / Bærum Hospital); may not generalize to other scopes or populations.
  • Positives and negatives come from related but distinct Kvasir collections, so some dataset-bias ("shortcut learning") risk remains.
  • Not a medical device. For research and educational use only.
Downloads last month
-
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Space using bekmon/kvasir-polyp-segmentation 1