Segmentation
A label for every pixel.
In segmentation, the network labels each individual pixel. Road, car, person, tumour. It's what lets a car see exactly where the road ends, and a phone blur the background behind your head. But there's a puzzle. Convolutional networks shrink images, yet here we need the full size back.
- Label every pixel
- Background blur, medical scans
- Networks shrink, we need full size
The U-Net
In twenty fifteen, a medical imaging team solved this with the U Net. The left side shrinks the image, finding what's there. The right side grows it back, finding where. And the key part is the skip connections across the U, which copy fine details from the early layers straight to the later ones. Look at torch cat in the forward function. Run it. The output has the same height and width as the input, with one score per class, for every pixel. The diffusion model behind image generators uses this same shape.
import torch
from torch import nn
def conv(i, o):
return nn.Sequential(nn.Conv2d(i, o, 3, padding=1), nn.ReLU(),
nn.Conv2d(o, o, 3, padding=1), nn.ReLU())
class UNet(nn.Module):
def __init__(self, classes=2):
super().__init__()
self.down1, self.down2 = conv(1, 16), conv(16, 32)
self.pool = nn.MaxPool2d(2)
self.mid = conv(32, 64)
self.up2 = nn.ConvTranspose2d(64, 32, 2, stride=2)
self.dec2 = conv(64, 32)
self.up1 = nn.ConvTranspose2d(32, 16, 2, stride=2)
self.dec1 = conv(32, 16)
self.out = nn.Conv2d(16, classes, 1)
def forward(self, x):
a = self.down1(x)
b = self.down2(self.pool(a))
m = self.mid(self.pool(b))
y = self.dec2(torch.cat([self.up2(m), b], 1))
y = self.dec1(torch.cat([self.up1(y), a], 1))
return self.out(y)
print("In: ", [1, 1, 64, 64])
print("Out:", list(UNet()(torch.rand(1, 1, 64, 64)).shape))In: [1, 1, 64, 64] Out: [1, 2, 64, 64]
Training a segmenter
Let's train a small one. We generate noisy images, each with a hidden circle, and the label is the exact mask of that circle. The loss is simply cross entropy, on each pixel. Run it. The model finds the shapes, even under heavy noise. We use I o U to score it, since most pixels are background, and pixel accuracy on its own can flatter you. There's our old friend, the imbalance trap.
import torch
from torch import nn
def make(n):
x = torch.zeros(n, 1, 32, 32)
y = torch.zeros(n, 32, 32, dtype=torch.long)
for i in range(n):
cx, cy, r = torch.randint(8, 24, (2,)).tolist() + [torch.randint(3, 8, ()).item()]
yy, xx = torch.meshgrid(torch.arange(32), torch.arange(32), indexing="ij")
m = (xx - cx) ** 2 + (yy - cy) ** 2 < r * r
y[i] = m.long()
x[i, 0] = m.float() + 0.3 * torch.randn(32, 32)
return x, y
torch.manual_seed(0)
model = nn.Sequential(nn.Conv2d(1, 16, 3, padding=1), nn.ReLU(),
nn.Conv2d(16, 16, 3, padding=2, dilation=2), nn.ReLU(),
nn.Conv2d(16, 2, 1))
opt = torch.optim.Adam(model.parameters(), lr=2e-3)
for step in range(300):
x, y = make(16)
loss = nn.functional.cross_entropy(model(x), y)
opt.zero_grad(); loss.backward(); opt.step()
x, y = make(100)
p = model(x).argmax(1)
inter = ((p == 1) & (y == 1)).sum().item()
union = ((p == 1) | (y == 1)).sum().item()
print(f"Pixel accuracy {(p == y).float().mean():.3f}")
print(f"Overlap (IoU) with the true shapes {inter / union:.3f}")Pixel accuracy 0.998 Overlap (IoU) with the true shapes 0.980
Summary
The U Net shrinks to understand, and grows back to locate, with skips for detail. Keep that shape in mind, because next, we'll return to diffusion, and see how a U Net and a sentence together make an image.
Key takeaways- Every pixel gets a label
- U-Net: down, up, and skips
- Use IoU, not pixel accuracy