This project presents a mask detector using TensorRT YOLOv4 on a Jetson Nano Developer Kit B01 4GB. The dataset used for YOLOv4 training is the Mask Dataset from MakeML. Cause the dataset is in Pascal VOC format, I have used the xml2yolo repository to convert the labels to Darknet format. YOLOv4 has been trained using a Google Colab Notebook based on the YOLOv4-Cloud-Tutorial repository. The conversion of the Darknet model from YOLOv4 to TensorRT has been carried out with the tensorrt_demos repository. The following sections explain in detail the steps to be followed to use the mask detector.
Run the following command to make sure you are in your home directory:
cd ~/Create a folder called project:
mkdir project
cd projectTo train the model I used the dataset Mask Dataset from MakeML. To download it run the following command:
wget https://arcraftimages.s3-accelerate.amazonaws.com/Datasets/Mask/MaskPascalVOC.zipCreate a folder called dataset to store all images and label files:
mkdir dataset
#Install the zip and unzip libraries if you don't have them
sudo apt install zip unzip
unzip MaskPascalVOC -d dataset/
rm MaskPascalVOC.zipInside the dataset folder there are two folders. One with the annotations and the other with the raw images.
├── annotations
│ ├── maksssksksss0.xml
│ ├── maksssksksss1.xml
│ ├── maksssksksss2.xml
│ ├── maksssksksss3.xml
│ └── ...
└── images
├── maksssksksss0.png
├── maksssksksss1.png
├── maksssksksss2.png
├── maksssksksss3.png
└── ...Clone the following repository:
cd ~/project
git clone https://github.com/jmudy/xml2yolo.git
cd xml2yoloCopy the label files to this directory and run the script convert.py to convert the labels from Pascal VOC format to Darknet format.
cp ../dataset/annotations/*.xml .
python3 convert.py
mv *.txt ../dataset/annotations/
rm ../dataset/annotations/*.xmlThe images for training and test have been divided in a ratio of 80-20% respectively.
Clone the following repository to download the files to be used in the training with Google Colab.
cd ~/project
git clone https://github.com/jmudy/mask-detector
cp -r mask-detector/yolov4-mask/ .
rm -r -f mask-detector/Create obj and test folders to store the training and test images and labels, with a ratio of 80% training - 20% test.
cd ~/project/dataset
mkdir obj
mkdir test
cd images
cp $(ls -v | head -n 682) ../obj
cp $(ls -v | tail -n 171) ../test
cd ../annotations
cp $(ls -v | head -n 682) ../obj
cp $(ls -v | tail -n 171) ../test
cd ../Shape of the dataset folder:
├── images
├── annotations
├── obj
│ ├── maksssksksss0.png
│ ├── maksssksksss0.txt
│ ├── maksssksksss1.png
│ ├── maksssksksss1.txt
│ ├── maksssksksss2.png
│ ├── maksssksksss2.txt
│ └── ...
└── test
├── maksssksksss682.png
├── maksssksksss682.txt
├── maksssksksss683.png
├── maksssksksss683.txt
├── maksssksksss684.png
├── maksssksksss684.txt
└── ...Compress obj and test folders and save in yolov4-mask folder.
zip -r ../yolov4-mask/obj.zip obj/
zip -r ../yolov4-mask/test.zip test/The following Google Colab Notebook can be used to train the model with the custom dataset (NOTE: Copy the yolov4-mask folder to the root of your Google Drive folder before using the Notebook)
https://colab.research.google.com/drive/1MriQiq8z7lxsDWkibTULqymypdeas_d-?usp=sharing
At the end of training rename the file /mydrive/yolov4-mask/backup/yolov4-mask_best.weights to yolov4-mask.weights.
Clone the following repository:
cd ~/project
git clone https://github.com/jkjung-avt/tensorrt_demos.git
cd tensorrt_demosBuild and install the following dependencies:
cd ssd
bash install_pycuda.sh
cd ../
wget https://raw.githubusercontent.com/jkjung-avt/jetson_nano/master/install_protobuf-3.8.0.sh
bash install_protobuf-3.8.0.sh
sudo pip3 install onnx==1.4.1Compile with make:
cd plugins
makeCopy in the ~/project/tensorrt_demos/yolo/ folder the yolov4-mask.cfg file you have used and the yolov4-mask.weights file that has been created in the training with Google Colab and convert the Darknet model to ONNX model and then to TensorRT engine.
cd ~/project/tensorrt_demos/yolo
python3 yolo_to_onnx.py -m yolov4-mask
python3 onnx_to_tensorrt.py -m yolov4-maskChange the COCO_CLASSES_LIST in the yolo_classes.py file located in the ~/project/tensorrt_demos/utils/ folder with the classes that have been trained:
"""yolo_classes.py
NOTE: Number of YOLO COCO output classes differs from SSD COCO models.
"""
COCO_CLASSES_LIST = [
'with_mask',
'without_mask',
'mask_weared_incorrect',
]
# For translating YOLO class ids (0~79) to SSD class ids (0~90)
yolo_cls_to_ssd = [
1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20,
21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40,
41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58,
59, 60, 61, 62, 63, 64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79,
80, 81, 82, 84, 85, 86, 87, 88, 89, 90,
]
def get_cls_dict(category_num):
"""Get the class ID to name translation dictionary."""
if category_num == 3:
return {i: n for i, n in enumerate(COCO_CLASSES_LIST)}
else:
return {i: 'CLS%d' % i for i in range(category_num)}Change in the trt_yolo.py file the default value of the classes that are detected and increase the confidence threshold:
cd ~/project/tensorrt_demos
sed -i '33s/default=80/default=3/' trt_yolo.py
sed -i '101s/conf_th=0.3/conf_th=0.8/' trt_yolo.pyRun the following command to display the results:
cd ~/project/tensorrt_demos
python3 trt_yolo.py --usb 0 --model yolov4-maskTo view the demo please click on the following YouTube link:
https://www.youtube.com/watch?v=YohIfmsn2Jg
This project is totally inspired by the following previous repositories: