This program trains a neural network to detect the use of facemasks in images. 850 images are used for testing and training with the Tensorflow Python package to train a neural network on a python notebook---which can be easily run on any machine.
Using the Jupyter Runtime hosted on Google Colab, this project can generate labels on face images such as:
Since all training was completed in a Jupyter Notebook, it's very easy to follow along with the steps used in this project and recreate an effective model! The training data and annotations are already included in this github directory.
All that's needed is to:
- Download the .ipynb file from this directory
- Open it in Google Colab
- Connect to a GPU under Edit->Notebook settings->Hardware accelerator->GPU
- Run the cells in order!
By the end, when training is complete, the project exports the trained model to the associated directory, and allows you to test it with images.
So, how does this project work? For this approach, I used the Python library Tensorflow, which streamlines the training of neural networks. Using face mask pictures is a great choice for object detection because their current relevance means there is an abundance of good data to obtain online.
This project takes a spin on the concept by introducing a third class: improperly worn face masks.
To start, the Python imports all required machine learning packages and downloads the required training data from this repository. This project uses the SSD ResNet50 V1 FPN 640x640 model---training this model requires .tfrecord files and a label map file.
Label Map:
The label_map.pbtxt works to describe what classes are used in the labeled objects. Alongside this, .tfrecord files are binary files used by Tensorflow to expedite the training process. These typically must be generated by the user, by loading all the data and writing it into a .record file; this is accomplished in the notebook.
After that, model_main_tf2.py should be run to train the model alongside this information. This can be pulled from the object detection directory from the official Tensorflow model directory information, which is downloaded at the beginning.
At the end, the user can load in their own images through Colab to test the efficiency of the model.
Cheers!
Max