x光下危险物品/违禁物品目标识别的模型训练与推理代码

本文链接：https://blog.csdn.net/matt45m/article/details/132129122

前言

1.安检在公共场合的重要性不言而喻，保障群众人身安全是其首要任务。在各种场合，安检都是不可或缺的环节。x光安检机作为安检的重要工具，尽管其具有人工监控判断成像的特性，但是其局限性也十分明显。
为了解决这一局限性为出发点，让人工智能介入，这里在Torch框架下使用YOLO算法实现了对x光图像中刀具、气罐、鞭炮等10类危险物品的目标检测功能。
2.源码地址：https://download.csdn.net/download/matt45m/88178088

数据集

1.数据集定义

每个场景的安检违禁品都是一样的，但从法律法规的角度来度，违禁品可以分10大类，60多种小类，具体如下：

枪支类：运动枪、机枪、防暴枪、电击枪、发令枪、麻醉注射枪、气枪、步枪、冲锋枪、射钉枪、消防灭火枪、猎枪、玩具枪、仿真枪、催泪枪、道具枪、钢珠枪、手枪
五金工具类：大扳手、铁锤子、铲子、警棍、斧子、臂力器、甩棍
宠物类：猫咪、仓鼠、水獭、蛇、刺猬、龙猫、兔子、狗狗（导盲犬除外）、乌龟、蜥蜴
消毒用品类：医用酒精、84消毒液、酒精免洗凝胶、酒精喷雾、过氧乙酸消毒液、双氧水消毒液、酒精消毒液
毒品、腐蚀物：盐酸、硒粉、农药、水银、氰化物、砒霜、苯酚、硫酸、老鼠药、硝酸、氢氧化钾、杀虫药、氢氧化钠
燃料、气体：煤油、乙烷、乙烯、天然气、甲烷、柴油、丙烯、乙炔、液化石油气、丁烷、汽油、氢气、一氧化碳
易燃物：红磷、固体酒精、白磷、赛璐珞、镁铝粉、钾、锂、钠、油漆、稀料、油纸、乙醚、黄磷、松香油、碳化钙（电石）、丙酮、苯、闪光粉
食品类：未密封的榴莲、自嗨锅、活鱼、活虾、自热米饭、活螃蟹、臭豆腐、自热火锅
爆炸物品类：燃烧弹、毒气弹、引火线、导火索、雷管、催泪弹、炸药、炸弹、照明弹、手雷、信号弹、各种烟花爆竹、手榴弹、烟幕弹、烟火药。
刀具用品类：美工刀、屠宰刀、尖头剪刀、陶瓷刀、水果刀、餐刀、瑞士军刀、菜刀。

2.数据集收集

X光的数据有两种形式，一种彩色的成像，一种是黑白成像，相对于彩色成像，黑白成像对于人类视觉，更难于分出违禁品。

在这里插入图片描述

3.数据标注

数据集使用了网上开源的数据和私人收集的部分数据集，数据集涵盖了基本上所有的违禁品，但由于人手与算力了的问题，只挑了4000张彩色成像的x光图像，标注了10个我国常见的违禁品，分别是有：‘lighter’,‘scissors’,‘powerbank’,‘pressure’,‘knife’,‘zippooil’,‘handcuffs’,‘slingshot’,‘firecrackers’,‘nailpolish’。数据标注LabelImg,标注格式是YOLO。
在这里插入图片描述

模型训练

训练和开发环境是win10,显卡RTX3080;cuda10.2,cudnn7.1;OpenCV4.5;yolov5用的是5s的模型，Anaconda 3.5。

1.创建环境

 conda create --name yolov5 python=3.8
 activate yolov5
 git clone https://github.com/ultralytics/yolov5.git
 cd yolov5
 pip install -r requirements.txt

或者

conda create --name yolov5 python=3.8
activate yolov5
git clone https://github.com/ultralytics/yolov5.git
cd yolov5
conda install pytorch torchvision cudatoolkit=10.2 -c pytorch
pip install cython matplotlib tqdm opencv-python tensorboard scipy pillow onnx pyyaml pandas seaborn

2.训练

打开model/yolov5s.yaml,更改nc数目。

