src/train_custom_model.py

from __future__ import annotations

import argparse
import json
import random
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Dict, List, Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.metrics import classification_report
from torch.utils.data import DataLoader, random_split
from torchvision import datasets, models, transforms


@dataclass
class TrainConfig:
    data_dir: str = "data/pokemon"
    output_model_path: str = "models/custom_resnet18.pth"
    output_metrics_path: str = "reports/custom_model_metrics.json"
    batch_size: int = 16
    num_epochs: int = 8
    learning_rate: float = 1e-4
    weight_decay: float = 1e-4
    val_split: float = 0.2
    image_size: int = 224
    seed: int = 42


def set_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def get_device() -> torch.device:
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


def build_transforms(image_size: int) -> Tuple[transforms.Compose, transforms.Compose]:
    train_tfms = transforms.Compose(
        [
            transforms.RandomResizedCrop(image_size),
            transforms.RandomHorizontalFlip(),
            transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ]
    )
    eval_tfms = transforms.Compose(
        [
            transforms.Resize(int(image_size * 1.14)),
            transforms.CenterCrop(image_size),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ]
    )
    return train_tfms, eval_tfms


def build_dataloaders(config: TrainConfig) -> Tuple[DataLoader, DataLoader, DataLoader, List[str]]:
    train_tfms, eval_tfms = build_transforms(config.image_size)

    train_root = Path(config.data_dir) / "train"
    test_root = Path(config.data_dir) / "test"

    full_train_dataset = datasets.ImageFolder(train_root, transform=train_tfms)
    eval_train_dataset = datasets.ImageFolder(train_root, transform=eval_tfms)
    test_dataset = datasets.ImageFolder(test_root, transform=eval_tfms)

    num_train = len(full_train_dataset)
    num_val = int(num_train * config.val_split)
    num_train_final = num_train - num_val

    train_subset, val_subset_indices = random_split(
        full_train_dataset,
        [num_train_final, num_val],
        generator=torch.Generator().manual_seed(config.seed),
    )

    # Recreate val subset with eval transforms by reusing indices from split.
    val_indices = val_subset_indices.indices
    val_subset = torch.utils.data.Subset(eval_train_dataset, val_indices)

    train_loader = DataLoader(train_subset, batch_size=config.batch_size, shuffle=True, num_workers=0)
    val_loader = DataLoader(val_subset, batch_size=config.batch_size, shuffle=False, num_workers=0)
    test_loader = DataLoader(test_dataset, batch_size=config.batch_size, shuffle=False, num_workers=0)

    classes = full_train_dataset.classes
    return train_loader, val_loader, test_loader, classes


def create_model(num_classes: int) -> nn.Module:
    model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
    in_features = model.fc.in_features
    model.fc = nn.Linear(in_features, num_classes)
    return model


def evaluate(model: nn.Module, loader: DataLoader, device: torch.device) -> Tuple[float, List[int], List[int]]:
    model.eval()
    correct = 0
    total = 0
    preds_all: List[int] = []
    labels_all: List[int] = []

    with torch.no_grad():
        for images, labels in loader:
            images = images.to(device)
            labels = labels.to(device)
            outputs = model(images)
            _, preds = torch.max(outputs, dim=1)

            total += labels.size(0)
            correct += (preds == labels).sum().item()
            preds_all.extend(preds.cpu().numpy().tolist())
            labels_all.extend(labels.cpu().numpy().tolist())

    acc = correct / total if total > 0 else 0.0
    return acc, preds_all, labels_all


def train(config: TrainConfig) -> Dict[str, object]:
    set_seed(config.seed)
    device = get_device()

    train_loader, val_loader, test_loader, classes = build_dataloaders(config)

    model = create_model(num_classes=len(classes)).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=config.learning_rate, weight_decay=config.weight_decay)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=4, gamma=0.5)

    best_val_acc = 0.0
    best_state = None
    history = []

    for epoch in range(config.num_epochs):
        model.train()
        running_loss = 0.0
        running_correct = 0
        running_total = 0

        for images, labels in train_loader:
            images = images.to(device)
            labels = labels.to(device)

            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item() * labels.size(0)
            _, preds = torch.max(outputs, dim=1)
            running_correct += (preds == labels).sum().item()
            running_total += labels.size(0)

        scheduler.step()

        train_loss = running_loss / max(1, running_total)
        train_acc = running_correct / max(1, running_total)
        val_acc, _, _ = evaluate(model, val_loader, device)

        history.append(
            {
                "epoch": epoch + 1,
                "train_loss": round(train_loss, 5),
                "train_acc": round(train_acc, 5),
                "val_acc": round(val_acc, 5),
            }
        )

        print(
            f"Epoch {epoch + 1}/{config.num_epochs} "
            f"- train_loss: {train_loss:.4f} "
            f"- train_acc: {train_acc:.4f} "
            f"- val_acc: {val_acc:.4f}"
        )

        if val_acc > best_val_acc:
            best_val_acc = val_acc
            best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}

    if best_state is not None:
        model.load_state_dict(best_state)

    test_acc, test_preds, test_labels = evaluate(model, test_loader, device)

    target_names = classes
    cls_report = classification_report(
        test_labels,
        test_preds,
        target_names=target_names,
        output_dict=True,
        zero_division=0,
    )

    output_model = Path(config.output_model_path)
    output_model.parent.mkdir(parents=True, exist_ok=True)
    torch.save(
        {
            "state_dict": model.state_dict(),
            "labels": classes,
            "image_size": config.image_size,
            "architecture": "resnet18",
        },
        output_model,
    )

    result = {
        "config": asdict(config),
        "device": str(device),
        "num_classes": len(classes),
        "labels": classes,
        "best_val_acc": round(best_val_acc, 5),
        "test_acc": round(test_acc, 5),
        "history": history,
        "classification_report": cls_report,
        "model_path": str(output_model),
    }

    output_metrics = Path(config.output_metrics_path)
    output_metrics.parent.mkdir(parents=True, exist_ok=True)
    output_metrics.write_text(json.dumps(result, indent=2), encoding="utf-8")

    print(f"Saved model to: {output_model}")
    print(f"Saved metrics to: {output_metrics}")
    print(f"Final test accuracy: {test_acc:.4f}")

    return result


def parse_args() -> TrainConfig:
    parser = argparse.ArgumentParser(description="Train transfer learning classifier on custom image data.")
    parser.add_argument("--data-dir", default="data/pokemon")
    parser.add_argument("--output-model", default="models/custom_resnet18.pth")
    parser.add_argument("--output-metrics", default="reports/custom_model_metrics.json")
    parser.add_argument("--batch-size", type=int, default=16)
    parser.add_argument("--epochs", type=int, default=8)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--weight-decay", type=float, default=1e-4)
    parser.add_argument("--val-split", type=float, default=0.2)
    parser.add_argument("--image-size", type=int, default=224)
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    return TrainConfig(
        data_dir=args.data_dir,
        output_model_path=args.output_model,
        output_metrics_path=args.output_metrics,
        batch_size=args.batch_size,
        num_epochs=args.epochs,
        learning_rate=args.lr,
        weight_decay=args.weight_decay,
        val_split=args.val_split,
        image_size=args.image_size,
        seed=args.seed,
    )


if __name__ == "__main__":
    cfg = parse_args()
    train(cfg)