Idiomatic TensorFlow on Android — Get started with the TensorFlow Support Library

Original article was published on Deep Learning on Medium

Idiomatic TensorFlow on Android — Get started with the TensorFlow Support Library

Working with data on Android

If you’ve used TensorFlow Lite on Android before, chances are that you’ve had to deal with the tedious task of pre-processing data, working with Float arrays in a statically typed language or resize, transform, normalize and do any of the other standard tasks required before the data is fit for consumption by the model.

Well, no more! The TFLite support library nightly is now available, and in this post, we’ll go over its usage and build a wrapper around a tflite model.

Note: A companion repository for this post is available here. Follow along, or jump straight into the source!

Scope of this post

This post is limited in scope to loading and creating a wrapper class around a tflite model; however, you can see a fully functional project in the repository linked above. The code is liberally commented and very straightforward.

If you still have any queries, please don’t hesitate to reach out to me and drop a comment. I’ll be glad to help you out.

Setting up the project

We’re going to be deploying a TFLite version of the popular YOLOv3 object detection model on an Android device. Without further ado, let’s jump into it.

Create a new project using Android studio, name it anything you like, and wait for the initial gradle sync to complete. Next, we’ll install the dependencies.

Adding dependencies

Add the following dependencies to your app-level build.gradle.

Let’s go through what they are, and why we need them.

  1. Quick Permissions: This is a great library to make granting permissions quick and easy.
  2. Tensorflow Lite: This is the core TFLite library.
  3. Tensorflow Lite Support: This is the support library we’ll be using to make our data-related tasks simpler.
  4. CameraView: This is a library that provides a simple API for accessing the camera.

Configuring the gradle project

Our project still needs a little more configuration before we’re ready for the code. In the app-level build.gradle file, add the following options under the android block.

The reason we need to add this is because we’ll be shipping the model inside our assets, which compresses it by default. This is an issue because compressed models cannot be loaded by the interpreter.

Note: After this initial configuration, run the gradle sync again to fetch all dependencies.

Jumping into the code

First things first; we need a model to load. The one I used can be found here. Place the model inside app/src/main/assets. This will enable us to load it at runtime.

The labels for detected objects can be found here. Place them in the same directory as the model.

Warning: If you plan to use your own custom models, a word of caution; the input and output shapes may not match the ones used in this project.

Creating a wrapper class

We’re going to wrap our model and its associated methods inside a class called YOLO. The initial code is as follows.

Let’s break this class down into its core functionality and behaviour.

  1. First, upon being created, the class loads the model from the app assets through the FileUtil class provided by the support library.
  2. Next, we have a class member. The interpreter is self explanatory, it’s an instance of a TFLite interpreter.
  3. Finally, we have some static variables. These are just the file names of the model and the labels inside our assets.

Moving on, let’s add a convenience method to load our labels from the assets.

Here we’ve declared a method that loads the label file, and lazily initialized a member var to the returned value.

Let’s get down to the brass tacks. We’re now going to define a method that takes in a bitmap, passes it into the model and returns the detected object classes.

Whoa, that’s a wall of code! Let’s go through it and break it down.

We’ve declared some new lazily initialized variables; an ImageProcessor and a TensorImage. These are classes exposed by the support library, to make loading images and processing them much simpler.

As is shown here, we can load a bitmap directly into the TensorImage and then pass it on to the ImageProcessor for further processing.

The ImageProcessor has several operations available, but the one we’ve used here is to resize our input images to 300 * 300. This is because our model’s input size requires a 300 * 300 image.

After processing the image, we create several TensorBuffers. These are representations of Tensors that we can manipulate and access easily. The shapes of these TensorBuffers is determined by the model. Take a look at the model summary to figure out the appropriate shapes.

We load the TensorImage into the input TensorBuffer, and then pass the input and output buffers into the interpreter.

Note: The YOLOv3 model has multiple outputs. This is the reason why we had to use multiple output buffers.

After running inference, the interpreter sets the internal FloatArrays of the output buffers. Right now, we’re only interested in the one that contains the predicted classes. Using the convenient kotlin map function, we map labels to the numerical classes output by the model and return them.

This class can now be used by our application to run inference on a bitmap. How convenient!

Conclusion

And that’s it! Compared to a project without using the support library; we’d have written much more code to resize the image, convert bitmaps to float arrays, allocate float arrays manually to store the output in, etc.

The TensorFlow support library thus makes life simpler for a developer; and despite being a nightly, it’s pretty stable in my experience.

To find out more, view the support library readme here. As of now, there aren’t any formal docs available, but the readme contains all the information a developer would need to get started quickly.

Stay safe, and have fun!