PyTorch NIMA: Neural IMage Assessment

PyTorch implementation of Neural IMage Assessment by Hossein Talebi and Peyman Milanfar. You can learn more from this post at Google Research Blog.

Installing

git clone https://github.com/truskovskiyk/nima.pytorch.git 
cd nima.pytorch
virtualenv -p python3.6 env
source ./env/bin/activate
pip install -r requirements/linux_gpu.txt

or You can just use ready Dockerfile

Dataset

The model was trained on the AVA (Aesthetic Visual Analysis) dataset You can get it from here Here are some examples of images with theire scores

Model

Used MobileNetV2 architecture as described in the paper Inverted Residuals and Linear Bottlenecks: Mobile Networks for Classification, Detection and Segmentation.

Pre-train model

You can use this pretrain-model with

val_emd_loss = 0.079
test_emd_loss = 0.080

Deployment

Deployed model on heroku URL is https://neural-image-assessment.herokuapp.com/ You can use it for testing in Your own images, but pay attention, that's free service, so it cannot handel too many requests. Here is simple curl command to test deployment models

curl  -X POST -F "file=@123.jpg" https://neural-image-assessment.herokuapp.com/api/get_scores

Please use our swagger for interactive testing

Usage

Clean and prepare dataset

export PYTHONPATH=.
python nima/cli.py prepare_dataset --path_to_ava_txt ./DATA/ava/AVA.txt \
                                    --path_to_save_csv ./DATA/ava \
                                    --path_to_images ./DATA/images/

Train model

export PYTHONPATH=.
python nima/cli.py train_model --path_to_save_csv ./DATA/ava/ \
                                --path_to_images ./DATA/images \
                                --batch_size 16 \
                                --num_workers 2 \
                                --num_epoch 15 \
                                --init_lr 0.009 \
                                --experiment_dir_name firts0.009

Use tensorboard to tracking training progress

tensorboard --logdir .

Validate model on val and test datasets

export PYTHONPATH=.
python nima/cli.py validate_model --path_to_model_weight ./pretrain-model.pth \
                                    --path_to_save_csv ./DATA/ava \
                                    --path_to_images ./DATA/images \
                                    --batch_size 16 \
                                    --num_workers 4

Get scores for one image

export PYTHONPATH=.
python nima/cli.py get_image_score --path_to_model_weight ./pretrain-model.pth \
                                       --path_to_image test_image.jpg

Contributing

Contributing are welcome

License

This project is licensed under the MIT License - see the LICENSE file for details

abliznyuk/nima.pytorch