Image Recognition for Android with a Custom TensorFlow Lite Model

Deep learning on Android using Google Colab, transfer learning, and TFLite

Thanks to TensorFlow Lite (TFLite), we can build deep learning models that work on mobile devices. In fact, models generated by TFLite are optimized specifically for mobile and edge deployment for that purpose. After a deep learning model is created in TensorFlow, developers can use the TensorFlow Lite converter to convert that model to a format that runs in mobile devices.

In this tutorial, we’re going to build a TensorFlow model for recognizing images on Android using a custom dataset and a convolutional neural network (CNN).

Rather than creating a CNN from scratch, we’ll use a pre-trained model and perform transfer learning to customize this model with our new dataset. The pre-trained model we’re going to use is MobileNetV1, and we’ll train a version of it on the Fruits360 dataset.

After we’ve built our TensorFlow model with transfer learning, we’ll use the TFLite converter to create a mobile-ready model variant. The model will then be used in an Android application that recognizes images captured by the camera.

The project built out of this tutorial is available at GitHub. It is tested and running properly. You are advised to download and use it in case something is not working well.

This tutorial creates a Jupyter Notebook running in Google Colab for building the custom model.

For a look at getting started with Jupyter Notebooks in Google Colab, check out my previous tutorial, which covers some of the core functionality:

The sections covered in this tutorial are as follows:

  • Preparing MobileNet for transfer learning
  • Downloading the dataset
  • Transfer learning
  • Downloading the Android Studio Project
  • Customizing the Android Studio Project

Preparing MobileNet for Transfer Learning

TensorFlow allow us to download several pre-trained deep learning models using the tensorflow.keras.applications module. Within this module, there are several classes, each responsible for working with a model.

For example, there is a class named MobileNet for downloading MobileNet models. The code below downloads the MobileNet model and saves it inside the base_model variable. The constructor of the MobileNet class is given 2 arguments: input_shape and include_top.

MobileNet is trained using fixed image sizes. When a new image is fed to such a pre-trained model, this image must be of the same size as the images used for training the model. Because the size of the images in the Fruits360 dataset is (100, 100), we’re going to use the version of MobileNet that’s trained on images of sizes (128, 128). The reason is that (128, 128) is the nearest to (100, 100)—other sizes are (160, 160) and (224, 224).

The second argument is set to False because the top layers are customized to work with the dataset on which the MobileNet model is trained (ImageNet).

As a result, such layers work with 1000 classes. The version of the Fruits360 dataset we’re going to use just has 102 classes, and thus we must remove such layers and add other layers that work with our dataset.

Note that the trainable property of the model is set to False. This means the weights of the model will not change but just be used.

from __future__ import absolute_import, division, print_function, unicode_literals

try:
  # The %tensorflow_version magic only works in colab.
  %tensorflow_version 2.x
except Exception:
  pass
import tensorflow as tf 

IMG_SHAPE = (128, 128, 3)

base_model = tf.keras.applications.MobileNet(input_shape=IMG_SHAPE, include_top=False)
base_model.trainable = False

model = tf.keras.Sequential([
  base_model,
  tf.keras.layers.GlobalAveragePooling2D(),
  tf.keras.layers.Dense(102, activation='sigmoid')
])

model.summary()

model.compile(optimizer=tf.keras.optimizers.RMSprop(lr=0.0001),
              loss='binary_crossentropy',
              metrics=['accuracy'])

In order to add layers at the top of the model architecture that work with our dataset, the Sequential class is used. It accepts a list of layers that creates a new model.

The first element in the list is the base_model. This means all layers in MobileNet [except for its top layers] are included in the new model. In addition to such layers, 2 new layers are added, which are average pooling and dense layers. The dense layer has 102 neurons—because our dataset has 102 classes.

The new model returned by the Sequential class is saved in the model variable and then compiled using the compile() method. This method accepts the optimizer, loss function, and metrics. The only metric used is classification accuracy.

After compiling the model, the next section prepares the dataset images to be fed to that model for transfer learning. Transfer learning here means training the newly-added layers at the top of the architecture so that they can work with the new dataset.

Downloading the Dataset

The dataset we’re going to use is available at this link:

On this page, you’ll find a download link. Click it, and the download will start. You’ll want to cancel the download and copy the link from which the download started. After that, replace the link in the code below with the link you fetched for downloading the dataset.

