Image Composition Grading

A CNN that predicts the composition quality score (1–5) of a photograph, trained on the CADB dataset.

Model

• Architecture: ResNet-18 (pretrained on ImageNet) with a custom regression head (512 → 256 → 1), fine-tuning only layer4 and fc
• Training: MSE loss, Adam optimizer (lr=1e-4, weight decay=1e-4), ReduceLROnPlateau scheduler, early stopping (patience 5)
• Data augmentation: Random resized crop, horizontal flip, rotation (±15°), color jitter
• Split: 80% train / 10% val / 10% test

Results

Best checkpoint at epoch 11 (out of 20), selected by validation loss:

Metric	Value
Test MAE	0.3591
Test RMSE	0.4604
Test MSE	0.2120

The model converges to its lowest validation MSE around 10–15 epochs; training further only reduces train loss (overfitting). Early stopping preserves the best bias–variance trade-off.

Files

File	Description
grading_nn.py	Training and model definition
grading_model.pth	Trained model weights
requirements.txt	Python dependencies

Setup

Requires Python 3.9+.

pip install -r requirements.txt

The CADB dataset (CADB_Dataset/ with images/ and JSON annotations) must be in the same directory. Download it from the CADB repo.

Usage

Train from scratch:

python grading_nn.py

Load the pretrained model:

import torch
from grading_nn import CompositionRegressor

model = CompositionRegressor()
model.load_state_dict(torch.load("grading_model.pth", map_location="cpu"))
model.eval()