scripts/train_deeplab_semseg.py

"""Train a semantic segmentation baseline from the YOLO polygon dataset.

Unlike YOLO instance segmentation, this model directly optimizes per-pixel
semantic classes, which matches the project mIoU target more closely.
"""

from __future__ import annotations

import argparse
import csv
import os
import sys
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from torch import nn
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from torchvision.models.segmentation import (
    DeepLabV3_ResNet101_Weights,
    DeepLabV3_ResNet50_Weights,
    deeplabv3_resnet101,
    deeplabv3_resnet50,
)


PROJECT_ROOT = Path(__file__).resolve().parents[1]
SCRIPT_DIR = Path(__file__).resolve().parent
if str(SCRIPT_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPT_DIR))
os.environ.setdefault("TORCH_HOME", str(PROJECT_ROOT / ".torch"))

from segmentation_utils import load_dataset_class_names, yolo_label_to_semantic_mask

DEFAULT_DATASET = PROJECT_ROOT / "data" / "processed" / "foodseg103_target_yolo"
DEFAULT_OUTPUT = PROJECT_ROOT / "runs" / "foodseg103_target" / "deeplabv3_resnet50_target"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dataset", type=Path, default=DEFAULT_DATASET)
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument("--epochs", type=int, default=20)
    parser.add_argument("--imgsz", type=int, default=512)
    parser.add_argument("--batch", type=int, default=8)
    parser.add_argument("--workers", type=int, default=0)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--weight-decay", type=float, default=1e-4)
    parser.add_argument("--limit-train", type=int)
    parser.add_argument("--limit-val", type=int)
    parser.add_argument("--no-pretrained", action="store_true")
    parser.add_argument("--no-class-weights", action="store_true")
    parser.add_argument("--class-weights", nargs="+", type=float)
    parser.add_argument("--no-amp", action="store_true")
    parser.add_argument("--resume-checkpoint", type=Path)
    parser.add_argument("--strong-aug", action="store_true")
    parser.add_argument("--backbone", choices=["resnet50", "resnet101"], default="resnet50")
    return parser.parse_args()


class YoloSemanticDataset(Dataset):
    def __init__(
        self,
        root: Path,
        split: str,
        imgsz: int,
        limit: int | None = None,
        train: bool = False,
        strong_aug: bool = False,
    ) -> None:
        self.root = root
        self.split = split
        self.imgsz = imgsz
        self.train = train
        self.strong_aug = strong_aug
        self.image_paths = sorted((root / "images" / split).glob("*.jpg"))
        if limit is not None:
            self.image_paths = self.image_paths[:limit]
        self.normalize = transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225],
        )

    def __len__(self) -> int:
        return len(self.image_paths)

    def __getitem__(self, index: int) -> tuple[torch.Tensor, torch.Tensor]:
        image_path = self.image_paths[index]
        label_path = self.root / "labels" / self.split / f"{image_path.stem}.txt"

        image = Image.open(image_path).convert("RGB")
        width, height = image.size
        mask_np = yolo_label_to_semantic_mask(
            label_path,
            height,
            width,
            background_value=0,
            class_offset=1,
        ).astype(np.uint8)
        mask = Image.fromarray(mask_np, mode="L")

        if self.train and torch.rand(()) < 0.5:
            image = image.transpose(Image.Transpose.FLIP_LEFT_RIGHT)
            mask = mask.transpose(Image.Transpose.FLIP_LEFT_RIGHT)

        if self.train and self.strong_aug:
            image, mask = self.apply_strong_aug(image, mask)

        image = image.resize((self.imgsz, self.imgsz), Image.Resampling.BILINEAR)
        mask = mask.resize((self.imgsz, self.imgsz), Image.Resampling.NEAREST)

        image_arr = np.asarray(image, dtype=np.float32) / 255.0
        image_tensor = torch.from_numpy(image_arr).permute(2, 0, 1)
        image_tensor = self.normalize(image_tensor)
        mask_tensor = torch.from_numpy(np.asarray(mask, dtype=np.int64))
        return image_tensor, mask_tensor

    def apply_strong_aug(self, image: Image.Image, mask: Image.Image) -> tuple[Image.Image, Image.Image]:
        if torch.rand(()) < 0.8:
            brightness = float(torch.empty(1).uniform_(0.75, 1.25))
            contrast = float(torch.empty(1).uniform_(0.75, 1.25))
            saturation = float(torch.empty(1).uniform_(0.75, 1.25))
            image = transforms.functional.adjust_brightness(image, brightness)
            image = transforms.functional.adjust_contrast(image, contrast)
            image = transforms.functional.adjust_saturation(image, saturation)

        if torch.rand(()) < 0.25:
            image = transforms.functional.gaussian_blur(image, kernel_size=3)

        if torch.rand(()) < 0.7:
            image, mask = self.random_resized_crop_pair(image, mask)

        return image, mask

    def random_resized_crop_pair(self, image: Image.Image, mask: Image.Image) -> tuple[Image.Image, Image.Image]:
        width, height = image.size
        scale = float(torch.empty(1).uniform_(0.75, 1.0))
        crop_w = max(1, int(width * scale))
        crop_h = max(1, int(height * scale))
        if crop_w == width and crop_h == height:
            return image, mask
        left = int(torch.randint(0, width - crop_w + 1, (1,)).item())
        top = int(torch.randint(0, height - crop_h + 1, (1,)).item())
        box = (left, top, left + crop_w, top + crop_h)
        return image.crop(box), mask.crop(box)


