Medical Image Segmentation

U-Net Architecture

Encoder-Decoder with Skip Connections

import torch.nn as nn

class UNet(nn.Module):
    def __init__(self, in_channels=1, n_classes=2):
        super().__init__()
        self.enc1 = nn.Sequential(nn.Conv2d(in_channels, 64, 3, padding=1),
                                  nn.BatchNorm2d(64), nn.ReLU(),
                                  nn.Conv2d(64, 64, 3, padding=1),
                                  nn.BatchNorm2d(64), nn.ReLU())
        self.pool = nn.MaxPool2d(2)
        self.enc2 = nn.Sequential(nn.Conv2d(64, 128, 3, padding=1),
                                  nn.BatchNorm2d(128), nn.ReLU())
        self.up = nn.ConvTranspose2d(128, 64, 2, stride=2)
        self.dec = nn.Sequential(nn.Conv2d(128, 64, 3, padding=1),
                                 nn.BatchNorm2d(64), nn.ReLU())
        self.output = nn.Conv2d(64, n_classes, 1)

    def forward(self, x):
        e1 = self.enc1(x)
        e2 = self.enc2(self.pool(e1))
        d = self.dec(torch.cat([self.up(e2), e1], dim=1))
        return self.output(d)

3D Segmentation

V-Net with Dice Loss

def dice_loss(pred, target, smooth=1.0):
    pred_flat = pred.view(-1)
    target_flat = target.view(-1)
    intersection = (pred_flat * target_flat).sum()
    return 1 - (2. * intersection + smooth) / (pred_flat.sum() + target_flat.sum() + smooth)

Active Learning for Annotation

Evaluation