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 envsource ./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.079test_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
export PYTHONPATH=.export PATH_TO_AVA_TXT=/storage/DATA/ava/AVA.txtexport PATH_TO_IMAGES=/storage/DATA/images/export PATH_TO_CSV=/storage/DATA/ava/export BATCH_SIZE=16export NUM_WORKERS=2export NUM_EPOCH=50export INIT_LR=0.0001export 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