Image Classification pipeline with Intel Neural Compute Stick 2 (NCS2)


This post was originally published by Oleksii Sheremet at Towards Data Science


What we’ll cover in this story:

  • Installing OpenVINO toolkit for Ubuntu.
  • Data preprocessing and model training with Google Colab.
  • Saving a TensorFlow model to protocol buffer (pb) format.
  • Converting a TensorFlow model to intermediate representation.
  • Running a trained model and prediction obtaining with NCS2 device.

In the past few years, there has been an active introduction of machine learning for industrial tasks solvings. Machine learning allows you to solve problems that humans cannot cope with. This can be work in hard-to-reach places, in hazardous chemical production, in conditions of increased radiation, etc. Machine learning is applicable in areas where human intelligence can also be applied, but is ineffective: predicting critical failures, preventing sudden equipment failures, condition-based maintenance, predicting the remaining life of equipment.

Often, in industrial enterprises there is no opportunity to use workstations for data analysis and processing. Therefore, special devices that can be easily connected to existing equipment (most often an industrial microcomputer) are in demand. The computational module NCS2 is one of such devices.

Our company AI Labs periodically receives requests for solving some industrial problems related to machine learning. Not so long ago, we had the opportunity to take advantage of the opportunities provided by NCS2. As a result, I developed a pipeline for image classification using NCS2. I would like to share the results of this work with the readers.

Installing OpenVINO toolkit for Ubuntu

First of all, you need to download and install the current version of OpenVINO. You can get OpenVINO from here. You can find a detailed installation guide here. Quick installation guide:

  1. Go to the folder (for example, Downloads) with the downloaded archive:
cd ~/Downloads/

2. Unpack the .tgz file (where <version> is the version of the downloaded archive, you will see it in the file name):

tar -xvzf l_openvino_toolkit_p_<version>.tgz

3. Go to the l_openvino_toolkit_p_<version> directory:

cd l_openvino_toolkit_p_<version>

4. Run the installation script as root:

sudo ./

5. Follow the instructions on your screen.

6. Set environment variables by editing the .bashrc file. Go to the user directory and open .bashrc file for editing (for example, with nano):

cd ~
nano .bashrc

Add the following line to the end of the .bashrc file, save it and reload:

source /opt/intel/openvino/bin/
source .bashrc

That’s all. All the necessary functionality for local work with NCS2 has been prepared. You can perform data preprocessing and model training.

Data preprocessing and model training with Google Colab

Animals-10 dataset was used to solve the classification problem. To simplify the problem, 5 classes (species of animals) were selected from Animals-10: cat, dog, chicken, horse, sheep. Images related to each of the classes are grouped into corresponding folders (‘cat’, ‘dog’, ‘chicken’, ‘horse’, ‘sheep’). The model was trained in Google Colab. Let’s take a closer look.

First of all, you need to install a version of TensorFlow suitable for working with NCS2:

!pip install tensorflow==1.15.2

Import libraries:

import tensorflow.compat.v1 as tf
import tensorflow_hub as hub
import matplotlib.pylab as plt
import numpy as np

Mount the Google Drive (with dataset) to Google Colab:

from google.colab import drive

Go to the directory where the data was uploaded (‘animals’):

%cd ‘drive/My Drive/animals’

The folder structure should be as follows:

Image is created by Oleksii Sheremet with tree command

Create a list of classes (labels):

image_dir = 'animals'
import osclass_labels = []
for x in tf.gfile.Walk(image_dir):
y = os.path.basename(x[0])
if y != 'animals':

We get the following list:

[‘cat’, ‘chicken’, ‘dog’, ‘horse’, ‘sheep’]

Based on this list, index ‘0’ corresponds to the image of a cat, ‘1’ — a chicken, ‘2’ — a dog, ‘3’ — a horse, ‘4’ — a sheep.

Set the size of the images (224 x 224). We use ImageDataGenerator to create a flow of images from folders when training the model.

IMAGE_SIZE = (224,224)image_generator = tf.keras.preprocessing.image.ImageDataGenerator(rescale=1/255, validation_split=0.2)training_set = image_generator.flow_from_directory(str(image_dir),                                              target_size=IMAGE_SIZE, subset='training')validation_set = image_generator.flow_from_directory(str(image_dir),                                              target_size=IMAGE_SIZE, subset='validation')

Retrieving pre-trained MobileNet network:

