train_segmentaion.py

import warnings
warnings.filterwarnings('ignore')
import logging
logging.getLogger('albumentations').setLevel(logging.ERROR)
logging.getLogger('torch').setLevel(logging.ERROR)

import os
import time
import random
import json
import numpy as np
from glob import glob
from collections import OrderedDict

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torch.optim.swa_utils import AveragedModel

import cv2
from PIL import Image
import albumentations as A
from albumentations.pytorch import ToTensorV2
import segmentation_models_pytorch as smp

# ═══════════════════════════════════════════════════════════════
# CONFIGURATION (Local Paths for GitHub Submission)
# ═══════════════════════════════════════════════════════════════
TRAIN_IMG_DIR     = './dataset/Offroad_Segmentation_Training_Dataset/train/Color_Images/'
TRAIN_MASK_DIR    = './dataset/Offroad_Segmentation_Training_Dataset/train/Segmentation/'
VAL_IMG_DIR       = './dataset/Offroad_Segmentation_Training_Dataset/val/Color_Images/'
VAL_MASK_DIR      = './dataset/Offroad_Segmentation_Training_Dataset/val/Segmentation/'

WEIGHTS_JSON_PATH = './rare_class_weights.json'
MODEL_SAVE_PATH   = './models/segmentation_head.pth'
SWA_SAVE_PATH     = './models/swa_model_v4.pth'

N_CLASSES   = 10
BATCH_SIZE   = 4
NUM_WORKERS  = 2
N_EPOCHS     = 25
PATIENCE     = 8
CUTMIX_PROB  = 0.3
ACCUMULATION_STEPS = 4
SEED         = 42

VALUE_MAP = {
    0: 0, 100: 1, 200: 2, 300: 3, 500: 4,
    550: 5, 600: 6, 700: 7, 800: 8, 7100: 0, 10000: 9
}

CLASS_WEIGHTS = torch.tensor(
    [1.0, 5.0, 8.0, 10.0, 12.0, 15.0, 12.0, 15.0, 20.0, 3.0],
    dtype=torch.float32
)

random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed_all(SEED)
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.deterministic = False

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# ═══════════════════════════════════════════════════════════════
# DATASET & TRANSFORMS (Dynamic Resizing supported)
# ═══════════════════════════════════════════════════════════════
def get_train_transforms(h, w):
    return A.Compose([
        A.Resize(h, w),
        A.HorizontalFlip(p=0.5),
        A.Affine(scale=(0.7, 1.3), translate_percent=(-0.1, 0.1),
                 rotate=(-15, 15), p=0.5),
        A.GridDistortion(p=0.2),
        A.CLAHE(clip_limit=4.0, tile_grid_size=(8, 8), p=0.5),
        A.RandomGamma(gamma_limit=(70, 130), p=0.3),
        A.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4,
                      hue=0.1, p=0.5),
        A.RandomBrightnessContrast(p=0.2),
        A.RandomShadow(p=0.2),
        A.CoarseDropout(num_holes_range=(1, 8),
                        hole_height_range=(8, 32),
                        hole_width_range=(8, 32), p=0.3),
        A.Normalize(mean=[0.485, 0.456, 0.406],
                    std=[0.229, 0.224, 0.225]),
        ToTensorV2(),
    ])

def get_val_transforms(h, w):
    return A.Compose([
        A.Resize(h, w),
        A.Normalize(mean=[0.485, 0.456, 0.406],
                    std=[0.229, 0.224, 0.225]),
        ToTensorV2(),
    ])

