Harish-JHR/WikiArtClassifySD

ArtExtract Task 1 — Checkpoints

Trained classifiers for painting style, artist, and genre classification using Stable Diffusion U-Net activations as features (WikiArt dataset).

Models

Three model types, each trained for three tasks (style / artist / genre):

ConvLSTM — spatial Conv-LSTM over the 16×16 activation grid from SD down_blocks.2. Treats each spatial position as a token in a sequence; attention pooling selects which regions matter for the prediction.

MLP probe — simple MLP on global-avg-pooled activations (3840-d). Fast linear probe used to verify the quality of the SD features.

ResNet50 — ImageNet-pretrained ResNet50 fine-tuned on raw paintings. Baseline that doesn't use diffusion features.

Checkpoint format

Each best.pt is a standard PyTorch checkpoint dict:

{
    "epoch":       int,
    "model_state": OrderedDict,   # load with model.load_state_dict()
    "opt_state":   OrderedDict,
    "val_acc":     float,
}

Loading example

import torch
from src.models import build_model

ckpt = torch.load("convlstm_style/best.pt", map_location="cpu")
model = build_model("convlstm", num_classes=27)
model.load_state_dict(ckpt["model_state"])
model.eval()
print(f"Loaded from epoch {ckpt['epoch']}, val_acc={ckpt['val_acc']:.4f}")

Training

SD features extracted from valhalla/sd-wikiart-v2 at timestep t=200, hooking down_blocks.2, mid_block, up_blocks.1.