How to train your own custom model with Tensorflow object detection API and deploy it into Android with TF Lite

In these article I will explain the steps of training your own model with your own data set using Google Colab’s GPU and Tensorflow’s object detection API. After getting the model trained you will learn how to use Tensorflow Lite converter to get the Lite model and then get the model running on a simple Android app. You can also check the article on my Medium

So yeah, let us just start!!

First: Get the model

Before doing anything make sure that you have python 3.5+ installed on your local machine I installed it with Anaconda.

After getting your python. I can say there we go !! :D . Fork and clone this repo to your local machine and then install the required libraries:

pip3 install -r requirements.txt

Second: Prepare your own data set

To get some good high quality pictures I used Google’s Open Images data set, to download a specific class from this data set you can follow instructions in this Github repo, After cloning the repo to your local machine,don’t forget to change the download.py file like the snippets down if you are working on Windows, otherwise you won’t be able to download the images:

Change the following line from:

urllib.request.urlretrieve(url, os.path.join(OUTPUT_DIR, url.split(“/”)[-1]))

To:

urllib.request.urlretrieve(url.replace("\\", "/"), os.path.join(OUTPUT_DIR, url.split("\\")[-1]))

Now, we have two choices the first is to train the model with all of the already annotated images of the class that we have downloaded, the second one which I will complete with, is choosing the images and generate their xml files manually with LablelImg since I have a really specific case. Anyways to use LabelImg you can go and download it from here.

After downloading it just go into the folder and double click on the labelImg named file and you will get the tool’s UI working.

In order not to overwhelm the model, we need to resize our images before annotating them. For that you can run the following snippet, but first save your photos, ideally with jpg extension to ./data/raw directory.

import os
import glob
import cv2

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(
        description="Resize raw images to uniformed target size."
    )
    parser.add_argument(
        "--raw-dir",
        help="Directory path to raw images.",
        default="./data/raw",
        type=str,
    )
    parser.add_argument(
        "--save-dir",
        help="Directory path to save resized images.",
        default="./data/images",
        type=str,
    )
    parser.add_argument(
        "--ext", help="Raw image files extension to resize.", default="jpg", type=str
    )
    parser.add_argument(
        "--target-size",
        help="Target size to resize as a tuple of 2 integers.",
        default="(800, 600)",
        type=str,
    )
    args = parser.parse_args()

    raw_dir = args.raw_dir
    save_dir = args.save_dir
    ext = args.ext
    target_size = eval(args.target_size)
    msg = "--target-size must be a tuple of 2 integers"
    assert isinstance(target_size, tuple) and len(target_size) == 2, msg
    fnames = glob.glob(os.path.join(raw_dir, "*.{}".format(ext)))
    os.makedirs(save_dir, exist_ok=True)
    print(
        "{} files to resize from directory `{}` to target size:{}".format(
            len(fnames), raw_dir, target_size
        )
    )
    for i, fname in enumerate(fnames):
        print(".", end="", flush=True)
        img = cv2.imread(fname)
        img_small = cv2.resize(img, target_size)
        new_fname = "{}.{}".format(str(i), ext)
        small_fname = os.path.join(save_dir, new_fname)
        cv2.imwrite(small_fname, img_small)
    print(
        "\nDone resizing {} files.\nSaved to directory: `{}`".format(
            len(fnames), save_dir
        )
    )

Run the file above in your local machine with cmd or anaconda prompt using this line of code:

python resize_images.py --raw-dir ./data/raw --save-dir ./data/images --ext jpg --target-size "(800, 600)"

Resized images will locate in ./data/images/

Next, we will split those files into two directories, ./data/images/train and ./data/images/test. The model will only use images in the "train" directory for training and images in "test" directory serve as additional data to evaluate the performance of the model. I used the 20%-80% percentage with 500 images, that means 400 were used tor train and 100 were used to test.

Okay! now, let us just start annotating our resized images. Open LabelImg tool and make sure that you choose PascalVOC format. Open the images directory and change the save directory to either training or test.

Now, grab one of the images and use the Create RedBox symbol to create a bound box for the object that you want your model to detect and save it after writing the name of it. You can bound more than an object.

You can now see generated xml files inside ./data/images/train and ./data/images/test directories.

Third: Run the model with your data on Colab

Now that you have all of your data annotated, you are ready to push all of the files to your forked Github repo, after doing so, open this notebook in Google Colaboratory and change the repo url to your forked repo.

You can find a detailed explanation of the notebook here. To get better results you may change the number of training steps to 200000 and the number of evaluation steps to 400. Ideally you need to get the loss function close to the zero so figure it out.

