"""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()