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.

Install TensorFlow

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()
>>> print(sess.run(hello))

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:

Hello, TensorFlow!

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:

Output from “Fatkun Batch Download Image” when search “husky” in 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:

python {$your-working_directory}/retrain.py
--bottleneck_dir=/{$your-working_directory}/bottlenecks         
--how_many_training_steps 500
--model_dir=/{$your-working_directory}/inception
--output_graph=/{$your-working_directory}/retrained_graph.pb
--output_labels=/{$your-working_directory}/retrained_labels.txt
--image_dir /{$your-working_directory}/${your_training_data_path}

Retrain command in Windows

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.

Creating bottlenecks

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.

Validating accuracy

Final test accuracy

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[1]
# 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
    in tf.gfile.GFile("D:/tensorflow_work/retrained_labels.txt")]
# Unpersists graph from file
with tf.gfile.FastGFile("D:/tensorflow_work/retrained_graph.pb", 'rb') as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
    _ = 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, 
     {'DecodeJpeg/contents:0': image_data})
     # Sort to show labels of first prediction in order of confidence
     top_k = predictions[0].argsort()[-len(predictions[0]):][::-1]
     for node_id in top_k:
         human_string = label_lines[node_id]
         score = predictions[0][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!

Result for test our classifier