UNet Image Segmentation

A deep learning application for semantic image segmentation using UNet architecture with ResNet34 encoder.

Overview

This project implements a U-Net model for semantic segmentation on the PASCAL VOC 2012 dataset. The model uses a ResNet34 backbone with ImageNet pre-trained weights to segment images into 21 different classes.

Features

• UNet architecture with ResNet34 encoder
• 21-class semantic segmentation
• FastAPI backend for model inference
• Streamlit frontend for user interaction
• Real-time image segmentation

Architecture

The model consists of:

• Encoder: ResNet34 (pre-trained on ImageNet)
• Decoder: UNet decoder with skip connections
• Segmentation head: 21-class output

Dataset

• Source: PASCAL VOC 2012 dataset
• Size: 3.52GB
• Images: 2,920 training images
• Classes: 21 semantic classes

Model Specifications

• Input shape: (batch_size, 3, 224, 224)
• Output shape: (batch_size, 21, 224, 224)
• Total parameters: 24.4M
• Trainable parameters: 3.2M

Setup

1. Install dependencies:

pip install uv
cd UNet_image_segmentation
uv sync

2. Start the backend API:

cd backend
uvicorn main:app --reload --port 8000

3. Run the frontend:

cd frontend
streamlit run main.py

Usage

1. Upload an image through the Streamlit interface
2. Click 'Predict' to perform segmentation
3. View the segmented output mask

Mathematical Details

For detailed mathematical formulations and theoretical background, please refer to the DOCUMENTATION.md file.

Technologies

• PyTorch
• FastAPI
• Streamlit
• Segmentation Models PyTorch
• ResNet34
• albumentation
• loss function: DiceLoss (multiclass) + CrossEntropyLoss
• optimizer: Adam
• Batchsize trained: 32
• fine tuned