README.md

---
license: mit
tags:
- pytorch
- image-segmentation
- unet
- medical-imaging
- resnet
- polyp-segmentation
datasets:
- RGarrido03/kvasir-seg-augmented
metrics:
- dice
- iou
model-index:
- name: UNet Residual ResNet34
  results:
  - task:
      type: image-segmentation
      name: Image Segmentation
    dataset:
      name: Kvasir-SEG (augmented)
      type: RGarrido03/kvasir-seg-augmented
    metrics:
    - type: dice
      value: 0.8991
      name: Dice Score
    - type: iou
      value: 0.8167
      name: IoU Score
---

# UNet Residual ResNet34

A U-Net architecture for polyp segmentation from colonoscopy images, combining:

- **ResNet34** pretrained encoder (ImageNet weights, early layers frozen during training)
- **Residual blocks** with Squeeze-and-Excitation (SE) attention in the decoder
- **ASPP** (Atrous Spatial Pyramid Pooling) bottleneck for multi-scale context

Inspired by the [ResUNet++](https://arxiv.org/abs/1911.07067) architecture.

## Model description

The encoder is a ResNet34 backbone that extracts multi-scale features at 5 resolutions
(H, H/2, H/4, H/8, H/16) with channel counts [64, 64, 128, 256, 512]. Channel adapters
project backbone features to the decoder's expected dimensions before skip connections.

The bottleneck applies a `ResidualSEBlock` followed by ASPP (rates [1, 6, 12, 18]) and a
projection convolution, giving the model a wide receptive field without losing spatial detail.

The decoder uses four `ResidualSEBlock` stages with transposed convolutions for upsampling
and concatenated skip connections from the encoder. A final 1×1 convolution produces the
binary segmentation logit map.

| Component        | Detail |
|------------------|--------|
| Input size       | 256×256×3 |
| Output size      | 256×256×1 (logits) |
| Encoder          | ResNet34 (pretrained, ImageNet) |
| Bottleneck       | ResidualSEBlock + ASPP |
| Decoder blocks   | ResidualSEBlock + ConvTranspose2d |
| Parameters       | ~28M |

## Training

Trained on the [Kvasir-SEG](https://datasets.simula.no/kvasir-seg/) dataset with 5× random
augmentation (geometric flips/rotations + brightness/contrast/saturation jitter), giving
~5 400 training samples.

| Hyperparameter          | Value |
|-------------------------|-------|
| Epochs                  | 8 |
| Batch size              | 128 |
| Optimizer               | AdamW |
| Learning rate           | 1e-3 (linear decay) |
| Weight decay            | 0.01 |
| Loss                    | 0.5 × BCE + 0.5 × Dice |

### Training metrics

| Epoch | Train loss | Val loss | Val Dice | Val IoU |
|-------|-----------|----------|----------|---------|
| 1     | 0.4281    | 0.2964   | 0.8118   | 0.6832  |
| 2     | 0.2200    | 0.4746   | 0.6293   | 0.4591  |
| 3     | 0.1249    | 0.1628   | 0.8586   | 0.7523  |
| 4     | 0.0819    | 0.1597   | 0.8556   | 0.7477  |
| 5     | 0.0735    | 0.1330   | 0.8815   | 0.7881  |
| 6     | 0.0577    | 0.1296   | 0.8873   | 0.7975  |
| 7     | 0.0502    | 0.1217   | 0.8919   | 0.8049  |
| **8** | **0.0399**| **0.1127**| **0.8991**| **0.8167**|

## Usage

```python
import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

# These files are also available in the repository
from models.backbones import ResNet34Backbone
from models.residual_unet import ResidualUNet

# Load model
backbone = ResNet34Backbone(pretrained=False)
model = ResidualUNet(in_channels=3, num_classes=1, backbone=backbone)

weights_path = hf_hub_download(
    repo_id="RGarrido03/unet-residual-resnet34",
    filename="model.safetensors",
)
model.load_state_dict(load_file(weights_path), strict=False)
model.eval()

# Inference
from torchvision.transforms import Compose, Resize, ToTensor
from PIL import Image

transform = Compose([Resize((256, 256)), ToTensor()])
image = transform(Image.open("your_image.jpg").convert("RGB")).unsqueeze(0)

with torch.no_grad():
    logits = model(image)          # (1, 1, 256, 256)
    mask = (logits.sigmoid() > 0.5).squeeze()  # binary mask
```

## Citation

If you use this model, please also cite the Kvasir-SEG dataset:

```bibtex
@inproceedings{Jha2020,
  title     = {Kvasir-{SEG}: A Segmented Polyp Dataset},
  author    = {Jha, Debesh and Smedsrud, Pia H. and Riegler, Michael A. and
               Halvorsen, P{\aa}l and de Lange, Thomas and Johansen, Dag and
               Johansen, H{\aa}vard D.},
  booktitle = {MultiMedia Modeling},
  year      = {2020},
}
```