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 "[email protected]" https://neural-image-assessment.herokuapp.com/api/get_scores

Please use our swagger for interactive testing

Usage

export PYTHONPATH=.
export PATH_TO_AVA_TXT=/storage/DATA/ava/AVA.txt
export PATH_TO_IMAGES=/storage/DATA/images/
export PATH_TO_CSV=/storage/DATA/ava/
export BATCH_SIZE=16
export NUM_WORKERS=2
export NUM_EPOCH=50
export INIT_LR=0.0001
export EXPERIMENT_DIR_NAME=/storage/experiment_n0001

Clean and prepare dataset

python nima/cli.py prepare_dataset --path_to_ava_txt $PATH_TO_AVA_TXT \
                                    --path_to_save_csv $PATH_TO_CSV \
                                    --path_to_images $PATH_TO_IMAGES

Train model

python nima/cli.py train_model --path_to_save_csv $PATH_TO_CSV \
                                --path_to_images $PATH_TO_IMAGES \
                                --batch_size $BATCH_SIZE \
                                --num_workers $NUM_WORKERS \
                                --num_epoch $NUM_EPOCH \
                                --init_lr $INIT_LR \
                                --experiment_dir_name $EXPERIMENT_DIR_NAME

Use tensorboard to tracking training progress

tensorboard --logdir .

Validate model on val and test datasets

python nima/cli.py validate_model --path_to_model_weight ./pretrain-model.pth \
                                    --path_to_save_csv $PATH_TO_CSV \
                                    --path_to_images $PATH_TO_IMAGES \
                                    --batch_size $BATCH_SIZE \
                                    --num_workers $NUM_EPOCH

Get scores for one image

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

Acknowledgments

• neural-image-assessment in keras
• Neural-IMage-Assessment in pytorch
• pytorch-mobilenet-v2
• origin NIMA article
• origin MobileNetV2 article
• Post at Google Research Blog
• Heroku: Cloud Application Platform