Image Classification Pipeline with Intel Neural Compute Stick 2 (NCS2) | by Oleksii Sheremet | Towards Data Science

Introduction
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:
Go to the folder (for example, Downloads) with the downloaded archive:
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 ./install.sh
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/setupvars.shsource .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 tfimport tensorflow_hub as hubimport matplotlib.pylab as pltimport numpy as np
Mount the Google Drive (with dataset) to Google Colab:
from google.colab import drivedrive.mount(‘/content/drive’)
Go to the directory where the data was uploaded (‘animals’):
%cd ‘drive/My Drive/animals’
The folder structure should be as follows:
Create a list of classes (labels):
image_dir = 'animals'import osclass_labels = []for x in tf.gfile.Walk(image_dir):    try:        y = os.path.basename(x[0])        if y != 'animals':            class_labels.append(y)    except:        passprint(class_labels)
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',                                 include_top=False,                                input_shape=(IMAGE_SIZE+(3,))) feature_extractor.trainable=False
Build model:
try:    del modelexcept:    passx=feature_extractor.outputx=tf.keras.layers.GlobalAveragePooling2D()(x)classifier=tf.keras.layers.Dense(label_batch.shape[1],activation='softmax')(x)model=tf.keras.Model(inputs=feature_extractor.input,outputs=classifier)model.build((None,)+IMAGE_SIZE+(3,))model.summary()
Compile the model:
model.compile(    optimizer=tf.keras.optimizers.Adam(),    loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True),    metrics=['acc'])
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):    self.batch_losses.append(logs['loss'])    self.batch_acc.append(logs['acc'])    try:        self.validation_losses.append(logs['val_loss'])        self.validation_acc.append(logs['val_acc'])    except:        self.validation_losses.append(None)        self.validation_acc.append(None)    self.model.reset_metrics()
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,                              steps_per_epoch=steps_per_epoch,                              validation_data=validation_set,                              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)    break
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:
plt.figure(figsize=(12,10))plt.subplots_adjust(hspace=0.5)for n in range(30):    plt.subplot(6,5,n+1)    plt.imshow(image_batch[n])    plt.title(f'pred: {class_labels[predicted_class[n]]}\norig: {class_labels[np.array(label_batch[n]).argmax()]}')    plt.axis('off')_ = plt.suptitle("Animals")
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.graphwith graph.as_default():    freeze_var_names = list(set(v.op.name for v in tf.global_variables()).difference([]))    output_names = [out.op.name for out in model.outputs]    output_names += [v.op.name 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 mo_tf.py):
cd ~/intel/openvino/deployment_tools/model_optimizer
Run mo_tf.py with following command (please note that you must use a scaling factor of 255, the same as in the image preprocessing stage):
python mo_tf.py --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 classification_sample.py script, which is included in OpenVINO (/home/user_name/intel/openvino/deployment_tools/inference_engine/samples/python/classification_sample/classification_sample.py). 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 classification_sample.py -m animals.xml -i 001.jpeg -d MYRIAD
Test image and output from NCS2 in console:
python classification_sample.py -m animals.xml -i 002.jpeg -d MYRIAD
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).
Conclusion
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.
References
Intel® Neural Compute Stick 2
OpenVINO toolkit
MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications