Weapon detection model using YOLOv5 in pytorch
https://github.com/ari-dasci/OD-WeaponDetection
Has two classes: Knife and Pistol
Images were labelled and split into Test (75%), Validation (10%) and Training set (15%)
Model trained using the YOLOv5 model on Google Colab
YOLOv5 repository was cloned into the cloud machine and train.py was run with the following parameters:
img 256, epochs 100, batch-size 32
Class | Images | Labels | P | R | mAP@.5 | mAP@.5:.95: 100% |
---|---|---|---|---|---|---|
all | 697 | 799 | 0.92 | 0.839 | 0.91 | 0.618 |
knife | 697 | 320 | 0.917 | 0.884 | 0.936 | 0.589 |
pistol | 697 | 479 | 0.922 | 0.793 | 0.884 | 0.647 |
Deepstream was installed on Jetson Nano
Detailed steps for installation
The trained .pt model was converted to .wts and .cfg file using the following:
https://github.com/marcoslucianops/DeepStream-Yolo#basic-usage
The .wts and .cfg files are moved to Jetson Nano and the model was run