| --- |
| license: mit |
| tags: |
| - art-classification |
| - stable-diffusion |
| - wikiart |
| --- |
| |
| # 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: |
| |
| ```python |
| { |
| "epoch": int, |
| "model_state": OrderedDict, # load with model.load_state_dict() |
| "opt_state": OrderedDict, |
| "val_acc": float, |
| } |
| ``` |
| |
| ## Loading example |
|
|
| ```python |
| 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`. |
|
|