Note that the download link expires after a few days—if it expires, you’ll need to fetch the link again.

import tensorflow as tf
import os

zip_file = tf.keras.utils.get_file(origin="https://storage.googleapis.com/kaggle-datasets/5857/414958/fruits-360_dataset.zip?GoogleAccessId=web-data@kaggle-161607.iam.gserviceaccount.com&Expires=1566314394&Signature=rqm4aS%2FIxgWrNSJ84ccMfQGpjJzZ7gZ9WmsKok6quQMyKin14vJyHBMSjqXSHR%2B6isrN2POYfwXqAgibIkeAy%2FcB2OIrxTsBtCmKUQKuzqdGMnXiLUiSKw0XgKUfvYOTC%2F0gdWMv%2B2xMLMjZQ3CYwUHNnPoDRPm9GopyyA6kZO%2B0UpwB59uwhADNiDNdVgD3GPMVleo4hPdOBVHpaWl%2F%2B%2BPDkOmQdlcH6b%2F983JHaktssmnCu8f0LVeQjzZY96d24O4H85x8wdZtmkHZCoFiIgCCMU%2BKMMBAbTL66QiUUB%2FW%2FpULPlpzN9sBBUR2yydB3CUwqLmSjAcwz3wQ%2FpIhzg%3D%3D", fname="fruits-360.zip", extract=True)
base_dir, _ = os.path.splitext(zip_file)
print(base_dir)

The dataset will be downloaded as a compressed file named fruits-360.zip, which is extracted automatically because the extract argument is set to True. The contents of the extracted file are as follows. The folders we’re going to use are Test and Training, which hold the test and training data, accordingly.

After downloading the dataset, next we need to feed its data to the MobileNet model for transfer learning, which we’ll cover in the next section.

Transfer Learning and Model Conversion

The code given below prepares the data generators that return batches of images for training the new model. In the end, the new model is saved with the name MobileNet_TransferLearning_Fruits360v48.h5.

import os

from __future__ import absolute_import, division, print_function, unicode_literals

try:
  # The %tensorflow_version magic only works in colab.
  %tensorflow_version 2.x
except Exception:
  pass
import tensorflow as tf

import numpy

zip_file = tf.keras.utils.get_file(origin="https://storage.googleapis.com/kaggle-datasets/5857/414958/fruits-360_dataset.zip?GoogleAccessId=web-data@kaggle-161607.iam.gserviceaccount.com&Expires=1568066045&Signature=apl0VAR2eb1vz8XOJAbxsf8KuSUmMqOCtoJ0HwDBh%2F2YkqtAnOtypYq5rpQF020%2ByqKGGTqbyVMZ1RC84CclDHxMratQlkHQP7Dh8W%2FBW5K4nxuNsInBG27kqIdx%2Ff%2Fere1WZOiqvUtXgf1uXNwbOjijqyIT77ypJFPhWt%2Bcg4cjjkbrFjumo1RAGbJ77u1A3nQu4VhxhH46tZ6xtEU0lrzSpMA3U8GkvHK0H65qTD3GzfZD%2FmjhDtD712mT3H5zPgvRFbbSenaGgE%2BdaqG5Z9EBK95FT19N2tSSOMQFGxhmG9VpiODHEn%2FIHDlyyH2E%2F7%2BpMDFMuB4n93NyFu%2FnMQ%3D%3D",
                                   fname="fruits-360.zip", extract=True)
base_dir, _ = os.path.splitext(zip_file)

train_dir = os.path.join(base_dir, 'Training')
validation_dir = os.path.join(base_dir, 'Test')

image_size = 128
batch_size = 32

train_datagen = tf.keras.preprocessing.image.ImageDataGenerator()

train_generator = train_datagen.flow_from_directory(directory=train_dir, target_size=(image_size, image_size), batch_size=batch_size)

validation_datagen = tf.keras.preprocessing.image.ImageDataGenerator()

validation_generator = validation_datagen.flow_from_directory(directory=validation_dir, target_size=(image_size, image_size), batch_size=batch_size)

IMG_SHAPE = (image_size, image_size, 3)

base_model = tf.keras.applications.MobileNet(input_shape=IMG_SHAPE, include_top=False)
base_model.trainable = False

