train.py

import os
import cv2
import torch
import albumentations as A
from albumentations.pytorch import ToTensorV2
from tqdm import tqdm
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader

from model import UNET
from dataset import CarvanaDataset  # <- use the updated multi-class–safe dataset
from utils import (
    save_checkpoint,
    save_predictions_as_imgs,  # keep using yours for now (ensure it argmaxes logits)
)

# -----------------------
# Hyperparameters etc.
# -----------------------
LEARNING_RATE = 1e-4
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
BATCH_SIZE = 16
NUM_EPOCHS = 10
NUM_WORKERS = 2
IMAGE_HEIGHT = 160  # 1280 originally
IMAGE_WIDTH = 240   # 1918 originally
PIN_MEMORY = (DEVICE == "cuda")
LOAD_MODEL = False  # not used below, but kept for parity

TRAIN_IMG_DIR = "yamaha_seg/train_images/"
TRAIN_MASK_DIR = "yamaha_seg/train_masks/"
VAL_IMG_DIR   = "yamaha_seg/test_images/"
VAL_MASK_DIR  = "yamaha_seg/test_masks/"

# -----------------------
# Transforms (label-safe)
# -----------------------
train_transform = A.Compose([
    A.Resize(height=IMAGE_HEIGHT, width=IMAGE_WIDTH, interpolation=cv2.INTER_NEAREST),
    A.ShiftScaleRotate(
        shift_limit=0.0,     # no shift
        scale_limit=0.0,     # no scale
        rotate_limit=35,     # only rotation
        border_mode=cv2.BORDER_CONSTANT,
        value=0,             # image fill
        mask_value=0,        # mask fill (keeps labels valid)
        p=1.0,
    ),
    A.HorizontalFlip(p=0.5),
    A.VerticalFlip(p=0.1),
    A.Normalize(mean=(0,0,0), std=(1,1,1), max_pixel_value=255.0),
    ToTensorV2(),
])

val_transform = A.Compose([
    A.Resize(height=IMAGE_HEIGHT, width=IMAGE_WIDTH, interpolation=cv2.INTER_NEAREST),
    A.Normalize(mean=(0,0,0), std=(1,1,1), max_pixel_value=255.0),
    ToTensorV2(),
])

# -----------------------
# Multi-class accuracy
# -----------------------
@torch.no_grad()
def check_accuracy_mc(loader, model, device="cuda"):
    model.eval()
    correct, total = 0, 0
    n_classes = model.out_channels
    iou_sum, iou_count = 0.0, 0

    for x, y in loader:
        x = x.to(device)
        y = y.long().to(device)        # [B,H,W]
        logits = model(x)              # [B,C,H,W]
        preds = logits.argmax(1)       # [B,H,W]

        correct += (preds == y).sum().item()
        total   += y.numel()

        for c in range(n_classes):
            p = (preds == c)
            t = (y == c)
            inter = (p & t).sum().item()
            union = (p | t).sum().item()
            if union > 0:
                iou_sum += inter / union
                iou_count += 1

    pix_acc = correct / total if total else 0.0
    miou = (iou_sum / iou_count) if iou_count else 0.0
    print(f"Pixel Acc: {pix_acc:.4f} | mIoU: {miou:.4f}")
    model.train()

# -----------------------
# Label-range sanity check
# -----------------------
def assert_valid_targets(loader, n_classes):
    for i, (_, y) in enumerate(loader):
        y = y.long()
        bad = (y < 0) | (y >= n_classes)
        if bad.any():
            import torch
            uniq = torch.unique(y[bad]).cpu().tolist()
            raise RuntimeError(
                f"[SanityCheck] Batch {i}: found invalid labels {uniq} not in [0, {n_classes-1}]. "
                "Check dataset mapping and that all mask transforms use INTER_NEAREST with mask_value=0."
            )
    print(f"[SanityCheck] All targets valid in [0, {n_classes-1}].")

@torch.no_grad()
def check_accuracy_mc(loader, model, n_classes, device="cuda"):
    model.eval()
    correct, total = 0, 0
    iou_sum, iou_count = 0.0, 0

    for x, y in loader:
        x = x.to(device)
        y = y.long().to(device)      # [B,H,W]
        logits = model(x)            # [B,C,H,W]
        preds = logits.argmax(1)     # [B,H,W]

        correct += (preds == y).sum().item()
        total   += y.numel()

        for c in range(n_classes):
            p = (preds == c)
            t = (y == c)
            inter = (p & t).sum().item()
            union = (p | t).sum().item()
            if union > 0:
                iou_sum += inter / union
                iou_count += 1

    pix_acc = correct / total if total else 0.0
    miou = (iou_sum / iou_count) if iou_count else 0.0
    print(f"Pixel Acc: {pix_acc:.4f} | mIoU: {miou:.4f}")
    model.train()

# -----------------------
# Train step (AMP updated)
# -----------------------
def train_fn(loader, model, optimizer, loss_fn, scaler):
    loop = tqdm(loader, ncols=100)
    for batch_idx, (data, targets) in enumerate(loop):
        data = data.to(device=DEVICE)
        targets = targets.long().to(device=DEVICE)  # [B,H,W]; no squeeze if dataset returns HxW

        optimizer.zero_grad(set_to_none=True)
        with torch.amp.autocast("cuda", enabled=(DEVICE == "cuda")):
            logits = model(data)             # [B,C,H,W]
            loss = loss_fn(logits, targets)  # CE expects class indices

        if scaler is not None:
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
        else:
            loss.backward()
            optimizer.step()

        loop.set_postfix(loss=f"{loss.item():.4f}")

# -----------------------
# Main
# -----------------------
def main():
    # Datasets & Loaders
    train_ds = CarvanaDataset(TRAIN_IMG_DIR, TRAIN_MASK_DIR, transform=train_transform)
    val_ds   = CarvanaDataset(VAL_IMG_DIR,   VAL_MASK_DIR,   transform=val_transform)

    NUM_CLASSES = len(train_ds.value_map)
    print("Classes (value→index):", train_ds.value_map)

    train_loader = DataLoader(
        train_ds, batch_size=BATCH_SIZE, shuffle=True,
        num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY
    )
    val_loader = DataLoader(
        val_ds, batch_size=BATCH_SIZE, shuffle=False,
        num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY
    )

    # Model / loss / opt
    model = UNET(in_channels=3, out_channels=NUM_CLASSES).to(DEVICE)
    loss_fn = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)

    # AMP scaler (new API)
    scaler = torch.amp.GradScaler("cuda") if DEVICE == "cuda" else None

    # Optional: cudnn tuning
    torch.backends.cudnn.benchmark = True

    # Sanity check labels before training
    assert_valid_targets(train_loader, NUM_CLASSES)
    assert_valid_targets(val_loader,   NUM_CLASSES)

    # Initial eval
    check_accuracy_mc(val_loader, model, NUM_CLASSES, device=DEVICE)

    # Train
    for epoch in range(NUM_EPOCHS):
        print(f"\nEpoch {epoch+1}/{NUM_EPOCHS}")
        train_fn(train_loader, model, optimizer, loss_fn, scaler)

        # Save model
        checkpoint = {"state_dict": model.state_dict(), "optimizer": optimizer.state_dict()}
        save_checkpoint(checkpoint)

        # Eval
        check_accuracy_mc(val_loader, model, NUM_CLASSES, device=DEVICE)

        # Sample predictions (your util should argmax logits internally; if not, update it)
        save_predictions_as_imgs(val_loader, model, folder="saved_images/", device=DEVICE)

if __name__ == "__main__":
    main()