Face Mask Detection - Generate and Train - Part 2

TF Lite

If you are just after the code, you can find it on Github.

Overview

In the previous post, I address the theory and design choices around this project, here I am going to delve into the code and explain what is going on.

As a reminder, you can find the code on Github.

The data generation and model training phases have been implemented in a series of Jupyter notebooks that are intended to be run on Google Colab although you can move these to another environment if you wish.

In order to have some consistency and save time each time a notebook is run, I have created a constants.py which defines a number of folders and locations for other resources we will require.

The directory structure created is as follows:

- project_dir
    - dataset
        - unprepapared
            - images
            - annotations
            - mask
        - prepared
            - record
            - train
                - images
                - annotations
            - test
                - images
                - annotations
        - tmp1
        - tmp2

    - scripts
        - preprocessing

    - training
        - checkpoints
        - config
        - trained_graphs
            - inference_graph
            - tf_lite_graph
        - trained-models
            - tflite
            - tfjs`

Notebooks Part 1-3 - Data Preparation

I have attempted to make the notebooks as self explanatory as possible.

The first two notebooks address the process of generating training data for the model. Part 1 looks to annotate unmasked category of images with bounding boxes (in the form of XML files in PASCAL VOC format).

Once the images are annotated, those images and the annotations are stored in Google Drive. But before they are stored, we need to confirm that every image got an annotation. Error do occur and not every image gets an annotation therefore we want to make sure only image files with annotations are saved therefore we perform a check first.

In Part 2, we have exactly the same actions going on as Part 1 but in this case we have an added step of superimposing images of face masks onto the annotated face images. We will use these images in our masked category.

And in Part 3 we have some additional tidying up to do. This is also the stage in which you should import any additional training data that you use to use as part of your project. You can find such data here although you will have to amend the labels.

We need to split out dataset into training and testing partitions before we then consolidate the data from the individual XML files into a single CSV file.

There is a final data preparation task to complete however due to it requiring the TensorFlow libraries it was easier practically to include it on Part 4 which is dedicated to Training.

Notebook Part 4 - Training

Before you start the process, you need to determin which pretrained model you wish to make use of. By default you will employ the MobileNet v2 SSD model but you could also use MobileNet v1 SSD or FaceNet SSD by commenting out the appropriate sections in the constants.py file.

This is also a good time to look at the pipeline config file, in this case mobilenet_v2_pipeline.config which stores all the model training configuration.

Make sure the pretrained model and training data locations are correct. Depending on how many evaluation images you have in your dataset, you might want to amend the num_examples parameter under eval_config.

You will most likely have to experiment with the learning rate and try different values depending on how the training process goes. You might even want to opt for a different optimizer all together.

The delay parameter under quantization is the number of training steps after which the model drops its precision down to 8 bits. You might need a few training runs to determine the numbers of steps after which your model starts to converge and then rerun the training with that amount as your delay.

We start off the notebook by installing all the necessary libraries. We have to make sure that TensorFlow 1.x is installed not 2.x as that is not supported for Object Detection at the time of writing.

Once we have extracted the data from the previous steps, we need to generate the TF Record files. A TF Record is essentially a simple format for storing a sequence of binary records and speeds up the training process.

We can now start the training job which will be a trial and error, iterative process. Once the model has converged to an acceptable level you must stop the cell from running. Be sure to have made a note of which checkpoint you want to restore from before you stop the cell.

With the saved checkpoints to hand, you must export them into an TF Lite Graph . This step is necessary before we can arrive at the TF Lite model which is the next step.

If we wanted to convert our checkpoint into a TF JS Model then there is a different graph export that we must undertake.