The Five Artists

Confirm the mapping and counts for all five artist IDs before training; preserve image IDs for submission.

Compare class counts and inspect representative images per artist rather than only the first few rows.

Use a stratified validation split and preserve spatial structure with a CNN.

Track validation accuracy, then refit on all 10,000 labeled images using the chosen epoch count and match the submission column names exactly.

Learn the five subtle drawing styles with a convolutional classifier; the complete training pipeline is included, but the accuracy remains unverified because the image files could not be downloaded.

#!/usr/bin/env python3
"""Train a CNN and write a Kaggle submission for Colombia's Five Artists task.

The script downloads the official Hugging Face datasets. It first uses a
stratified holdout to estimate accuracy and select the epoch count, then trains
on all labeled examples before predicting the 2,000 test images.
"""

from __future__ import annotations

import argparse
import copy
import random
from pathlib import Path

import numpy as np
import pandas as pd
import torch
from PIL import Image
from sklearn.metrics import accuracy_score, classification_report, f1_score
from sklearn.model_selection import train_test_split
from torch import nn
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms

TRAIN_DATASET = "eleon360/five-artists-dataset"
TEST_DATASET = "eleon360/five-artists-test-dataset"


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


class ArtistImages(Dataset):
    def __init__(self, source, indices, size: int, augment: bool):
        self.source = source
        self.indices = list(map(int, indices))
        base = [transforms.Resize((size, size), antialias=True)]
        if augment:
            base.extend(
                [
                    transforms.RandomRotation(degrees=6),
                    transforms.ColorJitter(brightness=0.10, contrast=0.10, saturation=0.08, hue=0.02),
                ]
            )
        base.extend(
            [
                transforms.ToTensor(),
                transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.25, 0.25, 0.25)),
            ]
        )
        self.transform = transforms.Compose(base)

    def __len__(self) -> int:
        return len(self.indices)

    def __getitem__(self, item: int):
        row = self.source[self.indices[item]]
        image = row["image"]
        if not isinstance(image, Image.Image):
            image = Image.fromarray(np.asarray(image))
        image = image.convert("RGB")
        tensor = self.transform(image)
        if "artist_id" in row:
            return tensor, int(row["artist_id"])
        return tensor, int(row["image_id"])


class ResidualBlock(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, stride: int):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.skip = (
            nn.Identity()
            if stride == 1 and in_channels == out_channels
            else nn.Sequential(
                nn.Conv2d(in_channels, out_channels, 1, stride=stride, bias=False),
                nn.BatchNorm2d(out_channels),
            )
        )
        self.activation = nn.ReLU(inplace=True)

    def forward(self, x):
        residual = self.skip(x)
        x = self.activation(self.bn1(self.conv1(x)))
        x = self.bn2(self.conv2(x))
        return self.activation(x + residual)


class ArtistCNN(nn.Module):
    def __init__(self, classes: int = 5):
        super().__init__()
        self.stem = nn.Sequential(
            nn.Conv2d(3, 48, 3, padding=1, bias=False),
            nn.BatchNorm2d(48),
            nn.ReLU(inplace=True),
        )
        self.features = nn.Sequential(
            ResidualBlock(48, 64, 2),
            ResidualBlock(64, 96, 2),
            ResidualBlock(96, 160, 2),
            ResidualBlock(160, 256, 2),
        )
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.head = nn.Sequential(nn.Flatten(), nn.Dropout(0.30), nn.Linear(256, classes))

    def forward(self, x):
        return self.head(self.pool(self.features(self.stem(x))))


def run_training(
    model: nn.Module,
    train_loader: DataLoader,
    valid_loader: DataLoader | None,
    device: torch.device,
    max_epochs: int,
    patience: int,
) -> tuple[nn.Module, int, dict[str, float] | None]:
    criterion = nn.CrossEntropyLoss(label_smoothing=0.04)
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max_epochs)
    best_state = copy.deepcopy(model.state_dict())
    best_epoch = max_epochs
    best_accuracy = -1.0
    best_metrics = None
    stale = 0

    for epoch in range(1, max_epochs + 1):
        model.train()
        total_loss = 0.0
        for images, labels in train_loader:
            images = images.to(device, non_blocking=True)
            labels = labels.to(device, dtype=torch.long, non_blocking=True)
            optimizer.zero_grad(set_to_none=True)
            logits = model(images)
            loss = criterion(logits, labels)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0)
            optimizer.step()
            total_loss += float(loss.detach()) * len(labels)
        scheduler.step()

        mean_loss = total_loss / max(len(train_loader.dataset), 1)
        if valid_loader is None:
            print(f"epoch={epoch:02d}/{max_epochs} train_loss={mean_loss:.4f}")
            best_state = copy.deepcopy(model.state_dict())
            continue

        y_true, y_pred = predict(model, valid_loader, device)
        accuracy = accuracy_score(y_true, y_pred)
        macro_f1 = f1_score(y_true, y_pred, average="macro", zero_division=0)
        print(f"epoch={epoch:02d} train_loss={mean_loss:.4f} valid_accuracy={accuracy:.5f} macro_f1={macro_f1:.5f}")
        if accuracy > best_accuracy:
            best_accuracy = float(accuracy)
            best_epoch = epoch
            best_state = copy.deepcopy(model.state_dict())
            best_metrics = {"accuracy": float(accuracy), "macro_f1": float(macro_f1)}
            stale = 0
        else:
            stale += 1
            if epoch >= 8 and stale >= patience:
                break

    model.load_state_dict(best_state)
    return model, best_epoch, best_metrics


