---
base_model: facebook/dinov2-small
tags:
  - pytorch
  - dino
  - dinov2
  - regression
  - aesthetics
---

# DINOv2-small + MLPHead Aesthetic Scorer

A lightweight aesthetic regression model built on a **frozen** `facebook/dinov2-small` backbone
with a trainable MLP head that predicts an aesthetic score from image embeddings.

## Architecture

| Component | Details |
|-----------|----------|
| Backbone  | `facebook/dinov2-small` (frozen, not included in this checkpoint) |
| Input     | CLS token — shape `(B, 384)` |
| Head      | `Linear(384->256) -> GELU -> Dropout(0.3) -> Linear(256->1)` |
| Output    | Scalar aesthetic score per image |

## Checkpoint Contents

This `.pt` file contains **only the MLPHead state dict** (4 tensors).
The DINOv2 backbone is loaded separately from `facebook/dinov2-small`.

## Usage

```python
import torch
import torch.nn as nn
from transformers import AutoImageProcessor, Dinov2Model
from huggingface_hub import hf_hub_download
from PIL import Image

class MLPHead(nn.Module):
    def __init__(self, embed_dim=384, hidden_dim=256, dropout_p=0.3):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(embed_dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout_p),
            nn.Linear(hidden_dim, 1),
        )
    def forward(self, x):
        return self.net(x).squeeze(-1)

# Load backbone
processor = AutoImageProcessor.from_pretrained("facebook/dinov2-small")
backbone  = Dinov2Model.from_pretrained("facebook/dinov2-small").eval()

# Load head
ckpt_path = hf_hub_download(repo_id="grantmwilkinson/dinov2-small-mlphead-aesthetic", filename="dinov2-small_MLPHead_best.pt")
head = MLPHead()
head.load_state_dict(torch.load(ckpt_path, map_location="cpu", weights_only=True))
head.eval()

# Inference
image  = Image.open("your_image.jpg").convert("RGB")
inputs = processor(images=image, return_tensors="pt")

with torch.no_grad():
    cls_token = backbone(**inputs).last_hidden_state[:, 0]  # (1, 384)
    score     = head(cls_token)                              # (1,)

print(f"Aesthetic score: {score.item():.4f}")
```