Real Time Helmet Detection

This is pytorch implementation of helmet detector based on CenterNet.

I used the SafetyHelmetWearing-Dataset(SHWD) to detect helmet or person.

I will continue to update the entries to be filled in TODO for research in nightly

Requirements

• imgaug (v0.4.0)
• torch (v1.6.0)
• torchvision (v0.7.0)
• torchsummary
• requirements.txt

Features

• Automatic Mixed Precision(AMP)
• Distributed Data Parallel(DDP)
• TorchScript

Usage

SafetyHelmetWearing-Dataset(SHWD) is need to train detector [Download].

The trained model weights and demo app can be used in release(v0.0).

Example Scripts

Train

$ python main.py --train-flag --gpu-no 0 --data ./DATA/VOC2028/ --save-path ./WEIGHTS/ --amp

Test

$ python main.py --gpu-no 0  --model-load ./WEIGHTS/check_point.pth --data ./DATA/VOC2028 --imsize 512 --save-path ./WEIGHTS/results --batch-size 8

Measure mAP

$ cd mAP
mAP$ python main.py -na -np --dr ../WEIGHTS/results/txt/

Demo

$ python evaluate.py --gpu-no 0 --model-load ./WEIGHTS/check_point.pth --data ./imgs/000019.jpg --imsize 512 --save-path ./imgs/000019 --topk 100 --conf-th 0.2 --nms-th 0.2 --fontsize 0

Results

Input	Output

Performance

Helmet (AP)	Person (AP)	mAP
88.16 %	88.71 %	88.43 %

The model trained by above Example Script. The performance may be improved by increasing model size (ex, --num-stack, --increase-ch, ...) or searching hyperparameters (ex, --hm-weight, --lr, ...) more carefully.

Loading A Pytorch Model in C++

1. Create TorchScript Code of Detector

$ python export.py --model-load ./WEIGHTS/check_point.pth --nms-th 0.5 --topk 100

2. Build App

3. Run App (Speed: 100FPS @(512x512, 1080 Ti))

PytorchToCpp/build$ ./main -m ../../jit_traced_model_gpu.pth -i ../../imgs/000019.jpg