feature_extractor = tf.keras.applications.MobileNet(weights='imagenet', 

Build model:

del model

Compile the model:


Callback method for collecting logs:

class CollectBatchStats(tf.keras.callbacks.Callback):
def __init__(self):
self.batch_losses = []
self.batch_acc = []
self.validation_losses = []
self.validation_acc = []def on_train_batch_end(self, batch, logs=None):

Train the model:

steps_per_epoch = np.ceil(training_set.samples/training_set.batch_size)batch_stats_callback = CollectBatchStats()history = model.fit_generator(training_set, epochs=5,
callbacks = [batch_stats_callback])

Check predictions:

for image_batch, label_batch in validation_set:
print("Image batch shape: ", image_batch.shape)
print("Label batch shape: ", label_batch.shape)

We have the following shapes:

Image batch shape:  (32, 224, 224, 3)
Label batch shape:  (32, 5)

We can get predictions and their classes:

predictions = model.predict(image_batch)
predicted_class = np.argmax(predictions, axis=-1)

Visualize the result:

for n in range(30):
plt.title(f'pred: {class_labels[predicted_class[n]]}norig: {class_labels[np.array(label_batch[n]).argmax()]}')
_ = plt.suptitle("Animals")

Image is created by Oleksii Sheremet with matplotlib module

Saving a TensorFlow model to protocol buffer (pb) format

Freeze and save model:

from tensorflow.compat.v1.keras import backend as Ksession = K.get_session()
graph = session.graph
with graph.as_default():
freeze_var_names = list(set( for v in tf.global_variables()).difference([]))
output_names = [ for out in model.outputs]
output_names += [ for v in tf.global_variables()]
input_graph_def = graph.as_graph_def()
for node in input_graph_def.node:
node.device = ""
frozen_graph = tf.graph_util.convert_variables_to_constants(session, input_graph_def, output_names, freeze_var_names)
tf.train.write_graph(frozen_graph, "animals", "animals.pb", as_text=False)

As a result, the trained model is saved to the animals.pb file.

Converting a TensorFlow model to intermediate representation

There is a detailed tutorial on converting TensorFlow models. To convert the trained model, you need to perform the operations described below.

Go to the model_optimizer folder (you need

cd ~/intel/openvino/deployment_tools/model_optimizer

Run with following command (please note that you must use a scaling factor of 255, the same as in the image preprocessing stage):

python --input_model <path_to_model>/animals.pb --batch 1 --scale 255

If the conversion is successful, you will see the following messages in the console:

[ SUCCESS ] Generated IR version 10 model.
[ SUCCESS ] XML file: /home/user_name/intel/openvino_2020.4.287/deployment_tools/model_optimizer/./animals.xml
[ SUCCESS ] BIN file: /home/user_name/intel/openvino_2020.4.287/deployment_tools/model_optimizer/./animals.bin
[ SUCCESS ] Total execution time: 31.81 seconds.
[ SUCCESS ] Memory consumed: 370 MB.

Thus, as a result of the conversion, we get 3 files: animals.xml and animals.bin, animals.mapping. It is advisable to copy these files to a separate folder (for example, test_folder).

Running a trained model and prediction obtaining with NCS2 device

To check the performance of the model on NCS2, you can use the script, which is included in OpenVINO (/home/user_name/intel/openvino/deployment_tools/inference_engine/samples/python/classification_sample/ Copy this script to the folder with 3 trained model files (animals.xml, animals.bin, animals.mapping). Also, to check the functionality, you can copy several images (for example, 001.jpeg and 002.jpeg) to the same folder.

Connect NCS2 to USB, go to the test_folder and run (MYRIAD key to run the model on NCS2):

python -m animals.xml -i 001.jpeg -d MYRIAD

Test image and output from NCS2 in console:

Image is created by Oleksii Sheremet with Adobe Photoshop

python -m animals.xml -i 002.jpeg -d MYRIAD

Image is created by Oleksii Sheremet with Adobe Photoshop

As you can see from the examples presented, the classification is correct. Image 001.jpeg shows a horse (classid=3), and image 002.jpeg shows a chicken (classid=1).


The presented pipeline allows you to quickly go through all the stages from data preprocessing and model training to its testing based on the NCS2 module. As practice has shown, NCS2 copes well with the task of image classification and can be recommended for solving simple industrial classification tasks.


Spread the word

This post was originally published by Oleksii Sheremet at Towards Data Science

Related posts