We recently launched our website aestheticstest.com, which as the name suggests, is an app that predicts how aesthetic an image is with the help of AI. Users can upload photos and can see the result within seconds. In this post, we describe how we trained the AI model and build a web app to use it.
Image Aesthetic Assessment
Aesthetic image analysis has attracted much attention in recent years. Such a tool has a lot of valuable use cases, such as helping you to decide which profile picture to use on social media platforms or which photos to add to your website.
Note that how aesthetic a photo is partly subjective since everyone has other preferences, but you can try to predict the average score of how a group of people would rate an image.
There is a paper from Google named Neural Image Assessment (NIMA) that tries basically the same that we wanted to archive, but it is already a few years old, so we thought that it would be better to build a model with the newest technologies that would probably be more accurate and also Google didn’t release the source code.
Dataset
There are multiple datasets for aesthetic photo assessments. Our initial experiments showed that a lot of training data is needed to get usable results, therefore we decided to use the AVA (Aesthetic Visual Analysis) dataset that consists of more than 200,000 images with labels for aesthetics scores on a scale from 1 to 10. The labels are the average rating of 10+ persons for each image.
In addition to a big dataset, we used a pre-trained model, precisely Xception that was pre-trained on the imagenet dataset. We fine-tuned the last layers of this model with the just mentioned AVA dataset. With the combination of a big dataset and a pre-trained Xception model, we were able to build a model with sufficient accuracy.
TensorFlow Model
Deep Learning models, on the base of a Convolutional Neural Network, have become the state-of-the-art choice for visual image analysis such as image classification or segmentation. Therefore, it was a natural decision to also use such an approach for our image aesthetics rating model.
As mentioned in the previous section, the heart of the model is a pre-trained Xception model that we optimized with images from the AVA dataset.
It’s a powerful architecture, but the code for the model is pretty simple thanks to Keras and TensorFlow:
We froze all layers from the pre-trained model and added a few more layers that we trained using the Adam optimizer. We Trained on Google Colab for multiple epochs, which took us 1–2 hours. After that, we unfroze all layers and trained the whole network on a low learning rate for a few epochs.
Results
These are examples from the test dataset that were not used for training. The average error (MAE) is 0.48 which we think is pretty good, considering that the training labels are subjective to a certain extent, meaning 100% accuracy is not possible.
It is also interesting to see is that the technical photo quality has a significant impact on the result. The left picture is the original that ranks higher compared to the same photo that is blurry, noisy, or has poor lighting.
If you want to try out this model with your own photos, click here.
API and Website
All tools we use to operate the model and a simple website to interact with it are part of the Google Cloud Platform. We created a Docker Image that uses the trained TensorFlow model and Django REST to provide an API. This image is deployed to a Google Compute Engine.
BTW, Google offers other tools bundled under the name AI Platform to train and operate machine learning models. The reason that we didn’t use this approach is that the operational costs would have been much higher (The VM we created costs just $12 per month).
For the website, we use Google App Engine, which is a managed environment that can be used to easily host simple web applications built with Flask. The Flask app uses the REST API described above to calculate the aesthetics score. The best thing is that GAE is completely free to use for one website.
If you like our project and want to try it with your own images, check out the website aestheticstest.com :)