class OffroadDataset(Dataset):
    def __init__(self, img_dir, mask_dir, transforms=None, is_train=False):
        self.img_paths = sorted(glob(os.path.join(img_dir, '*.png')) +
                                glob(os.path.join(img_dir, '*.jpg')))
        self.mask_paths = sorted(glob(os.path.join(mask_dir, '*.png')) +
                                 glob(os.path.join(mask_dir, '*.jpg')))
        assert len(self.img_paths) == len(self.mask_paths)
        self.transforms = transforms
        self.is_train = is_train

        self.sample_weights = []
        if self.is_train:
            if os.path.exists(WEIGHTS_JSON_PATH):
                with open(WEIGHTS_JSON_PATH, 'r') as f:
                    self.sample_weights = json.load(f)
                print(f"✅ Loaded {len(self.sample_weights)} pre-computed sample weights instantly.")
            else:
                raise FileNotFoundError(f"Missing weights file! Expected at: {WEIGHTS_JSON_PATH}")

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

    def _remap_mask(self, mask):
        out = np.zeros_like(mask, dtype=np.int64)
        for raw_id, mapped_id in VALUE_MAP.items():
            out[mask == raw_id] = mapped_id
        return out

    def __getitem__(self, idx):
        img = cv2.imread(self.img_paths[idx], cv2.IMREAD_COLOR)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        mask = cv2.imread(self.mask_paths[idx], cv2.IMREAD_UNCHANGED)
        
        if len(mask.shape) == 3:
            mask = mask[:, :, 0]
        mask = self._remap_mask(mask)

        if self.transforms:
            augmented = self.transforms(image=img, mask=mask)
            img = augmented['image']
            mask = augmented['mask'].long()

        if self.is_train and random.random() < CUTMIX_PROB:
            mix_idx = random.randint(0, len(self) - 1)
            mix_img = cv2.imread(self.img_paths[mix_idx], cv2.IMREAD_COLOR)
            mix_img = cv2.cvtColor(mix_img, cv2.COLOR_BGR2RGB)
            mix_mask = cv2.imread(self.mask_paths[mix_idx], cv2.IMREAD_UNCHANGED)
            
            if len(mix_mask.shape) == 3:
                mix_mask = mix_mask[:, :, 0]
            mix_mask = self._remap_mask(mix_mask)

            if self.transforms:
                mix_aug = self.transforms(image=mix_img, mask=mix_mask)
                mix_img = mix_aug['image']
                mix_mask = mix_aug['mask'].long()

            _, h, w = img.shape
            cut_h = random.randint(h // 4, h // 2)
            cut_w = random.randint(w // 4, w // 2)
            cy = random.randint(0, h - cut_h)
            cx = random.randint(0, w - cut_w)
            
            img[:, cy:cy+cut_h, cx:cx+cut_w] = mix_img[:, cy:cy+cut_h, cx:cx+cut_w]
            mask[cy:cy+cut_h, cx:cx+cut_w] = mix_mask[cy:cy+cut_h, cx:cx+cut_w]

        return img, mask

# ═══════════════════════════════════════════════════════════════
# ARCHITECTURE: DINOv2 + UPerNet
# ═══════════════════════════════════════════════════════════════
class PPM(nn.Module):
    def __init__(self, in_ch, out_ch, bins=(1, 2, 3, 6)):
        super().__init__()
        self.features = nn.ModuleList()
        for bin in bins:
            self.features.append(nn.Sequential(
                nn.AdaptiveAvgPool2d(bin),
                nn.Conv2d(in_ch, out_ch, 1, bias=False),
                nn.BatchNorm2d(out_ch),
                nn.ReLU(inplace=True)
            ))
        self.bottleneck = nn.Sequential(
            nn.Conv2d(in_ch + len(bins) * out_ch, out_ch, 3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        h, w = x.shape[2:]
        out = [x]
        for f in self.features:
            out.append(F.interpolate(f(x), size=(h, w), mode='bilinear', align_corners=False))
        out = torch.cat(out, 1)
        return self.bottleneck(out)

class DINOv2UPerNet(nn.Module):
    def __init__(self, n_classes=10):
        super().__init__()
        self.embed_dim = 768

        # ── Backbone ──
        self.backbone = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitb14', pretrained=True)
        for param in self.backbone.parameters():
            param.requires_grad_(False)
            
        for block in self.backbone.blocks[-4:]:
            for param in block.parameters():
                param.requires_grad_(True)

        if hasattr(self.backbone, 'set_grad_checkpointing'):
            self.backbone.set_grad_checkpointing(True)
        else:
            self._enable_grad_checkpointing()

        # ── UPerNet Decoder ──
        channels = 256
        self.ppm = PPM(self.embed_dim, channels, bins=(1, 2, 3, 6))

        self.fpn_in = nn.ModuleList([
            nn.Sequential(nn.Conv2d(self.embed_dim, channels, 1, bias=False),
                          nn.BatchNorm2d(channels), nn.ReLU(inplace=True))
            for _ in range(3)
        ])

        self.fpn_out = nn.ModuleList([
            nn.Sequential(nn.Conv2d(channels, channels, 3, padding=1, bias=False),
                          nn.BatchNorm2d(channels), nn.ReLU(inplace=True))
            for _ in range(3)
        ])

        self.fpn_bottleneck = nn.Sequential(
            nn.Conv2d(channels * 4, channels, 3, padding=1, bias=False),
            nn.BatchNorm2d(channels),
            nn.ReLU(inplace=True)
        )

        self.final_head = nn.Conv2d(channels, n_classes, 1)

    def _enable_grad_checkpointing(self):
        from torch.utils.checkpoint import checkpoint
        original_blocks = self.backbone.blocks
        class CheckpointedBlock(nn.Module):
            def __init__(self, block):
                super().__init__()
                self.block = block
            def forward(self, x):
                return checkpoint(self.block, x, use_reentrant=False)
                
        for i in range(len(original_blocks) - 4, len(original_blocks)):
            original_blocks[i] = CheckpointedBlock(original_blocks[i])

    def _extract_features(self, x):
        features = self.backbone.get_intermediate_layers(x, n=4, return_class_token=False)
        B, _, H, W = x.shape
        out = []
        for feat in features:
            feat = feat.reshape(B, H//14, W//14, self.embed_dim).permute(0, 3, 1, 2)
            out.append(feat)
        return out

    def _fpn_forward(self, f1, f2, f3, f4):
        p4 = self.ppm(f4)

        p3 = self.fpn_in[2](f3) + F.interpolate(p4, size=f3.shape[2:], mode='bilinear', align_corners=False)
        p3 = self.fpn_out[2](p3)

        p2 = self.fpn_in[1](f2) + F.interpolate(p3, size=f2.shape[2:], mode='bilinear', align_corners=False)
        p2 = self.fpn_out[1](p2)

        p1 = self.fpn_in[0](f1) + F.interpolate(p2, size=f1.shape[2:], mode='bilinear', align_corners=False)
        p1 = self.fpn_out[0](p1)

        p4_up = F.interpolate(p4, size=p1.shape[2:], mode='bilinear', align_corners=False)
        p3_up = F.interpolate(p3, size=p1.shape[2:], mode='bilinear', align_corners=False)
        p2_up = F.interpolate(p2, size=p1.shape[2:], mode='bilinear', align_corners=False)

        fpn_out = torch.cat([p1, p2_up, p3_up, p4_up], dim=1)
        return self.fpn_bottleneck(fpn_out)

    def forward(self, x):
        target_h, target_w = x.shape[2], x.shape[3]
        f1, f2, f3, f4 = self._extract_features(x)

        from torch.utils.checkpoint import checkpoint
        if self.training:
            out = checkpoint(self._fpn_forward, f1, f2, f3, f4, use_reentrant=False)
        else:
            out = self._fpn_forward(f1, f2, f3, f4)

        out = self.final_head(out)
        return F.interpolate(out, size=(target_h, target_w), mode='bilinear', align_corners=False)

# ═══════════════════════════════════════════════════════════════
# LOSSES & OHEM
# ═══════════════════════════════════════════════════════════════
class OHEMCELoss(nn.Module):
    def __init__(self, keep_ratio=0.3, class_weights=None, device=None, ignore_index=9):
        super().__init__()
        self.keep_ratio = keep_ratio
        self.ignore_index = ignore_index
        if class_weights is not None and device is not None:
             self.criterion = nn.CrossEntropyLoss(weight=class_weights.to(device), reduction='none', ignore_index=ignore_index, label_smoothing=0.1)
        else:
             self.criterion = nn.CrossEntropyLoss(reduction='none', ignore_index=ignore_index, label_smoothing=0.1)

    def forward(self, pred, target):
        loss = self.criterion(pred, target)
        valid_mask = (target != self.ignore_index)
        loss = loss[valid_mask]

        if loss.numel() == 0:
            return loss.sum()

        num_keep = int(self.keep_ratio * loss.numel())
        if num_keep > 0:
            loss, _ = torch.topk(loss, num_keep)
        return loss.mean()

class BoundaryLoss(nn.Module):
    def __init__(self):
        super().__init__()
        sobel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32).unsqueeze(0).unsqueeze(0)
        sobel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32).unsqueeze(0).unsqueeze(0)
        self.register_buffer('sobel_x', sobel_x)
        self.register_buffer('sobel_y', sobel_y)

    def _sobel_edges(self, x):
        B, C, H, W = x.shape
        x_flat = x.reshape(B * C, 1, H, W)
        gx = F.conv2d(x_flat, self.sobel_x.to(device=x.device, dtype=x.dtype), padding=1)
        gy = F.conv2d(x_flat, self.sobel_y.to(device=x.device, dtype=x.dtype), padding=1)
        edges = torch.sqrt(gx ** 2 + gy ** 2 + 1e-8)
        return edges.reshape(B, C, H, W)

    def forward(self, pred_logits, target):
        pred_prob = F.softmax(pred_logits, dim=1)
        target_oh = F.one_hot(target, num_classes=pred_logits.shape[1]).permute(0, 3, 1, 2).float()
        pred_edges = self._sobel_edges(pred_prob)
        target_edges = self._sobel_edges(target_oh)
        return F.mse_loss(pred_edges, target_edges)

class WeightedLovaszLoss(nn.Module):
    def __init__(self, class_weights, device, ignore_index=9):
        super().__init__()
        self.binary_lovasz = smp.losses.LovaszLoss(mode='binary', from_logits=False)
        self.class_weights = class_weights.to(device)
        self.ignore_index = ignore_index

    def forward(self, pred, target):
        probs = F.softmax(pred, dim=1)
        B, C, H, W = probs.shape
        loss = 0.0
        weight_sum = 0.0

        valid_mask = (target != self.ignore_index)

        for c in range(C):
            if c == self.ignore_index:
                continue

            target_c = (target == c).float()
            prob_c = probs[:, c, :, :]

            p_valid = prob_c[valid_mask]
            t_valid = target_c[valid_mask]

            if t_valid.numel() == 0:
                continue

            l_c = self.binary_lovasz(p_valid.view(1, 1, -1), t_valid.view(1, 1, -1))
            w = self.class_weights[c]
            loss += l_c * w
            weight_sum += w

        return loss / (weight_sum + 1e-8)

class CompositeLoss(nn.Module):
    def __init__(self, class_weights, device):
        super().__init__()
        self.ohem_ce = OHEMCELoss(keep_ratio=0.3, class_weights=class_weights, device=device, ignore_index=9)
        self.lovasz = WeightedLovaszLoss(class_weights=class_weights, device=device, ignore_index=9)
        self.focal = smp.losses.FocalLoss(mode='multiclass', gamma=2.0, ignore_index=9)
        self.boundary = BoundaryLoss()

    def forward(self, pred, target):
        ohem_loss = self.ohem_ce(pred, target)
        lovasz_loss = self.lovasz(pred, target)
        focal_loss = self.focal(pred, target)
        boundary_loss = self.boundary(pred, target)
        return 0.3 * ohem_loss + 0.4 * lovasz_loss + 0.2 * focal_loss + 0.1 * boundary_loss

# ═══════════════════════════════════════════════════════════════
# LOOKAHEAD OPTIMIZER
# ═══════════════════════════════════════════════════════════════
class Lookahead(torch.optim.Optimizer):
    def __init__(self, optimizer, k=5, alpha=0.5):
        self.optimizer = optimizer
        self.k = k
        self.alpha = alpha
        self.param_groups = self.optimizer.param_groups
        self.defaults = self.optimizer.defaults
        self.state = self.optimizer.state
        self._step_counter = 0

        self.slow_weights = {}
        for group in self.param_groups:
            for p in group['params']:
                self.slow_weights[p] = p.clone().detach()

    def step(self, closure=None):
        loss = self.optimizer.step(closure)
        self._step_counter += 1
        if self._step_counter % self.k == 0:
            with torch.no_grad():
                for group in self.param_groups:
                    for p in group['params']:
                        if p.grad is None:
                            continue
                        slow = self.slow_weights[p]
                        slow.add_(p - slow, alpha=self.alpha)
                        p.copy_(slow)
        return loss

    def zero_grad(self, set_to_none: bool = False):
        self.optimizer.zero_grad(set_to_none=set_to_none)

    def state_dict(self):
        return {
            'base_state': self.optimizer.state_dict(),
            'lookahead_step': self._step_counter,
            'slow_weights': self.slow_weights
        }

    def load_state_dict(self, state_dict):
        if 'base_state' in state_dict:
            self.optimizer.load_state_dict(state_dict['base_state'])
            self._step_counter = state_dict['lookahead_step']
            self.slow_weights = state_dict['slow_weights']
        else:
            self.optimizer.load_state_dict(state_dict)

# ═══════════════════════════════════════════════════════════════
# IOU
# ═══════════════════════════════════════════════════════════════
def compute_mean_iou(preds, targets, n_classes):
    ious = []
    for c in range(n_classes):
        pred_c = (preds == c)
        target_c = (targets == c)
        intersection = (pred_c & target_c).sum().float()
        union = (pred_c | target_c).sum().float()
        if union == 0: ious.append(float('nan'))
        else: ious.append((intersection / union).item())
    return float(np.nanmean(ious)), ious

# ═══════════════════════════════════════════════════════════════
# MAIN TRAINING
# ═══════════════════════════════════════════════════════════════
def main():
    print(f"Device: {DEVICE}")

    # Initial Progressive Resizing: Epochs 1-12 (252x448)
    train_ds = OffroadDataset(TRAIN_IMG_DIR, TRAIN_MASK_DIR, get_train_transforms(252, 448), is_train=True)
    val_ds = OffroadDataset(VAL_IMG_DIR, VAL_MASK_DIR, get_val_transforms(252, 448), is_train=False)

    from torch.utils.data import WeightedRandomSampler
    train_sampler = WeightedRandomSampler(
        weights=train_ds.sample_weights,
        num_samples=len(train_ds),
        replacement=True
    )

    train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=False, sampler=train_sampler,
                              num_workers=NUM_WORKERS, pin_memory=True, drop_last=True, persistent_workers=True)
    val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False,
                            num_workers=NUM_WORKERS, pin_memory=True, persistent_workers=True)

    model = DINOv2UPerNet(n_classes=N_CLASSES).to(DEVICE)

    backbone_params = []
    decoder_params = []
    for name, param in model.named_parameters():
        if not param.requires_grad: continue
        if 'backbone' in name: backbone_params.append(param)
        else: decoder_params.append(param)

    base_optimizer = torch.optim.AdamW([
        {'params': backbone_params, 'lr': 5e-6, 'weight_decay': 1e-4},
        {'params': decoder_params, 'lr': 1e-4, 'weight_decay': 1e-4},
    ])
    lookahead_optimizer = Lookahead(base_optimizer, k=5, alpha=0.5)

    update_steps_per_epoch = (len(train_loader) + ACCUMULATION_STEPS - 1) // ACCUMULATION_STEPS
    scheduler_backbone = torch.optim.lr_scheduler.CosineAnnealingLR(
        base_optimizer, T_max=N_EPOCHS, eta_min=1e-7
    )
    scheduler_decoder = torch.optim.lr_scheduler.OneCycleLR(
        base_optimizer, max_lr=[5e-6, 1e-4],
        epochs=N_EPOCHS, steps_per_epoch=update_steps_per_epoch, pct_start=0.2
    )

    criterion = CompositeLoss(CLASS_WEIGHTS, DEVICE).to(DEVICE)
    scaler = torch.amp.GradScaler('cuda')

    start_epoch = 0
    best_iou = 0.0
    patience_counter = 0
    top_k_checkpoints = []

    if os.path.exists(MODEL_SAVE_PATH):
        print(f"Resuming from {MODEL_SAVE_PATH}")
        ckpt = torch.load(MODEL_SAVE_PATH, weights_only=False, map_location=DEVICE)
        model.load_state_dict(ckpt['model_state_dict'])
        if 'optimizer_state_dict' in ckpt:
            lookahead_optimizer.load_state_dict(ckpt['optimizer_state_dict'])
        start_epoch = ckpt.get('epoch', 0) + 1
        best_iou = float(ckpt.get('best_iou', 0.0))
        patience_counter = ckpt.get('patience_counter', 0)
        top_k_checkpoints = ckpt.get('top_k_checkpoints', [])

    for epoch in range(start_epoch, N_EPOCHS):
        # ── Progressive Resizing Update ──
        if epoch == 12:
            print("Switching resolution to 896x504 (Progressive Resizing)")
            train_loader.dataset.transforms = get_train_transforms(504, 896)
            val_loader.dataset.transforms = get_val_transforms(504, 896)

        model.train()
        epoch_loss = 0.0
        n_batches = 0
        t0 = time.time()

        lookahead_optimizer.zero_grad()

        for batch_idx, (images, masks) in enumerate(train_loader):
            images = images.to(DEVICE, non_blocking=True)
            masks = masks.to(DEVICE, non_blocking=True)

            with torch.amp.autocast('cuda'):
                main_out = model(images)

            main_out = main_out.float()
            total_loss = criterion(main_out, masks) / ACCUMULATION_STEPS

            scaler.scale(total_loss).backward()

            if (batch_idx + 1) % ACCUMULATION_STEPS == 0 or (batch_idx + 1) == len(train_loader):
                scaler.unscale_(lookahead_optimizer.optimizer)
                grad_norm = torch.nn.utils.clip_grad_norm_(
                    [p for p in model.parameters() if p.requires_grad], max_norm=1.0)

                if grad_norm.isnan() or grad_norm.isinf():
                    lookahead_optimizer.zero_grad()
                    scaler.update()
                    continue

                scaler.step(lookahead_optimizer)
                scaler.update()

                lookahead_optimizer.zero_grad()
                scheduler_decoder.step()

            epoch_loss += total_loss.item() * ACCUMULATION_STEPS
            n_batches += 1

            if batch_idx % 50 == 0:
                lr_bb = lookahead_optimizer.param_groups[0]['lr']
                lr_dec = lookahead_optimizer.param_groups[1]['lr']
                print(f"  [{epoch+1}/{N_EPOCHS}] batch {batch_idx}/{len(train_loader)} loss={total_loss.item() * ACCUMULATION_STEPS:.4f} lr_bb={lr_bb:.2e} lr_dec={lr_dec:.2e}")

        scheduler_backbone.step()
        avg_loss = epoch_loss / max(n_batches, 1)

        # ── VALIDATION ──
        model.eval()
        all_preds = []
        all_targets = []
        with torch.no_grad():
            for images, masks in val_loader:
                images = images.to(DEVICE, non_blocking=True)
                masks = masks.to(DEVICE, non_blocking=True)
                outputs = model(images)
                preds = outputs.argmax(dim=1)
                all_preds.append(preds.cpu())
                all_targets.append(masks.cpu())

        all_preds = torch.cat(all_preds)
        all_targets = torch.cat(all_targets)
        val_iou, per_class_iou = compute_mean_iou(all_preds, all_targets, N_CLASSES)

        print(f"Epoch {epoch+1}/{N_EPOCHS} | Loss: {avg_loss:.4f} | mIoU: {val_iou:.4f} | Time: {time.time()-t0:.1f}s")

        state_copy = OrderedDict({k: v.cpu().clone() for k, v in model.state_dict().items()})
        top_k_checkpoints.append((val_iou, state_copy))
        top_k_checkpoints.sort(key=lambda x: x[0], reverse=True)
        if len(top_k_checkpoints) > 5:
            top_k_checkpoints = top_k_checkpoints[:5]

        if val_iou > best_iou:
            best_iou = val_iou
            patience_counter = 0
            os.makedirs(os.path.dirname(MODEL_SAVE_PATH), exist_ok=True)
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': lookahead_optimizer.state_dict(),
                'best_iou': best_iou,
                'patience_counter': patience_counter,
                'top_k_checkpoints': top_k_checkpoints,
                'per_class_iou': per_class_iou,
            }, MODEL_SAVE_PATH)
            print(f"  ★ New best model! IoU: {best_iou:.4f}")
        else:
            patience_counter += 1
            print(f"  No improvement ({patience_counter}/{PATIENCE})")

        if patience_counter >= PATIENCE:
            break

    # ── SWA ──
    print("\n" + "=" * 60)
    print("Computing SWA from top-5 checkpoints...")
    if len(top_k_checkpoints) >= 2:
        avg_state = OrderedDict()
        n_ckpts = len(top_k_checkpoints)
        ref_state = top_k_checkpoints[0][1]

        for k in ref_state.keys():
            if ref_state[k].is_floating_point():
                avg_state[k] = sum(ckpt[1][k].float() for ckpt in top_k_checkpoints) / n_ckpts
                avg_state[k] = avg_state[k].to(ref_state[k].dtype)
            else:
                avg_state[k] = ref_state[k].clone()

        torch.save({
            'model_state_dict': avg_state,
            'best_iou': float(best_iou),
        }, SWA_SAVE_PATH)
        print(f"SWA saved to {SWA_SAVE_PATH}")

if __name__ == '__main__':
    main()