Before running the cells make sure that you change the run time of the notebook to Python 3 and GPU you can do so by clicking the Runtime in the tool bar on the upper left of the notebook. Then, run all of teh celles one by one and check the results.

Fourth: Get the model’s lite version with TFLite

Well, things are getting more exciting right?

Let us add a couple of cells to the notebook:

First you need to export the SSD_Model to tflite

!python /content/models/research/object_detection/export_tflite_ssd_graph.py — pipeline_config_path {pipeline_fname} — trained_checkpoint_prefix “/content/models/research/training/model.ckpt-1000” — output_directory “/content/object_detection_demo/” — add_postprocessing_op True — max_detections 10

Next we’ll use TensorFlow Lite to get the optimized model by using TOCO, the TensorFlow Lite Optimizing Converter. This will convert the resulting frozen graph (tflite_graph.pb) to the TensorFlow Lite flatbuffer format (detect.tflite) via the following command. For a quantized model, run this from the tensorflow/ directory:

!toco \
 — graph_def_file=”/content/object_detection_demo/tflite_graph.pb” \
 — output_file=”/content/object_detection_demo/detect.tflite” \
 — input_shapes=1,300,300,3 \
 — input_arrays=normalized_input_image_tensor \
 — output_arrays=’TFLite_Detection_PostProcess’,’TFLite_Detection_PostProcess:1',’TFLite_Detection_PostProcess:2',’TFLite_Detection_PostProcess:3' \
 — inference_type=FLOAT \
 — allow_custom_ops

This command takes the input tensor normalized_input_image_tensor after resizing each camera image frame to 300x300 pixels. The outputs of the quantized model are named ‘TFLite_Detection_PostProcess’, ‘TFLite_Detection_PostProcess:1’, ‘TFLite_Detection_PostProcess:2’, and ‘TFLite_Detection_PostProcess:3’ and represent four arrays: detection_boxes, detection_classes, detection_scores, and num_detections. The documentation for other flags used in this command is here. If things ran successfully, you should now see a third file in the directory that you choose called detect.tflite. This file contains the graph and all model parameters and can be run via the TensorFlow Lite interpreter on the Android device.

Fifth: Let us run our model on Android

Yayyyy! you’re almost there, hold on :D

I assume that you already have Android Studio installed, if not you can just install it.

Clone the tflite repo to get the Android tflite project, open your android studio and and click on the open an existing project, then from the Open File or Project window that appears, navigate to and select the tensorflow-lite/examples/object_detection/android directory from wherever you cloned the TensorFlow Lite sample GitHub repo. Click OK.

Now copy the detect.tflite and labelmap.txt files under the assets package in the project

The labelmap.txt file should look like this:

???
banana
apple

Since I choosed the inference type to be float then real-numbers arrays will be of type float in the output file. If they were quantized in the input file, then they get dequantized. So I need to set the value of TF_OD_API_IS_QUANTIZED variable to false in the DetectorActivity.java check the snippet:

  // Configuration values for the prepackaged SSD model.
	  private static final int TF_OD_API_INPUT_SIZE = 300;
	  private static final boolean TF_OD_API_IS_QUANTIZED = false;
	  private static final String TF_OD_API_MODEL_FILE = "detect.tflite";
	  private static final String TF_OD_API_LABELS_FILE = "file:///android_asset/labelmap.txt";

	  private static final DetectorMode MODE = DetectorMode.TF_OD_API;

Also you need to edit the file TFLiteObjectDetectionAPIModel.java to :

-      recognitions.add(
-          new Recognition(
-              "" + i,
-              labels.get((int) outputClasses[0][i] + labelOffset),
-              outputScores[0][i],
-              detection));
+        final int classLabel = (int) outputClasses[0][i] + labelOffset;
+        if (inRange(classLabel, labels.size(), 0) && inRange(outputScores[0][i], 1, 0)) {
+            recognitions.add(
+                    new Recognition(
+                            "" + i,
+                            labels.get(classLabel),
+                            outputScores[0][i],
+                            detection));
+        }
    }
    Trace.endSection(); // "recognizeImage"
    return recognitions;
  }
+  private boolean inRange(float number, float max, float min) {
+    return number < max && number >= min;
+  }
+

And last but not least got your build.gradle and comment the following line

//apply from:’download_model.gradle’

Sync gradle and plug you phone usb and run the app, and voilà:

Conclusion

I have used many resources to get these results, I have added the links to the main resources, but if you get some errors in any step you may go and Google it :D. Nothing comes on a plate of gold :D..