def build_model(num_classes: int, pretrained: bool, backbone: str = "resnet50") -> nn.Module:
    if backbone == "resnet101":
        weights = DeepLabV3_ResNet101_Weights.DEFAULT if pretrained else None
        model = deeplabv3_resnet101(weights=weights, weights_backbone=None, aux_loss=True)
    else:
        weights = DeepLabV3_ResNet50_Weights.DEFAULT if pretrained else None
        model = deeplabv3_resnet50(weights=weights, weights_backbone=None, aux_loss=True)
    model.classifier[-1] = nn.Conv2d(256, num_classes, kernel_size=1)
    if model.aux_classifier is not None:
        model.aux_classifier[-1] = nn.Conv2d(256, num_classes, kernel_size=1)
    return model


def load_compatible_state_dict(model: nn.Module, checkpoint: Path, device: torch.device) -> None:
    state_dict = torch.load(checkpoint, map_location=device)
    model_state = model.state_dict()
    compatible = {
        key: value
        for key, value in state_dict.items()
        if key in model_state and value.shape == model_state[key].shape
    }
    skipped = sorted(set(state_dict) - set(compatible))
    model.load_state_dict(compatible, strict=False)
    if skipped:
        print(f"Skipped incompatible checkpoint keys: {len(skipped)}")


def compute_class_weights(dataset: Dataset, num_classes: int) -> torch.Tensor:
    counts = torch.zeros(num_classes, dtype=torch.float64)
    for _, mask in dataset:
        counts += torch.bincount(mask.reshape(-1), minlength=num_classes).double()

    counts = torch.clamp(counts, min=1.0)
    weights = torch.sqrt(counts.sum() / counts)
    weights = weights / weights.mean()
    weights[0] = min(weights[0].item(), 0.5)
    return weights.float()


def compute_miou(confusion: torch.Tensor, class_names: dict[int, str]) -> tuple[dict[str, float], float]:
    ious = {}
    for class_id, name in class_names.items():
        target_id = class_id + 1
        tp = confusion[target_id, target_id].item()
        fp = confusion[:, target_id].sum().item() - tp
        fn = confusion[target_id, :].sum().item() - tp
        denom = tp + fp + fn
        ious[name] = tp / denom if denom else float("nan")
    return ious, float(np.nanmean(list(ious.values())))