# parameters
nc: 10  # 检测总类别
depth_multiple: 0.33  # model depth multiple 网络的深度系数
width_multiple: 0.50  # layer channel multiple 卷积核的系数

# anchors 候选框，可以改成自己目标的尺寸，也可以增加候选框
anchors:
  - [10,13, 16,30, 33,23]  # P3/8
  - [30,61, 62,45, 59,119]  # P4/16
  - [116,90, 156,198, 373,326]  # P5/32

# YOLOv5 backbone
backbone: #特征提取模块
  # [from, number, module, args]
  # from - 输入是什么，-1：上一层的输出结果;
  # number - 该层的重复的次数，要乘以系数，小于1则等于1 源码（ n = max(round(n * gd), 1) if n > 1 else n）
  # module - 层的名字
  # args - 卷积核的个数
  [[-1, 1, Focus, [64, 3]],  # 0-P1/2  # 64要乘以卷积核的个数 64*0.5 = 32个特征图
   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
   [-1, 3, BottleneckCSP, [128]],
   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
   [-1, 9, BottleneckCSP, [256]],
   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
   [-1, 9, BottleneckCSP, [512]],
   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
   [-1, 1, SPP, [1024, [5, 9, 13]]],
   [-1, 3, BottleneckCSP, [1024, False]],  # 9
  ]

# YOLOv5 head
head:
  [[-1, 1, Conv, [512, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 6], 1, Concat, [1]],  # cat backbone P4
   [-1, 3, BottleneckCSP, [512, False]],  # 13

   [-1, 1, Conv, [256, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 4], 1, Concat, [1]],  # cat backbone P3
   [-1, 3, BottleneckCSP, [256, False]],  # 17 (P3/8-small)

   [-1, 1, Conv, [256, 3, 2]],
   [[-1, 14], 1, Concat, [1]],  # cat head P4
   [-1, 3, BottleneckCSP, [512, False]],  # 20 (P4/16-medium)

   [-1, 1, Conv, [512, 3, 2]],
   [[-1, 10], 1, Concat, [1]],  # cat head P5
   [-1, 3, BottleneckCSP, [1024, False]],  # 23 (P5/32-large)

   [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5) [17,20,23] #17层、20层、23层;
  ]

在data目录下添加一个dangerous.yaml训练数据配置文件，文件内容如下：

# download command/URL (optional)
download: bash data/scripts/get_voc.sh

# 训练集txt与验证集txt路径
train: data/xxx/train.txt
val: data/xxx/val.txt

# 总类别数
nc: 10

# 类名
names: ['lighter','scissors','powerbank','pressure','knife','zippooil','handcuffs','slingshot','firecrackers','nailpolish']

开始训练
训练命令

单卡：

python train.py --cfg models/yolov5s.yaml --data data/ODID.yaml --hyp data/hyps/hyp.scratch.yaml --epochs 100 --multi-scale --device 0

多卡：

python train.py --cfg models/yolov5s.yaml --data data/ODID.yaml --hyp data/hyps/hyp.scratch.yaml --epochs 100 --multi-scale --device 0，1

测试模型

python test.py --weights runs/train/exp/weights/best.pt --data data/ODID.yaml --device 0 --verbose
--weights: 训练得到的模型
--data：数据配置文件.txt
--device：选择gpu进行评测
--verbose：是否打印每一类的评测指标

模型推理

对于简单的项目需求，可以使用Gradio来快速部署。但是，当项目的复杂性增加，界面布局变得更复杂时，使用Gradio来管理和设置各个控件的位置会变得非常困难。在这种情况下，就可以考虑使用Qt。PyQt是一个Python绑定的Qt框架，它利用Qt的布局和设计工具，将拖放和设置好的界面转换为Python代码，从而简化了项目开发过程。
py qt5安装

pip install PyQt5
pip install PyQt5-tools

安装好PyQt5之后就可以开始编写推理代码了，pyqt能快速创建一个一推理的UI界面，项目分三种输入，图像、摄像头、视频。

#! /usr/bin/env python
# -*- coding: utf-8 -*-
"""
Run a YOLO_v3 style detection model on test images.
"""
import numpy as np
import argparse

import cv2
import torch
import torch.backends.cudnn as cudnn
from numpy import random

from models.experimental import attempt_load
from utils.general import check_img