Need help with Neural-IMage-Assessment?
Click the “chat” button below for chat support from the developer who created it, or find similar developers for support.

About the developer

kentsyx
318 Stars 70 Forks Other 70 Commits 16 Opened issues

Description

A PyTorch Implementation of Neural IMage Assessment

Services available

!
?

Need anything else?

Contributors list

# 133,710
Python
image-e...
photo-e...
pytorch
68 commits
# 9,126
Python
explain...
explain...
mariadb
1 commit

NIMA: Neural IMage Assessment

Python 3.6+ MIT License

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.

Implementation Details

  • The model was trained on the AVA (Aesthetic Visual Analysis) dataset containing 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.~~ Use provided CSVs which have already done this for you.

  • Dataset is split 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. Should be easy to plug in the other two options (MobileNet and Inception-v2).

  • The learning rate setting differs from the original paper. Can't seem to get the model to converge using the original params. Also didn't do much hyper-param tuning therefore you could probably get better results. Other settings are all directly mirrored from the paper.

Requirements

Code is written using PyTorch 1.8.1 with CUDA 11.1. You can recreate the environment I used with conda by

conda env create -f env.yml

to install the dependancies.

Usage

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 5e-4 --dense_lr 5e-3 --decay --ckpt_path /path/to/ckpts --epochs 100 --early_stoppping_patience 10

For inference, do

python -W ignore test.py --model /path/to/your_model --test_csv /path/to/test_labels.csv --test_images /path/to/images --predictions /path/to/save/predictions

See

predictions/
for dumped predictions as an example.

Training Statistics

Training is done with early stopping. Here I set

early_stopping_patience=10
.

Pretrained Model

~0.069 EMD on validation. Not fully converged yet (constrained by resources). To continue training, download the pretrained weights and add

--warm_start --warm_start_epoch 34
to your args.

Google Drive

Annotation CSV Files

Train Validation Test

Example Results

  • Here first shows some good predictions from the test set. Each image title starts with ground-truth rating followed by the predicted mean and std in the parentheses.

  • Also some failure cases, it would seem that the model usually fails at images with low/high aesthetic ratings.

  • The predicted aesthetic ratings from training on the AVA dataset are sensitive to contrast adjustments, preferring images with higher contrast. Below top row is the reference image with contrast
    c=1.0
    , while bottom images are enhanced with contrast
    [0.25, 0.75, 1.25, 1.75]
    . Contrast adjustment is done using
    ImageEnhance.Contrast
    from
    PIL
    (in this case pillow-simd).

License

MIT

We use cookies. If you continue to browse the site, you agree to the use of cookies. For more information on our use of cookies please see our Privacy Policy.