A Custom Image Classification Android App using TensorFlow Lite

If you want to build your own Android app for image classification right from scratch, I recommend you read this article.

Today we’ll learn how to build an Android app for classifying images on custom data right from scratch. We’ll be building our own image classification model and deploy it on Android using TensorFlow Lite.

So, what are we waiting for? Let’s go.

Things we need before we begin?

1. Google Teachable Machine:

If you have no experience in training Machine Learning models, now you no need to worry. We will use Google Teachable Machine; an online, web-based tool to train ML models.

If you want to familiarize yourself with Teachable Machine, feel free to visit this link.

2. Android Studio:

To build an Android app and deploy our trained model on it, we’ll need Android Studio. If you have already installed Android Studio on your machine, you’re good to go. If not, you can download it from this link.

Once you’ve set up everything, let’s move ahead.

Step 1: Create a dataset.

I’ll be creating an animal classifier app by collecting images of different animals and training a model. Feel free to choose your own custom dataset. The overall process remains the same.

To give an idea about how my dataset looks like, here are some images I have collected.

The dataset consists of Six classes – Elephant, Dog, Cow, Chicken, Cat, and Butterfly.

Step 2: Train the classification model.

Alright, we’re already good to start training our model.

Open Teachable Machine on your web browser. Since we’re building an image classification project, we need to choose the Image Project option.

Now, you’ll be prompted to edit class names and import their respective images. The image below will help you to get going with it. You can either upload images from your local machine or import them from Google Drive. I’ll be importing the dataset from my local machine.

I have imported about 60 images for each class. Remember,Garbage in, garbage out”If you want to achieve good accuracy, make sure you provide enough images for model training.

Note: Teachable Machine is a web-based application. Every image imported for training is an HTTP request sent to it to which it responds. Flooding it with too many images can cause it to become unresponsive.

Once you’ve imported all the images, you may go ahead with the configuration of training hyperparameters.


It is the number of times the learning algorithm will work through the entire training dataset. With every next incoming epoch, the learning algorithm tries to learn better from its previous epoch and increase the model’s accuracy. By default, the no. of epochs is set to 50. This means that the learning algorithm will run for 50 passes before training is terminated.

Batch Size:

It refers to the no. of training samples passed to the algorithm during training for every pass. The default batch size is 16. This means that 16 training images will be passed to the algorithm at once for it to look at it and learn what’s inside those images.

For example, for the model to identify Doraemon in an image, the features it should ideally learn are Doraemon’s blue color and his round shape. Similarly, it learns features of other classes as well.

Learning Rate:

It is a hyperparameter that determines the step size at each iteration while minimizing the loss function. Lower the loss better will be the model accuracy. In simple words, it is the rate at which the model ‘learns’ the features of the images passed to it. If the learning rate is too high, then the model accuracy may decrease as it may fail to converge to the minima and if it is too low, it may take a lot of time to train the model. So, choose the learning rate wisely. Here, by default, the learning rate is set to 0.001.

I’m keeping all the hyperparameters’ default values untouched as of now. Feel free to play around with different values and see the results for yourself.

Once you’ve configured all the hyperparameters, click on Train Model to start training the model.

Depending upon the no. of images in your dataset and the no. of epochs you’ve set, the time to train the model would vary.

Step 3: Test the model.

Once the model training is finished, test the model with some images. Turn the Input toggle button ON and choose either the Webcam or File option to import the test image. I’m importing the test image as a file from the local machine.

And well, here are the results. 100% Cat. If you’re satisfied with the results obtained, you’re good to go to the next step. If still not satisfied, reconfigure the hyperparameters by changing epochs, batch size, and learning rate and possibly add more training images and retrain the model.

Step 4: Export model.

Next, if you feel your model is giving out expected results, you can export the model.

Click on the Export Model option. In the pop-up that appears, choose the TensorFlow Lite tab. After that, select the Quantized option.

But, what is quantization, and why do we need to do it?

Since mobile phones have limited hardware specifications compared to a computer, running Deep Learning models on them can cause them to crash down.

To deal with this problem, models are quantized post-training. It includes general techniques that reduce CPU and hardware accelerator latency, processing, power, and model size with little degradation in model accuracy. These techniques can be performed on an already-trained float TensorFlow model and applied during TensorFlow Lite conversion.

Once the model is quantized and converted to TensorFlow Lite, a file named converted_tflite_quantized.zip will be downloaded automatically. The downloaded ZIP contains two files – models.tflite and labels.txt.

Now, we can deploy this TF Lite model on our Android app.

But wait, have we created an app yet for the TF Lite model to be deployed? Find it out in the next step.

Step 5: Creating an Android app.

I have built a skeleton app already. You can find it on this link.

Now, navigate to Cartoon-Classification-on-Android-using-TF-Lite/app/src/main/assets directory and copy models.tflite and labels.txt files here.

After copying these two files, open Android Studio, open this repository as a new project, and build itself for some time. 

Alright, once the project is built successfully, open the MainActivity.java file and all you have to do is make the following changes.

Edit Line 169 by replacing cartoon_model.tflite with the name of your TF Lite model (model.tflite).

Edit Line 231 by replacing cartoon_labels.txt with the name of your model's label file. (labels.txt).

After doing the aforementioned changes, build the project again.

Next, simply connect your mobile phone to your machine and install the application on it and enjoy using this image classification app of yours.

Congratulations on building your own image classification Android app. Feel free to update the source app by adding more exciting features and do share your works in the comment sections. I’d love to have a look at them.

Post a Comment