train.py

import os
import yaml
import random
import numpy as np
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm
import albumentations as A

from .model import UNet
from .data import SegDataset
from .utils import BCEDiceLoss, dice_coef, iou_score


def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def get_transforms(size):
    return A.Compose([
        A.HorizontalFlip(p=0.5),
        A.VerticalFlip(p=0.1),
        A.RandomRotate90(p=0.5),
        A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.1, rotate_limit=15, p=0.5),
        A.GaussNoise(p=0.1),
        A.ColorJitter(p=0.2),
    ])


def train_one_epoch(model, loader, loss_fn, optimizer, device):
    model.train()
    epoch_loss = 0.0
    dms, ious = [], []
    for imgs, msks in tqdm(loader, desc='train', leave=False):
        imgs = torch.from_numpy(imgs).to(device)
        msks = torch.from_numpy(msks).to(device)
        optimizer.zero_grad()
        logits = model(imgs)
        loss = loss_fn(logits, msks)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()
        dms.append(dice_coef(logits.detach(), msks))
        ious.append(iou_score(logits.detach(), msks))
    return epoch_loss/len(loader), float(np.mean(dms)), float(np.mean(ious))


def eval_one_epoch(model, loader, loss_fn, device):
    model.eval()
    epoch_loss = 0.0
    dms, ious = [], []
    with torch.no_grad():
        for imgs, msks in tqdm(loader, desc='val', leave=False):
            imgs = torch.from_numpy(imgs).to(device)
            msks = torch.from_numpy(msks).to(device)
            logits = model(imgs)
            loss = loss_fn(logits, msks)
            epoch_loss += loss.item()
            dms.append(dice_coef(logits, msks))
            ious.append(iou_score(logits, msks))
    return epoch_loss/len(loader), float(np.mean(dms)), float(np.mean(ious))


def main(config_path:str):
    with open(config_path, 'r') as f:
        cfg = yaml.safe_load(f)

    # Defensive casts to avoid type issues
    seed = int(cfg['seed'])
    num_epochs = int(cfg['num_epochs'])
    batch_size = int(cfg['batch_size'])
    image_size = int(cfg['image_size'])
    lr = float(cfg['learning_rate'])
    wd = float(cfg['weight_decay'])
    num_workers = int(cfg['num_workers'])

    set_seed(seed)

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

    os.makedirs(cfg['paths']['checkpoints_dir'], exist_ok=True)

    aug = get_transforms(image_size)

    train_ds = SegDataset(cfg['paths']['train_images'], cfg['paths']['train_masks'], image_size=image_size, aug=aug)
    val_ds   = SegDataset(cfg['paths']['val_images'], cfg['paths']['val_masks'], image_size=image_size, aug=None)

    # Guard against empty datasets
    if len(train_ds) == 0:
        raise FileNotFoundError(f"No training samples found in {cfg['paths']['train_images']} and {cfg['paths']['train_masks']}")
    if len(val_ds) == 0:
        raise FileNotFoundError(f"No validation samples found in {cfg['paths']['val_images']} and {cfg['paths']['val_masks']}")

    def collate_fn(batch):
        imgs, msks = zip(*batch)
        import numpy as np
        return np.stack(imgs), np.stack(msks)

    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=num_workers, collate_fn=collate_fn)
    val_loader   = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=num_workers, collate_fn=collate_fn)

    model = UNet(
        in_channels=int(cfg['model']['in_channels']),
        out_channels=int(cfg['model']['out_channels']),
        features=tuple(cfg['model']['features'])
    ).to(device)

    loss_fn = BCEDiceLoss(float(cfg['loss']['bce_weight']), float(cfg['loss']['dice_weight']))
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=wd)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', patience=3, factor=0.5)

    best_iou = -1
    for epoch in range(num_epochs):
        tr_loss, tr_dice, tr_iou = train_one_epoch(model, train_loader, loss_fn, optimizer, device)
        va_loss, va_dice, va_iou = eval_one_epoch(model, val_loader, loss_fn, device)
        scheduler.step(va_iou)
        print(f"Epoch {epoch+1}/{num_epochs} | train loss {tr_loss:.4f} dice {tr_dice:.4f} iou {tr_iou:.4f} | val loss {va_loss:.4f} dice {va_dice:.4f} iou {va_iou:.4f}")
        if va_iou > best_iou:
            best_iou = va_iou
            best_path = os.path.join(cfg['paths']['checkpoints_dir'], 'best.pth')
            torch.save({'model_state': model.state_dict(), 'cfg': cfg}, best_path)
            print(f"Saved best to {best_path}")

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', type=str, default='configs/config.yaml')
    args = parser.parse_args()
    main(args.config)