@torch.no_grad()
def predict(model: nn.Module, loader: DataLoader, device: torch.device) -> tuple[list[int], list[int]]:
    model.eval()
    targets: list[int] = []
    predictions: list[int] = []
    for images, labels in loader:
        logits = model(images.to(device, non_blocking=True))
        predictions.extend(logits.argmax(dim=1).cpu().tolist())
        targets.extend(torch.as_tensor(labels).cpu().tolist())
    return targets, predictions


@torch.no_grad()
def predict_test(model: nn.Module, loader: DataLoader, device: torch.device) -> tuple[list[int], list[int]]:
    model.eval()
    image_ids: list[int] = []
    predictions: list[int] = []
    for images, ids in loader:
        logits = model(images.to(device, non_blocking=True))
        predictions.extend(logits.argmax(dim=1).cpu().tolist())
        image_ids.extend(torch.as_tensor(ids).cpu().tolist())
    return image_ids, predictions


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--train-dataset", default=TRAIN_DATASET)
    parser.add_argument("--test-dataset", default=TEST_DATASET)
    parser.add_argument("--output", type=Path, default=Path("submission.csv"))
    parser.add_argument("--checkpoint", type=Path, default=Path("artist_cnn.pt"))
    parser.add_argument("--size", type=int, default=128)
    parser.add_argument("--batch-size", type=int, default=64)
    parser.add_argument("--epochs", type=int, default=30)
    parser.add_argument("--patience", type=int, default=6)
    parser.add_argument("--workers", type=int, default=2)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--prediction-column", default="artist_id", help="Use the Kaggle data-page target column")
    args = parser.parse_args()

    from datasets import load_dataset

    seed_everything(args.seed)
    train_data = load_dataset(args.train_dataset, split="train")
    test_data = load_dataset(args.test_dataset, split="train")
    labels = np.asarray(train_data["artist_id"], dtype=np.int64)
    if not set(np.unique(labels)).issubset(set(range(5))):
        raise ValueError("Expected artist_id labels 0 through 4")
    counts = np.bincount(labels, minlength=5)
    print(f"train={len(train_data):,}; test={len(test_data):,}; image=256x256 per Hugging Face viewer")
    print(f"artist_id class counts={counts.tolist()}")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"device={device}")
    indices = np.arange(len(train_data))
    train_idx, valid_idx = train_test_split(
        indices, test_size=0.15, random_state=args.seed, stratify=labels
    )
    pin = device.type == "cuda"
    train_loader = DataLoader(
        ArtistImages(train_data, train_idx, args.size, augment=True),
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=args.workers,
        pin_memory=pin,
    )
    valid_loader = DataLoader(
        ArtistImages(train_data, valid_idx, args.size, augment=False),
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=args.workers,
        pin_memory=pin,
    )

    model = ArtistCNN().to(device)
    model, best_epoch, metrics = run_training(
        model, train_loader, valid_loader, device, args.epochs, args.patience
    )
    y_valid, y_valid_hat = predict(model, valid_loader, device)
    print(f"best_epoch={best_epoch}; holdout_accuracy={accuracy_score(y_valid, y_valid_hat):.6f}")
    print(f"holdout_macro_f1={f1_score(y_valid, y_valid_hat, average='macro', zero_division=0):.6f}")
    print(classification_report(y_valid, y_valid_hat, labels=list(range(5)), zero_division=0))

    # Refit on all labeled images for the number of epochs selected by holdout accuracy.
    seed_everything(args.seed)
    full_loader = DataLoader(
        ArtistImages(train_data, indices, args.size, augment=True),
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=args.workers,
        pin_memory=pin,
    )
    final_model = ArtistCNN().to(device)
    final_model, _, _ = run_training(final_model, full_loader, None, device, best_epoch, args.patience)
    args.checkpoint.parent.mkdir(parents=True, exist_ok=True)
    torch.save(
        {"state_dict": final_model.state_dict(), "size": args.size, "best_epoch": best_epoch, "seed": args.seed},
        args.checkpoint,
    )

    test_loader = DataLoader(
        ArtistImages(test_data, np.arange(len(test_data)), args.size, augment=False),
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=args.workers,
        pin_memory=pin,
    )
    image_ids, test_predictions = predict_test(final_model, test_loader, device)
    output = pd.DataFrame({"image_id": image_ids, args.prediction_column: test_predictions})
    output.to_csv(args.output, index=False)
    print(f"wrote {len(output):,} predictions to {args.output}")


if __name__ == "__main__":
    main()

Download solution.py