Week04 Polyp Segmentation Unet Multiclass

--- license: mit tags: - image-segmentation - medical-imaging - polyp-segmentation - unet - pytorch - multi-class datasets: - Angelou0516/kvasir-seg pipeline_tag: image-segmentation --- # U-Net for Gastrointestinal Polyp Segmentation (Multi-Class) Multi-class (3-class) segmentation model trained on the [Kvasir-SEG](https://huggingface.co/datasets/Angelou0516/kvasir-seg) dataset. See also the [binary segmentation variant](https://huggingface.co/sebastiao-teixeira/week04-polyp-segmentation-unet). ## Architecture Same U-Net encoder-decoder as the binary variant, with output changed to 3 channels: - **Encoder**: 4 stages (64 → 128 → 256 → 512) + 1024-channel bottleneck. - **Decoder**: 4 upsampling stages with skip connections (512 → 256 → 128 → 64). - **Reduction head**: Conv2d(64→32) + BatchNorm + ReLU before the final 1×1 conv. - **Output**: 3-channel logits (multi-class segmentation). - Input: 3-channel RGB, 256×256. ## Loss Function **CombinedLoss** = 0.5 × CrossEntropy + 0.5 × Multi-class Dice Loss - **CrossEntropy**: replaces BCE for the multi-class case, providing per-pixel classification gradients across all classes. - **Multi-class Dice Loss**: computes per-class Dice using softmax probabilities and one-hot targets, then averages across classes. Ensures balanced optimization even with class imbalance. ## 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.9036 | — | | Test | 0.9169 | 0.3322 | ## Usage ```python import torch from model import UNet from safetensors.torch import load_file model = UNet(in_channels=3, num_classes=3) model.load_state_dict(load_file("model.safetensors")) model.eval() # input: (B, 3, 256, 256) float tensor in [0, 1] # output: (B, 3, 256, 256) logits — apply argmax(dim=1) for class map ``` Or download from Hub: ```python 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-multiclass", filename="model.safetensors", ) m...