@torch.no_grad()
def evaluate(model: nn.Module, loader: DataLoader, device: torch.device, num_classes_with_bg: int, class_names: dict[int, str]) -> tuple[float, dict[str, float]]:
    model.eval()
    confusion = torch.zeros((num_classes_with_bg, num_classes_with_bg), dtype=torch.int64, device=device)

    for images, masks in loader:
        images = images.to(device)
        masks = masks.to(device)
        logits = model(images)["out"]
        if logits.shape[-2:] != masks.shape[-2:]:
            logits = F.interpolate(logits, size=masks.shape[-2:], mode="bilinear", align_corners=False)
        preds = logits.argmax(dim=1)
        valid = (masks >= 0) & (masks < num_classes_with_bg)
        indices = masks[valid] * num_classes_with_bg + preds[valid]
        confusion += torch.bincount(indices, minlength=num_classes_with_bg**2).reshape(num_classes_with_bg, num_classes_with_bg)

    ious, miou = compute_miou(confusion.cpu(), class_names)
    return miou, ious


def main() -> None:
    args = parse_args()
    class_names = load_dataset_class_names(args.dataset)
    num_classes_with_bg = len(class_names) + 1
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    args.output.mkdir(parents=True, exist_ok=True)

    train_ds = YoloSemanticDataset(args.dataset, "train", args.imgsz, args.limit_train, train=True, strong_aug=args.strong_aug)
    val_ds = YoloSemanticDataset(args.dataset, "val", args.imgsz, args.limit_val, train=False)
    train_loader = DataLoader(train_ds, batch_size=args.batch, shuffle=True, num_workers=args.workers, pin_memory=torch.cuda.is_available())
    val_loader = DataLoader(val_ds, batch_size=args.batch, shuffle=False, num_workers=args.workers, pin_memory=torch.cuda.is_available())

    model = build_model(num_classes_with_bg, pretrained=not args.no_pretrained, backbone=args.backbone).to(device)
    if args.resume_checkpoint:
        load_compatible_state_dict(model, args.resume_checkpoint, device)
        print(f"Loaded checkpoint: {args.resume_checkpoint}")
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
    if args.class_weights:
        if len(args.class_weights) != num_classes_with_bg:
            raise ValueError(f"Expected {num_classes_with_bg} class weights, got {len(args.class_weights)}")
        class_weights = torch.tensor(args.class_weights, dtype=torch.float32)
    elif args.no_class_weights:
        class_weights = None
    else:
        class_weights = compute_class_weights(train_ds, num_classes_with_bg)
    if class_weights is not None:
        print(f"class_weights={class_weights.tolist()}")
        class_weights = class_weights.to(device)
    criterion = nn.CrossEntropyLoss(weight=class_weights)
    use_amp = torch.cuda.is_available() and not args.no_amp
    scaler = torch.amp.GradScaler("cuda", enabled=use_amp)

    log_path = args.output / "metrics.csv"
    best_miou = -1.0
    with log_path.open("w", newline="", encoding="utf-8") as f:
        fieldnames = ["epoch", "train_loss", "val_miou", *[f"iou_{name}" for name in class_names.values()]]
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()

        for epoch in range(1, args.epochs + 1):
            model.train()
            running_loss = 0.0
            seen = 0
            for images, masks in train_loader:
                images = images.to(device)
                masks = masks.to(device)
                optimizer.zero_grad(set_to_none=True)
                with torch.amp.autocast("cuda", enabled=use_amp):
                    logits = model(images)["out"]
                    if logits.shape[-2:] != masks.shape[-2:]:
                        logits = F.interpolate(logits, size=masks.shape[-2:], mode="bilinear", align_corners=False)
                    loss = criterion(logits, masks)
                scaler.scale(loss).backward()
                scaler.step(optimizer)
                scaler.update()
                running_loss += loss.item() * images.size(0)
                seen += images.size(0)

            train_loss = running_loss / max(seen, 1)
            val_miou, ious = evaluate(model, val_loader, device, num_classes_with_bg, class_names)
            row = {"epoch": epoch, "train_loss": round(train_loss, 6), "val_miou": round(val_miou, 6)}
            row.update({f"iou_{name}": round(value, 6) for name, value in ious.items()})
            writer.writerow(row)
            f.flush()
            print(row)

            torch.save(model.state_dict(), args.output / "last.pt")
            if val_miou > best_miou:
                best_miou = val_miou
                torch.save(model.state_dict(), args.output / "best.pt")

    print(f"Best val mIoU: {best_miou:.4f}")
    print(f"Wrote metrics: {log_path}")


if __name__ == "__main__":
    main()