model = tf.keras.Sequential([
  base_model,
  tf.keras.layers.GlobalAveragePooling2D(),
  tf.keras.layers.Dense(102, activation='softmax')
])

model.summary()

model.compile(optimizer=tf.keras.optimizers.Adam(),
              loss='categorical_crossentropy',
              metrics=['accuracy'])

epochs = 1
steps_per_epoch = numpy.ceil(train_generator.n / batch_size)
validation_steps = numpy.ceil(validation_generator.n / batch_size)

history = model.fit_generator(generator=train_generator,
                              steps_per_epoch = steps_per_epoch,
                              epochs=epochs,
                              validation_data=validation_generator,
                              validation_steps=validation_steps)

After customizing MobileNet for working with the Fruits360 dataset, the customized MobileNet remains a TensorFlow model—we still need to convert it to TensorFlow Lite in order to use it on Android. The code shown below uses the TFLiteConverter to convert the model to TFLite.

In order to build an Android app that classifies images according to the classes in the Fruits360 dataset, a text file must be available holding the class names. The next code snippet creates such a text file. After the model makes a prediction on an input image, the prediction score is used to find the class label of the image.

At this point, we’ve created our TFLite model and the class names are available. The complete code that downloads the dataset and the model, performs transfer learning on MobileNet, and saves the returned model in the TFLite format, is pasted below. To run this code in Colab, all you need to change is the dataset download link:

import os

from __future__ import absolute_import, division, print_function, unicode_literals

try:
  # The %tensorflow_version magic only works in colab.
  %tensorflow_version 2.x
except Exception:
  pass
import tensorflow as tf

import numpy

zip_file = tf.keras.utils.get_file(origin="https://storage.googleapis.com/kaggle-datasets/5857/414958/fruits-360_dataset.zip?GoogleAccessId=web-data@kaggle-161607.iam.gserviceaccount.com&Expires=1568066045&Signature=apl0VAR2eb1vz8XOJAbxsf8KuSUmMqOCtoJ0HwDBh%2F2YkqtAnOtypYq5rpQF020%2ByqKGGTqbyVMZ1RC84CclDHxMratQlkHQP7Dh8W%2FBW5K4nxuNsInBG27kqIdx%2Ff%2Fere1WZOiqvUtXgf1uXNwbOjijqyIT77ypJFPhWt%2Bcg4cjjkbrFjumo1RAGbJ77u1A3nQu4VhxhH46tZ6xtEU0lrzSpMA3U8GkvHK0H65qTD3GzfZD%2FmjhDtD712mT3H5zPgvRFbbSenaGgE%2BdaqG5Z9EBK95FT19N2tSSOMQFGxhmG9VpiODHEn%2FIHDlyyH2E%2F7%2BpMDFMuB4n93NyFu%2FnMQ%3D%3D",
                                   fname="fruits-360.zip", extract=True)
base_dir, _ = os.path.splitext(zip_file)

train_dir = os.path.join(base_dir, 'Training')
validation_dir = os.path.join(base_dir, 'Test')

image_size = 128
batch_size = 32

train_datagen = tf.keras.preprocessing.image.ImageDataGenerator()

train_generator = train_datagen.flow_from_directory(directory=train_dir, target_size=(image_size, image_size), batch_size=batch_size)

validation_datagen = tf.keras.preprocessing.image.ImageDataGenerator()

validation_generator = validation_datagen.flow_from_directory(directory=validation_dir, target_size=(image_size, image_size), batch_size=batch_size)

IMG_SHAPE = (image_size, image_size, 3)

base_model = tf.keras.applications.MobileNet(input_shape=IMG_SHAPE, include_top=False)
base_model.trainable = False

model = tf.keras.Sequential([
  base_model,
  tf.keras.layers.GlobalAveragePooling2D(),
  tf.keras.layers.Dense(102, activation='softmax')
])

model.summary()

model.compile(optimizer=tf.keras.optimizers.Adam(),
              loss='categorical_crossentropy',
              metrics=['accuracy'])

epochs = 25
steps_per_epoch = numpy.ceil(train_generator.n / batch_size)
validation_steps = numpy.ceil(validation_generator.n / batch_size)

history = model.fit_generator(generator=train_generator,
                              steps_per_epoch = steps_per_epoch,
                              epochs=epochs,
                              validation_data=validation_generator,
                              validation_steps=validation_steps)


