Image segmentation with U-Net

import numpy as np
import torch
from torch import nn

def make_scans(n, seed, size=32):
    """Synthetic scans with one or two irregular 'lesions', and the true pixel mask for each."""
    rng = np.random.default_rng(seed)
    yy, xx = np.mgrid[:size, :size]
    images, masks = rng.normal(0, 1, (n, size, size)), np.zeros((n, size, size))
    for i in range(n):
        for _ in range(rng.integers(1, 3)):
            cy, cx = rng.integers(6, size - 6, 2)
            ry, rx = rng.uniform(2, 5, 2)
            masks[i] = np.maximum(masks[i], ((yy - cy) / ry) ** 2 + ((xx - cx) / rx) ** 2 < 1)
        images[i] += 1.5 * masks[i]
    to_t = lambda a: torch.tensor(a, dtype=torch.float32).unsqueeze(1)
    return to_t(images), to_t(masks)

def block(c_in, c_out):
    return nn.Sequential(nn.Conv2d(c_in, c_out, 3, padding=1), nn.ReLU(),
                         nn.Conv2d(c_out, c_out, 3, padding=1), nn.ReLU())

class TinyUNet(nn.Module):
    """Encoder-decoder with a skip connection: the U-Net idea at the smallest useful size."""
    def __init__(self):
        super().__init__()
        self.enc1, self.enc2 = block(1, 8), block(8, 16)
        self.pool, self.up = nn.MaxPool2d(2), nn.ConvTranspose2d(16, 8, 2, stride=2)
        self.dec1 = block(16, 8)                       # 8 upsampled + 8 from the skip connection
        self.out = nn.Conv2d(8, 1, 1)
    def forward(self, x):
        e1 = self.enc1(x)                              # full resolution: detail
        e2 = self.enc2(self.pool(e1))                  # half resolution: context
        d1 = self.dec1(torch.cat([self.up(e2), e1], dim=1))
        return self.out(d1)                            # one logit per pixel

def dice(pred_mask, true_mask, eps=1e-6):
    inter = (pred_mask * true_mask).sum(dim=(1, 2, 3))
    return ((2 * inter + eps) / (pred_mask.sum(dim=(1, 2, 3)) + true_mask.sum(dim=(1, 2, 3)) + eps))

X_train, M_train = make_scans(1500, seed=1)
X_test, M_test = make_scans(300, seed=2)
torch.manual_seed(0)
unet = TinyUNet()
opt = torch.optim.Adam(unet.parameters(), lr=2e-3)
for epoch in range(1, 9):
    for idx in torch.randperm(len(X_train)).split(32):
        opt.zero_grad()
        nn.functional.binary_cross_entropy_with_logits(unet(X_train[idx]), M_train[idx]).backward()
        opt.step()
    if epoch in (1, 4, 8):
        with torch.no_grad():
            d = dice((unet(X_test) > 0).float(), M_test)
        print(f"epoch {epoch}: mean Dice on test scans {d.mean():.3f}")

# epoch 1: mean Dice on test scans 0.000
# epoch 4: mean Dice on test scans 0.864
# epoch 8: mean Dice on test scans 0.891

with torch.no_grad():
    d = dice((unet(X_test) > 0).float(), M_test)
    threshold_only = dice((X_test > 0.75).float(), M_test)       # naive: bright pixels = lesion
print(f"Simple brightness threshold: mean Dice {threshold_only.mean():.3f}")
print(f"U-Net: mean Dice {d.mean():.3f}; worst 5% of scans below {torch.quantile(d, 0.05):.3f}")

# Simple brightness threshold: mean Dice 0.270
# U-Net: mean Dice 0.891; worst 5% of scans below 0.769