grading_nn.py

import json
import os

import numpy as np
import torch
import torch.nn as nn
from PIL import Image
from sklearn.model_selection import train_test_split
from torch.utils.data import Dataset, DataLoader
from torchvision import models, transforms

DATASET_ROOT = os.path.join(os.path.dirname(__file__), "CADB_Dataset")

def main():
    torch.manual_seed(42)
    np.random.seed(42)

    with open(os.path.join(DATASET_ROOT, "composition_scores.json")) as f:
        scores_data = json.load(f)
    image_dir = os.path.join(DATASET_ROOT, "images")

    all_names = list(scores_data.keys())
    all_scores = [scores_data[n]["mean"] for n in all_names]

    X_train, X_tmp, y_train, y_tmp = train_test_split(all_names, all_scores, test_size=0.20, random_state=42)
    X_val, X_test, y_val, y_test = train_test_split(X_tmp, y_tmp, test_size=0.50, random_state=42)

    print(f"Train: {len(X_train)}  Val: {len(X_val)}  Test: {len(X_test)}")

    train_tf = transforms.Compose([
        transforms.RandomResizedCrop(224, scale=(0.7, 1.0)),
        transforms.RandomHorizontalFlip(),
        transforms.RandomRotation(15),
        transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ])
    eval_tf = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ])

    train_ds = CADBDataset(X_train, y_train, image_dir, train_tf)
    val_ds = CADBDataset(X_val, y_val, image_dir, eval_tf)
    test_ds = CADBDataset(X_test, y_test, image_dir, eval_tf)

    train_loader = DataLoader(train_ds, batch_size=32, shuffle=True, num_workers=2, pin_memory=True)
    val_loader = DataLoader(val_ds, batch_size=32, shuffle=False, num_workers=2, pin_memory=True)
    test_loader = DataLoader(test_ds, batch_size=32, shuffle=False, num_workers=2, pin_memory=True)

    device = "mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Device: {device}\n")

    model = CompositionRegressor()
    train_model(model, train_loader, val_loader, epochs=20, lr=1e-4, device=device)

    model.eval()
    all_preds, all_targets = [], []
    with torch.no_grad():
        for imgs, scores in test_loader:
            all_preds.append(model(imgs.to(device)).cpu())
            all_targets.append(scores)

    preds = torch.cat(all_preds)
    targets = torch.cat(all_targets)
    mae = (preds - targets).abs().mean().item()
    rmse = ((preds - targets) ** 2).mean().sqrt().item()

    print(f"\nTest MAE:  {mae:.4f}")
    print(f"Test RMSE: {rmse:.4f}")
    print(f"Test MSE:  {rmse**2:.4f}")

    torch.save(model.state_dict(), "composition_model.pth")
    print("Model saved as composition_model.pth")


class CADBDataset(Dataset):
    def __init__(self, image_names, scores, image_dir, transform=None):
        self.image_names = image_names
        self.scores = scores
        self.image_dir = image_dir
        self.transform = transform

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

    def __getitem__(self, i):
        img = Image.open(os.path.join(self.image_dir, self.image_names[i])).convert("RGB")
        if self.transform:
            img = self.transform(img)
        return img, torch.tensor(self.scores[i], dtype=torch.float32)


class CompositionRegressor(nn.Module):
    def __init__(self):
        super().__init__()
        backbone = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)

        for name, param in backbone.named_parameters():
            if not name.startswith("fc") and not name.startswith("layer4"):
                param.requires_grad = False

        backbone.fc = nn.Sequential(
            nn.Linear(backbone.fc.in_features, 256),
            nn.ReLU(),
            nn.Dropout(0.4),
            nn.Linear(256, 1)
        )
        self.net = backbone

    def forward(self, x):
        return self.net(x).squeeze(1)


def train_model(model, train_loader, val_loader, epochs=20, lr=1e-4, device="cpu", patience=5):
    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=1e-4)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=3, factor=0.5)
    model.to(device)

    best_val_loss = float("inf")
    best_weights = None
    epochs_no_improve = 0

    for epoch in range(epochs):
        model.train()
        train_loss = 0.0
        for imgs, scores in train_loader:
            imgs, scores = imgs.to(device), scores.to(device)
            optimizer.zero_grad()
            loss = criterion(model(imgs), scores)
            loss.backward()
            optimizer.step()
            train_loss += loss.item() * len(scores)
        train_loss /= len(train_loader.dataset)

        model.eval()
        val_loss = 0.0
        with torch.no_grad():
            for imgs, scores in val_loader:
                imgs, scores = imgs.to(device), scores.to(device)
                val_loss += criterion(model(imgs), scores).item() * len(scores)
        val_loss /= len(val_loader.dataset)

        lr_before = optimizer.param_groups[0]["lr"]
        scheduler.step(val_loss)
        lr_after = optimizer.param_groups[0]["lr"]

        improved = val_loss < best_val_loss
        if improved:
            best_val_loss = val_loss
            best_weights = {k: v.cpu().clone() for k, v in model.state_dict().items()}
            epochs_no_improve = 0
        elif lr_after < lr_before:
            epochs_no_improve = 0
        else:
            epochs_no_improve += 1

        print(f"Epoch {epoch+1:>3}/{epochs} | Train MSE: {train_loss:.4f} | Val MSE: {val_loss:.4f} | LR: {lr_after:.2e}{' *' if improved else ''}")

        if epochs_no_improve >= patience:
            print(f"\nEarly stopping at epoch {epoch+1} (no improvement for {patience} epochs)")
            break

    model.load_state_dict(best_weights)
    print(f"Restored best weights (Val MSE: {best_val_loss:.4f})")


if __name__ == "__main__":
    main()