Image Segmentation: Semantic Segmentation, Instance Segmentation & U-Net

import torch
import torch.nn as nn

class SoftDiceLoss(nn.Module):
    def __init__(self, eps: float = 1e-6):
        super().__init__()
        self.eps = eps

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        # logits: [N, 1, H, W] for binary segmentation
        # targets: [N, 1, H, W] float {0.0, 1.0}
        probs = torch.sigmoid(logits)
        
        probs_flat = probs.view(probs.shape[0], -1)
        targets_flat = targets.view(targets.shape[0], -1)

        intersection = (probs_flat * targets_flat).sum(dim=1)
        cardinality = probs_flat.sum(dim=1) + targets_flat.sum(dim=1)

        dice = (2.0 * intersection + self.eps) / (cardinality + self.eps)
        return 1.0 - dice.mean()

import torch
import torch.nn as nn
import torch.nn.functional as F

class DoubleConv(nn.Module):
    # (Conv2d -> BatchNorm -> ReLU) * 2 with padding=1 to preserve spatial dimensions.
    def __init__(self, in_channels: int, out_channels: int):
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.block(x)


class UNet(nn.Module):
    def __init__(self, in_channels: int = 3, num_classes: int = 3):
        super().__init__()
        # Encoder (Contracting Path)
        self.enc1 = DoubleConv(in_channels, 32)
        self.enc2 = DoubleConv(32, 64)
        self.enc3 = DoubleConv(64, 128)
        self.enc4 = DoubleConv(128, 256)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

        # Bottleneck
        self.bottleneck = DoubleConv(256, 512)

        # Decoder (Expanding Path) with Transposed Convolutions
        self.up4 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
        self.dec4 = DoubleConv(512, 256)

        self.up3 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
        self.dec3 = DoubleConv(256, 128)

        self.up2 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
        self.dec2 = DoubleConv(128, 64)

        self.up1 = nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2)
        self.dec1 = DoubleConv(64, 32)

        # 1x1 Convolution Segmentation Head
        self.out = nn.Conv2d(32, num_classes, kernel_size=1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Encoder forward pass
        e1 = self.enc1(x)                # [N, 32, H, W]
        e2 = self.enc2(self.pool(e1))    # [N, 64, H/2, W/2]
        e3 = self.enc3(self.pool(e2))    # [N, 128, H/4, W/4]
        e4 = self.enc4(self.pool(e3))    # [N, 256, H/8, W/8]

        # Bottleneck
        b = self.bottleneck(self.pool(e4)) # [N, 512, H/16, W/16]

        # Decoder stage 4 with skip connection from e4
        d4 = self.up4(b)
        if d4.shape[-2:] != e4.shape[-2:]:
            d4 = F.interpolate(d4, size=e4.shape[-2:], mode="bilinear", align_corners=False)
        d4 = self.dec4(torch.cat([d4, e4], dim=1)) # Concat channels: 256 + 256 = 512

        # Decoder stage 3 with skip connection from e3
        d3 = self.up3(d4)
        if d3.shape[-2:] != e3.shape[-2:]:
            d3 = F.interpolate(d3, size=e3.shape[-2:], mode="bilinear", align_corners=False)
        d3 = self.dec3(torch.cat([d3, e3], dim=1)) # 128 + 128 = 256

        # Decoder stage 2 with skip connection from e2
        d2 = self.up2(d3)
        if d2.shape[-2:] != e2.shape[-2:]:
            d2 = F.interpolate(d2, size=e2.shape[-2:], mode="bilinear", align_corners=False)
        d2 = self.dec2(torch.cat([d2, e2], dim=1)) # 64 + 64 = 128

        # Decoder stage 1 with skip connection from e1
        d1 = self.up1(d2)
        if d1.shape[-2:] != e1.shape[-2:]:
            d1 = F.interpolate(d1, size=e1.shape[-2:], mode="bilinear", align_corners=False)
        d1 = self.dec1(torch.cat([d1, e1], dim=1)) # 32 + 32 = 64

        # Output projection head produces raw logits: [N, num_classes, H, W]
        return self.out(d1)

# Sanity check output dimensions
model = UNet(in_channels=3, num_classes=4)
x = torch.randn(2, 3, 256, 256)
logits = model(x)
print(f"Logits shape: {logits.shape}") # Expect: [2, 4, 256, 256]

import numpy as np
from PIL import Image
from torch.utils.data import Dataset, DataLoader

class SegmentationDataset(Dataset):
    def __init__(self, image_paths, mask_paths, size=(256, 256)):
        self.image_paths = image_paths
        self.mask_paths = mask_paths
        self.size = size

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        # 1. Load image and mask
        image = Image.open(self.image_paths[idx]).convert("RGB")
        mask = Image.open(self.mask_paths[idx]).convert("L")

        # 2. Resize with correct interpolation!
        w, h = self.size[1], self.size[0]
        image = image.resize((w, h), Image.Resampling.BILINEAR)
        mask = mask.resize((w, h), Image.Resampling.NEAREST)

        # 3. Convert to PyTorch tensors
        img_np = np.asarray(image, dtype=np.float32) / 255.0  # Normalize to [0, 1]
        mask_np = np.asarray(mask, dtype=np.int64)            # Discrete class IDs

        image_tensor = torch.from_numpy(img_np).permute(2, 0, 1) # [C, H, W]
        mask_tensor = torch.from_numpy(mask_np)                 # [H, W]

        return image_tensor, mask_tensor

# Setup Training Loop
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = UNet(in_channels=3, num_classes=3).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)

# Training Epoch
model.train()
for images, masks in train_loader:
    images, masks = images.to(device), masks.to(device)
    
    optimizer.zero_grad(set_to_none=True)
    logits = model(images)               # [N, 3, 256, 256]
    loss = criterion(logits, masks)      # masks is [N, 256, 256] with dtype torch.long
    loss.backward()
    optimizer.step()

@torch.no_grad()
def evaluate_miou(model, dataloader, num_classes, device):
    model.eval()
    confusion = torch.zeros(num_classes, num_classes, dtype=torch.long, device=device)

    for images, targets in dataloader:
        images, targets = images.to(device), targets.to(device)
        logits = model(images)
        preds = logits.argmax(dim=1) # [N, H, W]

        # Valid pixel mask excluding ignore_index
        valid = (targets >= 0) & (targets < num_classes)
        indices = num_classes * targets[valid].long() + preds[valid].long()
        confusion += torch.bincount(indices, minlength=num_classes**2).view(num_classes, num_classes)

    # Compute per-class IoU from confusion matrix diagonal
    tp = confusion.diag()
    fp = confusion.sum(dim=0) - tp
    fn = confusion.sum(dim=1) - tp
    denominator = tp + fp + fn

    ious = tp.float() / denominator.clamp(min=1)
    present_classes = denominator > 0
    mean_iou = ious[present_classes].mean().item()

    return ious.cpu().numpy(), mean_iou