Create a simple image classifier using Tensorflow
In this post, we are going to show you how to build an image classifier using deep learning library, created by Google, called TensorFlow. Here, for simplicity, we only demonstrate the CPU-only version.
Install Python 3.5.x & Pip
According to official document here, for Ubuntu users, we need Python and Pip both installed, if not, please enter following command:
$ sudo apt-get install python-pip python-dev
For Windows users, TensorFlow only supports version 3.5.x of Python. Here is the download link. Remember to set your system environment, adding python 3.5.x to path to directly use in command line.
For Ubuntu users, install TensorFlow by following commands according to your Python version:
$ pip install tensorflow # Python 2.7; CPU support support only
$ pip3 install tensorflow # Python 3.n; CPU support support only
For Windows users:
C:\> pip3 install --upgrade tensorflow
Validate Your Installation
For both users, invoke python from your shell as follows:
$ python # Ubuntu
C:\> python # Windows
Enter the following short program inside the python interactive shell:
>>> import tensorflow as tf
>>> hello = tf.constant('Hello, TensorFlow!')
>>> sess = tf.Session()
If you have problem in import tensorflow at first line, download and install Visual C++ Redistributable for Visual Studio 2015 then try again.
If the system outputs the following, then Tensorflow has been successfully installed:
Prepare Your Training Data
Training data is what you want your computer to learn, to recognize the objects in the images. We need at least two hundred images for each class. There is a super convenient Google plug-in called “Fatkun Batch Download Image” may help us downloading amount of images at once. Here is an example, I search for ‘husky’ from Google:
Place the images into different folder for each class as following:
Start Retraining Process
At first, get the latest sample code from git TensorFlow repo. The sample code will be in
/tensorflow/tensorflow/examples/image_retraining/retrain.py Retraining also known as Transform Learning train the customized graph based on Inception which is one of the best Google image classifier. It help us training more accuracy classifier with less effort according our input data.
For training, type the following command by invoking python:
While retrain begins, you will see some bottlenecks being created similar to the screenshot below. The bottlenecks are the last layer before the final layer of the Inception model which will classify your training images.
After the bottlenecks training is over, the final layer gets trained and you will see validation accuracy as shown below. At the end, you will see “Final test accuracy” and there would be two more files: “retrained_graph.pb” and “retrained_labels.txt” in your working directory.
Test Your Image Classifier
Go to your work directory and create a file called
label_image.py. Type the following lines of code and save it.
import tensorflow as tf, sys
image_path = sys.argv
# Read in the image_data
image_data = tf.gfile.FastGFile(image_path, 'rb').read()
# Loads label file, strips off carriage return
label_lines = [line.rstrip() for line
# Unpersists graph from file
with tf.gfile.FastGFile("D:/tensorflow_work/retrained_graph.pb", 'rb') as f:
graph_def = tf.GraphDef()
_ = tf.import_graph_def(graph_def, name='')
# Feed the image_data as input to the graph and get first prediction
with tf.Session() as sess:
softmax_tensor = sess.graph.get_tensor_by_name('final_result:0')
predictions = sess.run(softmax_tensor,
# Sort to show labels of first prediction in order of confidence
top_k = predictions.argsort()[-len(predictions):][::-1]
for node_id in top_k:
human_string = label_lines[node_id]
score = predictions[node_id]
print('%s (score = %.5f)' % (human_string, score))
Run the following command in your terminal to test, for example, an image containing husky: (Here we use training data as input, but you should use others images for validating accuracy.)
Now, I get a score of 0.99933 for the test image. This indicates that the classifier is 99.9 percent sure that the image contains a husky. Congratulation! You now can train your own image classifier!