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++ 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.

Training

Trained on the Kvasir-SEG dataset with 5× random augmentation (geometric flips/rotations + brightness/contrast/saturation jitter), giving ~5 400 training samples.

Training metrics

Usage

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:

@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},
}