saved_model_dir = '/content/TFLite'
tf.saved_model.save(model, saved_model_dir)

converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
tflite_model = converter.convert()

with open('model.tflite', 'wb') as f:
  f.write(tflite_model)

labels = 'n'.join(sorted(train_generator.class_indices.keys()))

with open('labels.txt', 'w') as f:
  f.write(labels)

After the above code runs, the TFlite model and the text file will be listed in the Jupyter notebook storage, according to the next image. You must retain to the file names (model.tflite and labels.txt). The reason for naming these files with such names will be clear in the next section.

Now, we have the data required for TensorFlow Lite to make predictions. The next step is to use these files in an Android Studio project. Inside the project, you’ll need to access the camera, capture images, feed the image to the model for making predictions, and finally show the result on the screen.

Rather than building the Android app from scratch, we can make use of an existing project that implements the above features. All we need to do is insert the TFLite model and the class labels in the right locations. Let’s start by downloading and testing the Android Studio project in the next section.

Downloading the Android Studio Project

The Android Studio project that uses MobileNet for image recognition can be downloaded from the set of examples available on TensorFlow’s examples repo on GitHub.

You can either clone this project, which includes more than the Android Studio project we’ll use. It has several projects that are ready to use by just importing them into Android Studio. To do that, just issue the following command:

After the project is cloned successfully, you can find several Android Studio projects that are ready to use in this directory: /examples/lite/examples. These projects are not just available for Android, but also for iOS and Raspberry Pi.

The project we need is available in this directory /examples/lite/examples/image_classification/android.

You can also download this project directory from this location:

After the download is complete, open Android Studio and import the project. Importing the project takes a bit of time wile the Grade sync completes. After the project is imported successfully, you can run it to see how things work.

After running the application, you can just place the camera in front of an object and the app will return a few class labels that are likely to match the object, in addition the classification probability score for each. The next figure gives an example. This is just for using the project as it is. But what about using the model we generated previously? We’ll implement in the next section.

Customizing our Android Studio Project

The app searches for 2 files names model.tflite and labels.txt in this directory:

/examples/lite/examples/image_classification/android/app/src/main/assets.

Go to this directory and paste the 2 files there. But this isn’t everything we need to do—there are some required changes in the project to allow the new model to be used.

When the TF model is converted into a TFLite model, the generated TFLite model isn’t optimized at all. The model is saved as a float model. This means the weights are saved as float numbers. In order to inform the app to work with the float model, there are 3 changes to make.

1. The first is to edit the strings.xml resource file. The file is available in this directory:

/examples/lite/examples/image_classification/android/app/src/main/res/values/strings.xml.

Just replace the <string-array> element by the following:

2. The second change is in the ClassifierFloatMobileNet.java file that’s available in this directory:

examples/lite/examples/image_classification/android/app/src/main/java/org/tensorflow/lite/examples/classification/tflite/ClassifierFloatMobileNet.java.

Change the return statement inside the getModelPath() method to the following statement:

3. The third change depends on the input image size. If you’re using the version of the MobileNet model that accepts images of size (224, 224), then you have nothing to do.

Otherwise, you’ll need to specify the size by editing the ClassifierFloatMobileNet.java file to return the X size in the getImageSizeX() method, and edit the getImageSizeY() method to specify the Y size of the image. In our case, we’re use an image of size (128, 128), and thus we must edit these 2 methods.

In case there is something wrong while applying these changes, please check the GitHub project associated with the tutorial: https://github.com/ahmedfgad/AndroidTFLite

After making these changes, the next step is to rerun the application. Open the app, and you’ll find that it predicts the class labels of the Fruits360 dataset. The next figure gives an example:

Conclusion

This tutorial discusses how to build a deep convolutional neural network that runs on Android devices using the TensorFlow Lite model format. We achieved this by transferring the learning of a pre-trained model (MobileNet) to a new dataset using TensorFlow.

The model returned from this transfer learning process was then converted into a TensorFlow Lite model, which was implemented in an Android Studio project to predict the class labels of new images.

Fritz

Our team has been at the forefront of Artificial Intelligence and Machine Learning research for more than 15 years and we're using our collective intelligence to help others learn, understand and grow using these new technologies in ethical and sustainable ways.

Comments 0 Responses

Leave a Reply

Your email address will not be published. Required fields are marked *

wix banner square