U-Net

"""A one-level U-Net forward pass and the Dice score - the structure, written out.

Feature maps are lists of channels, and a channel is a 2-D list [h][w]. Nothing
is vectorized, so every function maps onto one box of the encoder / decoder /
skip diagram. Concatenation across channels is literally list concatenation.
"""

import math


def conv3x3_same(maps, kernels, bias):
    """kernels[out][in] is a 3x3 grid; zero padding keeps height and width."""
    h, w = len(maps[0]), len(maps[0][0])
    output = []

    for out_ch in range(len(kernels)):
        channel = [[bias[out_ch]] * w for _ in range(h)]
        for in_ch in range(len(maps)):
            kernel = kernels[out_ch][in_ch]
            for y in range(h):
                for x in range(w):
                    for u in range(3):
                        for v in range(3):
                            yy, xx = y + u - 1, x + v - 1
                            if 0 <= yy < h and 0 <= xx < w:   # zero padding
                                channel[y][x] += maps[in_ch][yy][xx] * kernel[u][v]
        output.append(channel)

    return output


def relu(maps):
    return [[[max(0.0, v) for v in row] for row in channel] for channel in maps]


def max_pool2(maps):
    """Contracting path: halve resolution, keep the strongest response."""
    h, w = len(maps[0]), len(maps[0][0])
    return [
        [
            [
                max(channel[2 * i + u][2 * j + v] for u in range(2) for v in range(2))
                for j in range(w // 2)
            ]
            for i in range(h // 2)
        ]
        for channel in maps
    ]


def up_conv2x2(maps, kernels, bias):
    """Expanding path: transposed conv, stride 2. Each input pixel paints a 2x2 block."""
    h, w = len(maps[0]), len(maps[0][0])
    output = [[[bias[o]] * (2 * w) for _ in range(2 * h)] for o in range(len(kernels))]

    for out_ch in range(len(kernels)):
        for in_ch in range(len(maps)):
            for y in range(h):
                for x in range(w):
                    for u in range(2):
                        for v in range(2):
                            output[out_ch][2 * y + u][2 * x + v] += (
                                maps[in_ch][y][x] * kernels[out_ch][in_ch][u][v]
                            )
    return output


def concat(skip_maps, decoder_maps):
    """The U-Net skip: stack channels. ResNet would ADD them element-wise instead."""
    return skip_maps + decoder_maps


def conv1x1(maps, weights, bias):
    h, w = len(maps[0]), len(maps[0][0])
    return [
        [[bias[o] + sum(weights[o][c] * maps[c][y][x] for c in range(len(maps))) for x in range(w)] for y in range(h)]
        for o in range(len(weights))
    ]


def softmax_per_pixel(logits):
    h, w = len(logits[0]), len(logits[0][0])
    probs = [[[0.0] * w for _ in range(h)] for _ in logits]
    for y in range(h):
        for x in range(w):
            peak = max(channel[y][x] for channel in logits)
            exps = [math.exp(channel[y][x] - peak) for channel in logits]
            total = sum(exps)
            for c, value in enumerate(exps):
                probs[c][y][x] = value / total
    return probs


def unet_forward(image, params):
    """encoder -> downsample -> bottleneck -> upsample -> CONCAT skip -> decoder -> head."""
    enc = relu(conv3x3_same(image, *params["enc"]))     # full resolution
    down = max_pool2(enc)                                # context up, detail lost
    mid = relu(conv3x3_same(down, *params["mid"]))
    up = up_conv2x2(mid, *params["up"])                  # back to full resolution, but blurry
    merged = concat(enc, up)                             # the skip returns the lost detail
    dec = relu(conv3x3_same(merged, *params["dec"]))
    return softmax_per_pixel(conv1x1(dec, *params["head"]))


def cross_entropy(probs, labels):
    """Mean over every pixel of -log p(true class)."""
    h, w = len(labels), len(labels[0])
    total = 0.0
    for y in range(h):
        for x in range(w):
            total -= math.log(max(probs[labels[y][x]][y][x], 1e-12))
    return total / (h * w)


def dice_set(predicted, truth):
    """2|A n B| / (|A| + |B|), with masks as sets of (row, col) pixels."""
    if not predicted and not truth:
        return 1.0                          # both empty: a perfect match by convention
    return 2 * len(predicted & truth) / (len(predicted) + len(truth))


def jaccard_set(predicted, truth):
    """|A n B| / |A u B| - the IoU used as the evaluation metric."""
    union = predicted | truth
    return len(predicted & truth) / len(union) if union else 1.0


def soft_dice_loss(probs, labels, eps=1e-6):
    """1 - mean over classes of 2*sum(p*y) / (sum(p) + sum(y)), on soft probabilities."""
    h, w = len(labels), len(labels[0])
    scores = []
    for c in range(len(probs)):
        intersection = prob_sum = truth_sum = 0.0
        for y in range(h):
            for x in range(w):
                is_class = 1.0 if labels[y][x] == c else 0.0
                intersection += probs[c][y][x] * is_class
                prob_sum += probs[c][y][x]
                truth_sum += is_class
        scores.append((2 * intersection + eps) / (prob_sum + truth_sum + eps))
    return 1.0 - sum(scores) / len(scores)


def imbalance_demo():
    """A 2%-foreground mask: predicting nothing scores 98% accuracy and 0 overlap."""
    truth = {(r, c) for r in range(10) for c in range(20)}     # 200 of 10,000 pixels
    predicted = set()                                          # "background everywhere"
    accuracy = 1 - len(predicted ^ truth) / 10_000
    return accuracy, dice_set(predicted, truth), jaccard_set(predicted, truth)

"""U-Net - typed modules, concat skips with explicit shape contracts, Dice + CE."""

from dataclasses import dataclass

import torch
import torch.nn.functional as F
from torch import Tensor, nn


@dataclass(frozen=True)
class UNetConfig:
    in_channels: int = 1
    num_classes: int = 2
    base_width: int = 32
    depth: int = 4

    @property
    def size_multiple(self) -> int:
        """Each downsampling halves H and W, so inputs must be a multiple of 2**depth."""
        return 2**self.depth


class DoubleConv(nn.Module):
    """(Conv3x3 -> BatchNorm -> ReLU) twice. Bias is off: BatchNorm re-centres anyway."""

    def __init__(self, in_ch: int, out_ch: int) -> None:
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
        )

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


class Down(nn.Module):
    def __init__(self, in_ch: int, out_ch: int) -> None:
        super().__init__()
        self.pool = nn.MaxPool2d(2)
        self.conv = DoubleConv(in_ch, out_ch)

    def forward(self, x: Tensor) -> Tensor:
        return self.conv(self.pool(x))


class Up(nn.Module):
    """Upsample, CONCATENATE the encoder skip, then convolve.

    The concat is why the decoder conv takes (in_ch // 2) + skip_ch channels:
    the two sources are stacked and the conv learns how to weigh them. A
    ResNet-style add would need equal channel counts and mix them blindly.
    """

    def __init__(self, in_ch: int, skip_ch: int, out_ch: int) -> None:
        super().__init__()
        self.up = nn.ConvTranspose2d(in_ch, in_ch // 2, kernel_size=2, stride=2)
        self.conv = DoubleConv(in_ch // 2 + skip_ch, out_ch)

    def forward(self, x: Tensor, skip: Tensor) -> Tensor:
        x = self.up(x)
        if x.shape[-2:] != skip.shape[-2:]:
            raise ValueError(
                f"skip {tuple(skip.shape[-2:])} and upsampled {tuple(x.shape[-2:])} "
                "feature maps disagree; pad the input to a multiple of 2**depth"
            )
        return self.conv(torch.cat([skip, x], dim=1))


class UNet(nn.Module):
    def __init__(self, config: UNetConfig) -> None:
        super().__init__()
        if config.depth < 1:
            raise ValueError(f"depth must be >= 1, got {config.depth}")
        if config.base_width < 1:
            raise ValueError(f"base_width must be >= 1, got {config.base_width}")

        self.config = config
        widths = [config.base_width * 2**level for level in range(config.depth + 1)]

        self.stem = DoubleConv(config.in_channels, widths[0])
        self.downs = nn.ModuleList(
            Down(widths[level], widths[level + 1]) for level in range(config.depth)
        )
        self.ups = nn.ModuleList(
            Up(widths[level + 1], widths[level], widths[level])
            for level in reversed(range(config.depth))
        )
        self.head = nn.Conv2d(widths[0], config.num_classes, kernel_size=1)

    def forward(self, x: Tensor) -> Tensor:
        if x.dim() != 4:
            raise ValueError(f"expected (batch, channels, height, width), got {tuple(x.shape)}")

        multiple = self.config.size_multiple
        height, width = x.shape[-2:]
        if height % multiple or width % multiple:
            raise ValueError(
                f"height and width must be multiples of {multiple}, got {height}x{width}; "
                "use pad_to_multiple() and crop the output"
            )

        skips: list[Tensor] = []
        x = self.stem(x)
        for down in self.downs:
            skips.append(x)       # held until the matching decoder level needs it
            x = down(x)
        for up in self.ups:
            x = up(x, skips.pop())
        return self.head(x)       # logits: (batch, num_classes, H, W)


def pad_to_multiple(x: Tensor, multiple: int) -> tuple[Tensor, tuple[int, int]]:
    """Reflect-pad bottom/right so H and W divide evenly; return the padding to undo it.

    Padding then cropping the OUTPUT is preferred to cropping the input, which
    would throw away real pixels at the border.
    """
    height, width = x.shape[-2:]
    pad_h = (-height) % multiple
    pad_w = (-width) % multiple
    if pad_h >= height or pad_w >= width:
        raise ValueError("input is too small to reflect-pad; use mode='constant' instead")
    return F.pad(x, (0, pad_w, 0, pad_h), mode="reflect"), (pad_h, pad_w)


def crop_padding(x: Tensor, padding: tuple[int, int]) -> Tensor:
    pad_h, pad_w = padding
    return x[..., : x.shape[-2] - pad_h, : x.shape[-1] - pad_w]


class DiceCrossEntropy(nn.Module):
    """Cross-entropy + soft Dice. CE is smooth early; Dice fixes class imbalance."""

    def __init__(self, dice_weight: float = 1.0, eps: float = 1e-6) -> None:
        super().__init__()
        if dice_weight < 0.0:
            raise ValueError(f"dice_weight must be >= 0, got {dice_weight}")
        self.dice_weight = dice_weight
        self.eps = eps

    def forward(self, logits: Tensor, target: Tensor) -> Tensor:
        """logits (N, C, H, W); target (N, H, W) of integer class ids."""
        if target.dim() != 3:
            raise ValueError(f"target must be (N, H, W), got {tuple(target.shape)}")

        cross_entropy = F.cross_entropy(logits, target)

        probs = logits.float().softmax(dim=1)
        one_hot = F.one_hot(target, logits.shape[1]).permute(0, 3, 1, 2).to(probs.dtype)

        # Pooled over batch AND space, per class. Per-image Dice is 0/0 on an
        # image with no foreground and its gradient blows up as mass -> 0.
        dims = (0, 2, 3)
        intersection = (probs * one_hot).sum(dims)
        cardinality = probs.sum(dims) + one_hot.sum(dims)
        dice = (2.0 * intersection + self.eps) / (cardinality + self.eps)

        return cross_entropy + self.dice_weight * (1.0 - dice.mean())


@dataclass(frozen=True)
class SegmentationScores:
    dice: float
    iou: float


@torch.no_grad()
def class_scores(pred: Tensor, target: Tensor, class_id: int) -> SegmentationScores | None:
    """Hard-mask Dice and IoU for one class; None when the class is absent from both."""
    predicted = pred == class_id
    truth = target == class_id
    intersection = int((predicted & truth).sum())
    total = int(predicted.sum()) + int(truth.sum())
    if total == 0:
        return None     # undefined, not perfect: skip rather than inflate the mean
    union = total - intersection
    return SegmentationScores(dice=2 * intersection / total, iou=intersection / union)

"""U-Net - channels-last + bf16 + compile for training; batched, blended tiling for inference."""

import torch
from torch import Tensor, nn


def prepare_for_training(model: nn.Module, device: str) -> nn.Module:
    """channels_last layout and torch.compile.

    Tensor cores want NHWC; PyTorch defaults to NCHW, so cuDNN inserts a
    transpose around every convolution. Converting once removes it everywhere.
    """
    model = model.to(device, memory_format=torch.channels_last)
    return torch.compile(model)   # fuses BatchNorm + ReLU into the surrounding kernels


def train_step(model: nn.Module, optimizer, loss_fn, images: Tensor, masks: Tensor, device: str) -> float:
    images = images.to(device, non_blocking=True, memory_format=torch.channels_last)
    masks = masks.to(device, non_blocking=True)

    optimizer.zero_grad(set_to_none=True)
    with torch.autocast(device_type=device, dtype=torch.bfloat16):
        logits = model(images)
    # Dice sums many small probabilities: do that reduction in float32.
    loss = loss_fn(logits.float(), masks)

    loss.backward()
    optimizer.step()
    return float(loss.item())


def gaussian_window(tile: int, device: str, sigma_scale: float = 0.125) -> Tensor:
    """Weight tile centres more than edges, where the receptive field lacks context."""
    coords = torch.arange(tile, dtype=torch.float32, device=device) - (tile - 1) / 2
    profile = torch.exp(-(coords**2) / (2 * (tile * sigma_scale) ** 2))
    window = profile[:, None] * profile[None, :]
    return window / window.max()


def tile_origins(length: int, tile: int, stride: int) -> list[int]:
    """Window starts that cover the axis, with a final tile flush against the edge."""
    if length < tile:
        raise ValueError(f"image side {length} is smaller than the tile {tile}; pad first")
    last = length - tile
    origins = list(range(0, last + 1, stride))
    if origins[-1] != last:
        origins.append(last)
    return origins


@torch.no_grad()
def sliding_window_predict(
    model: nn.Module,
    image: Tensor,           # (C, H, W), already padded to be at least tile x tile
    num_classes: int,
    tile: int = 256,
    overlap: float = 0.5,
    batch_size: int = 8,
) -> Tensor:
    model.eval()
    device = image.device.type
    _, height, width = image.shape
    stride = max(1, int(tile * (1.0 - overlap)))
    window = gaussian_window(tile, device)

    # Pre-allocated accumulators; every tile is added in place, no per-tile result list.
    prob_sum = torch.zeros(num_classes, height, width, device=image.device)
    weight_sum = torch.zeros(1, height, width, device=image.device)

    origins = [
        (top, left)
        for top in tile_origins(height, tile, stride)
        for left in tile_origins(width, tile, stride)
    ]

    for start in range(0, len(origins), batch_size):
        batch_origins = origins[start : start + batch_size]
        # One forward over a stack of tiles instead of one call per tile.
        batch = torch.stack(
            [image[:, top : top + tile, left : left + tile] for top, left in batch_origins]
        ).contiguous(memory_format=torch.channels_last)

        with torch.autocast(device_type=device, dtype=torch.bfloat16):
            logits = model(batch)
        probs = logits.float().softmax(dim=1)

        for index, (top, left) in enumerate(batch_origins):
            prob_sum[:, top : top + tile, left : left + tile].add_(probs[index] * window)
            weight_sum[:, top : top + tile, left : left + tile].add_(window)

    return prob_sum / weight_sum