This is a PyTorch implementation of the paper NIMA: Neural IMage Assessment (accepted at IEEE Transactions on Image Processing) by Hossein Talebi and Peyman Milanfar. You can learn more from this post at Google Research Blog.
-
The model was trained on the AVA (Aesthetic Visual Analysis) dataset, which contains roughly 255,500 images. You can get it from here. Note: there may be some corrupted images in the dataset, remove them first before you start training.
-
The dataset is into 229,981 images for training, 12,691 images for validation and 12,818 images for testing.
-
An ImageNet pretrained VGG-16 is used as the base network of the model, for which I got a ~0.072 EMD loss on the validation set. Haven't tried the other two options (MobileNet and Inception-v2) in the paper. You are very welcome to make your own extensions.
-
The learning rate setting differs from the original paper. I can't seem to get the model to converge with momentum SGD using an lr of 3e-7 for the conv base and 3e-6 for the dense block. Also I didn't do much hyper-param tuning therefore you could probably get better results. Other settings are all directly mirrored from the paper.
It is recommeded to use conda to manage your env. For example do
conda create -n nima python=3.6
conda activate nima
pip install -r requirements.txt
to install the dependancies.
To start training on the AVA dataset, first download the dataset from the link above and decompress which should create a directory named images/
. Then download the curated annotation CSVs below
which already splits the dataset (You can create your own split of course). Then do
python main.py --img_path /path/to/images/ --train --train_csv_file /path/to/train_labels.csv --val_csv_file /path/to/val_labels.csv --conv_base_lr 3e-4 --dense_lr 3e-3 --decay --ckpt_path /path/to/ckpts --epochs 100 --early_stoppping_patience 10
For inference, here the predicted score mean and std is generated. See predictions/
for an example format.
python test.py --model /path/to/your_model --test_csv /path/to/test_labels.csv --test_images /path/to/images --predictions /path/to/save/predictions
Training is done with early stopping monitoring. Here I set early_stopping_patience=10
.
- Here shows the predicted mean scores of some images from the validation set. Each image title starts with ground-truth rating followed by the predicted mean and std in the parentheses.
- Also some failure cases...
- The predicted aesthetic ratings from training on the AVA dataset are sensitive to contrast adjustments. Below images from left to right in a row-major order are with progressively sharper contrast, with lower leftmost being the original input. Contrast adjustment is done using
ImageEnhance
fromPIL
.
MIT