Skip to content

tyui592/Real_Time_Helmet_Detection

Repository files navigation

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

Open RTHD in Colab

Requirements

Features

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
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
  1. Build App

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

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