Skip to main content

✂️ Image Segmentation

Boxes aren't enough for self-driving cars. They need to know the exact pixel boundaries of the road, pedestrians, and streetlights!

🖍️ The Coloring Book​

Image Segmentation is essentially coloring inside the lines. Instead of the CNN outputting a tiny 1D array of guesses, it uses a U-Net architecture to expand the image back up to its original size, where every single pixel is given a class (e.g. Pixel [5, 10] = "Road").

🐍 Python Implementation​

import torch
import torchvision.models.segmentation as seg_models

# Load a massive DeepLabV3 segmentation model
model = seg_models.deeplabv3_resnet50(weights=seg_models.DeepLabV3_ResNet50_Weights.DEFAULT)
model.eval()

# Input Image
dummy_img = torch.randn(1, 3, 224, 224)

# The output is a massive mask of the same size!
output = model(dummy_img)['out']
print("Segmentation Mask Shape:", output.shape) # (1, 21 classes, 224, 224)