Image Segmentation

Project

Table of contents

• Overview
• Details of Dataset
• Data Preparation
• Architecture of Network
• Training and Test
• Evaluation
• Additional Sections
• Conclusion
• Resources
• Project repository

---

Overview

This project demonstrates how to implement an image segmentation model using the U-Net architecture. The goal is to segment medical images, such as brain MRI scans, to identify regions of interest (e.g., lesions or abnormalities). The workflow includes data loading, preprocessing (including augmentation and normalization), model definition, training, evaluation, and visualization.

---

Details of Dataset

• Dataset Used: Brain MRI dataset for FLAIR abnormality segmentation
• Data Format: 2D image slices (often .tif files) and corresponding masks
• Channels: 3 (input images)
• Segmentation Classes: 1 (binary mask: abnormality vs. background)

---

Data Preparation

Preprocessing Steps:

• Cropping to the smallest enclosing volume
• Padding to square shape
• Resizing to a fixed size (e.g., 256x256)
• Normalization (channel-wise mean/std)

Augmentation:

• Random scaling, rotation, and horizontal flipping are applied to increase data diversity and improve model robustness.

from torchvision.transforms import Compose
from skimage.transform import rescale, rotate
import numpy as np

# Example augmentation pipeline
def transforms(scale=None, angle=None, flip_prob=None):
    transform_list = []
    if scale is not None:
        transform_list.append(Scale(scale))
    if angle is not None:
        transform_list.append(Rotate(angle))
    if flip_prob is not None:
        transform_list.append(HorizontalFlip(flip_prob))
    return Compose(transform_list)

---

Architecture of Network

U-Net Architecture:

• Encoder-decoder structure with skip connections
• Each block consists of two convolutional layers, batch normalization, and ReLU activations
• Downsampling via max pooling; upsampling via transposed convolutions
• Final layer uses a 1x1 convolution and sigmoid activation for binary segmentation

import torch
import torch.nn as nn
from collections import OrderedDict

class UNet(nn.Module):
    def __init__(self, in_channels=3, out_channels=1, init_features=32):
        super(UNet, self).__init__()
        features = init_features
        self.encoder1 = UNet._block(in_channels, features, name="enc1")
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.encoder2 = UNet._block(features, features * 2, name="enc2")
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.encoder3 = UNet._block(features * 2, features * 4, name="enc3")
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.encoder4 = UNet._block(features * 4, features * 8, name="enc4")
        self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.bottleneck = UNet._block(features * 8, features * 16, name="bottleneck")
        self.upconv4 = nn.ConvTranspose2d(features * 16, features * 8, kernel_size=2, stride=2)
        self.decoder4 = UNet._block((features * 8) * 2, features * 8, name="dec4")
        self.upconv3 = nn.ConvTranspose2d(features * 8, features * 4, kernel_size=2, stride=2)
        self.decoder3 = UNet._block((features * 4) * 2, features * 4, name="dec3")
        self.upconv2 = nn.ConvTranspose2d(features * 4, features * 2, kernel_size=2, stride=2)
        self.decoder2 = UNet._block((features * 2) * 2, features * 2, name="dec2")
        self.upconv1 = nn.ConvTranspose2d(features * 2, features, kernel_size=2, stride=2)
        self.decoder1 = UNet._block(features * 2, features, name="dec1")
        self.conv = nn.Conv2d(in_channels=features, out_channels=out_channels, kernel_size=1)
    def forward(self, x):
        enc1 = self.encoder1(x)
        enc2 = self.encoder2(self.pool1(enc1))
        enc3 = self.encoder3(self.pool2(enc2))
        enc4 = self.encoder4(self.pool3(enc3))
        bottleneck = self.bottleneck(self.pool4(enc4))
        dec4 = self.upconv4(bottleneck)
        dec4 = torch.cat((dec4, enc4), dim=1)
        dec4 = self.decoder4(dec4)
        dec3 = self.upconv3(dec4)
        dec3 = torch.cat((dec3, enc3), dim=1)
        dec3 = self.decoder3(dec3)
        dec2 = self.upconv2(dec3)
        dec2 = torch.cat((dec2, enc2), dim=1)
        dec2 = self.decoder2(dec2)
        dec1 = self.upconv1(dec2)
        dec1 = torch.cat((dec1, enc1), dim=1)
        dec1 = self.decoder1(dec1)
        return torch.sigmoid(self.conv(dec1))
    @staticmethod
    def _block(in_channels, features, name):
        return nn.Sequential(OrderedDict([
            (name + "conv1", nn.Conv2d(in_channels, features, 3, padding=1, bias=False)),
            (name + "norm1", nn.BatchNorm2d(features)),
            (name + "relu1", nn.ReLU(inplace=True)),
            (name + "conv2", nn.Conv2d(features, features, 3, padding=1, bias=False)),
            (name + "norm2", nn.BatchNorm2d(features)),
            (name + "relu2", nn.ReLU(inplace=True)),
        ]))

---

Training and Test

Loss function
Binary Cross Entropy Loss (nn.BCELoss) is used for binary segmentation tasks.

Optimizer
Adam optimizer (optim.Adam) is used for its adaptive learning rate and efficient convergence.

Settings

• Batch size: 2
• Learning rate: 0.001
• Epochs: (set as needed, e.g., 20)
• Device: Uses GPU (CUDA) if available, otherwise CPU

import torch.optim as optim

model = UNet()
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.BCELoss()

for epoch in range(20):
    model.train()
    for images, masks in train_loader:
        images, masks = images.to(device), masks.to(device)
        outputs = model(images)
        loss = criterion(outputs, masks)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

---

Evaluation

• Dice Similarity Coefficient (DSC) is used to evaluate segmentation performance.
• Visual inspection of predicted masks is also common.

def dsc(y_pred, y_true):
    y_pred = (y_pred > 0.5).astype(int)
    y_true = y_true.astype(int)
    return np.sum(y_pred[y_true == 1]) * 2.0 / (np.sum(y_pred) + np.sum(y_true))

---

Additional Sections

Saving the Model:
The trained model is saved using:

torch.save(model.state_dict(), "unet_model.pth")

Visualization:
Visualize input images, ground truth masks, and predicted masks for qualitative assessment.

---

Conclusion

This project covers the essential steps for building a medical image segmentation model using U-Net:

• Data loading, preprocessing, and augmentation
• Model architecture design
• Training and evaluation
• Model saving and visualization

You can extend this project by experimenting with different architectures, datasets, and hyperparameters.

---

Resources

• U-Net: Convolutional Networks for Biomedical Image Segmentation (Original Paper, arXiv)
• U-Net Paper PDF
• MedPy - Medical Image Processing in Python
• PyTorch Documentation
• scikit-image Documentation
• MedMNIST Dataset Collection
• McMedHacks

Project repository

GitHub Code: Medical Image Classification, Segmentation using U-Net, and Tabular Data Classification