基于Yolov5的吸烟行为识别检测

---

前言

  上一篇《win10系统下Yolov5目标检测环境搭建（Anaconda3+Pytorch+Yolov5，CPU，无GPU）》中，介绍了win10系统下yolov5的环境搭建步骤，本篇在此环境下，实现吸烟行为识别检测。

一、数据集smoke_data

1、Annotations文件夹

存放labelimg标注工具生成的.xml标注文件，每一张图片对应一个xml文件

2、ImageSets/Main文件夹

存放train.txt，val.txt，test.txt和trainval.txt文件

3、JPEGImages文件夹

存放待训练的图片文件，我使用的所有图片都为jpg格式

二、数据集划分

新建split_train_val.py文件，其作用主要是生成train.txt，val.txt，test.txt和trainval.txt文件，生成的这些文件，存放在ImageSets/Main文件夹下。py脚本内容如下：

# coding:utf-8
 
import os
import random
import argparse
 
parser = argparse.ArgumentParser()
#xml文件的地址，根据自己的数据进行修改 xml一般存放在Annotations下
parser.add_argument('--xml_path', default='smoke_data/Annotations', type=str, help='input xml label path')
#数据集的划分，地址选择自己数据下的ImageSets/Main
parser.add_argument('--txt_path', default='smoke_data/ImageSets/Main', type=str, help='output txt label path')
opt = parser.parse_args()
 
trainval_percent = 0.9
train_percent = 0.9
xmlfilepath = opt.xml_path
txtsavepath = opt.txt_path
total_xml = os.listdir(xmlfilepath)
if not os.path.exists(txtsavepath):
    os.makedirs(txtsavepath)
 
num = len(total_xml)
list_index = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list_index, tv)
train = random.sample(trainval, tr)
 
file_trainval = open(txtsavepath + '/trainval.txt', 'w')
file_test = open(txtsavepath + '/test.txt', 'w')
file_train = open(txtsavepath + '/train.txt', 'w')
file_val = open(txtsavepath + '/val.txt', 'w')
 
for i in list_index:
    name = total_xml[i][:-4] + '\n'
    if i in trainval:
        file_trainval.write(name)
        if i in train:
            file_train.write(name)
        else:
            file_val.write(name)
    else:
        file_test.write(name)
 
file_trainval.close()
file_train.close()
file_val.close()
file_test.close()

执行命令，如下：

python split_train_val.py

可以看到上述txt文件，已经生成，如下：

三、使用labelimg工具标注

进入之前建立的pytorch环境，安装labelimg，命令如下：

pip install labelimg

输入labelimg命令，启动，如下：

选择使用“PascalVOC”格式；
点击“open dir”按钮，选择JPEGImages文件夹目录；
点击“change save dir”按钮，选择Annotations文件夹目录；
点击“create RectBox”按钮，选择区域进行标注；
标注完成后，点击“save”按钮，进行保存，xml会存放在Annotations文件夹下。

四、将数据集格式转化成yolo_txt格式

  即将xml文件转化成txt文件，每一张图片对应一个txt文件，文件中每一行对应一个目标信息，包括class, x_center, y_center, width和height。
  新建voc_label.py脚本，将classes 里面的内容改成自己的类别，若是多个类别，则用逗号分隔，内容如下：

# -*- coding: utf-8 -*-
import xml.etree.ElementTree as ET
import os
from os import getcwd
 
sets = ['train', 'val', 'test']
classes = [
    "smoke"]   # 改成自己的类别
abs_path = os.getcwd()
print(abs_path)
 
def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = (box[0] + box[1]) / 2.0 - 1
    y = (box[2] + box[3]) / 2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return x, y, w, h
 
def convert_annotation(image_id):
    in_file = open('F:/SVN-ZJKY/YiFeiShouJiRobot/yolov5-master/smoke_data/Annotations/%s.xml' % (image_id), encoding='UTF-8')
    out_file = open('F:/SVN-ZJKY/YiFeiShouJiRobot/yolov5-master/smoke_data/labels/%s.txt' % (image_id), 'w')
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)
    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        # difficult = obj.find('Difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
             float(xmlbox.find('ymax').text))
        b1, b2, b3, b4 = b
        # 标注越界修正
        if b2 > w:
            b2 = w
        if b4 > h:
            b4 = h
        b = (b1, b2, b3, b4)
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
 
wd = getcwd()
for image_set in sets:
    if not os.path.exists('F:/SVN-ZJKY/YiFeiShouJiRobot/yolov5-master/smoke_data/labels/'):
        os.makedirs('F:/SVN-ZJKY/YiFeiShouJiRobot/yolov5-master/smoke_data/labels/')
    image_ids = open('F:/SVN-ZJKY/YiFeiShouJiRobot/yolov5-master/smoke_data/ImageSets/Main/%s.txt' % (image_set)).read().strip().split()
    list_file = open('smoke_data/%s.txt' % (image_set), 'w')
    for image_id in image_ids:
        list_file.write(abs_path + '/smoke_data/images/%s.jpg\n' % (image_id))
        convert_annotation(image_id)
    list_file.close()

执行命令，如下：

python voc_label.py

运行该脚本，会生成labels文件夹和三个包含数据集的txt文件，即该脚本的作用：将训练集、验证集、测试集生成label标签（训练中要用到），同时将数据集路径导入txt文件中。如下：

其中，labels文件夹，存放xml文件转化后的txt文件；
train.txt文件，记录的是训练集图像绝对路径；
test.txt文件，记录的是测试集图像绝对路径；
val.txt文件，记录的是验证集图像绝对路径；

五、增加自定义数据集文件

在yolov5-master/data/文件夹下，新建smoke_data.yaml文件，内容如下：

train: F:/SVN-ZJKY/YiFeiShouJiRobot/yolov5-master/smoke_data/train.txt
val: F:/SVN-ZJKY/YiFeiShouJiRobot/yolov5-master/smoke_data/val.txt
test: F:/SVN-ZJKY/YiFeiShouJiRobot/yolov5-master/smoke_data/test.txt
 
# Classes
nc: 1  # number of classes
names: ['smoke']  # class names

注意：修改自己对应的类别个数和类别名。

六、修改模型文件yolov5s.yaml

我选择使用的是yolov5s模型，进入yolov5-master/model/文件夹下，找到yolov5s.yaml文件，修改对应的类别个数nc，如下：

# YOLOv5  by Ultralytics, GPL-3.0 license

# Parameters
#nc: 80  # number of classes
nc: 1
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.50  # layer channel multiple
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 v6.0 backbone
backbone:
  # [from, number, module, args]
  [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2
   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
   [-1, 3, C3, [128]],
   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
   [-1, 6, C3, [256]],
   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
   [-1, 9, C3, [512]],
   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
   [-1, 3, C3, [1024]],
   [-1, 1, SPPF, [1024, 5]],  # 9
  ]

# YOLOv5 v6.0 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, C3, [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, C3, [256, False]],  # 17 (P3/8-small)

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

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

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

七、训练

修改train.py，主要修改内容，如下：

def parse_opt(known=False):
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='initial weights path')
    parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml', help='model.yaml path')
    parser.add_argument('--data', type=str, default=ROOT / 'data/smoke_data.yaml', help='dataset.yaml path')
    parser.add_argument('--hyp', type=str, default=ROOT / 'data/hyps/hyp.scratch.yaml', help='hyperparameters path')
    parser.add_argument('--epochs', type=int, default=50)
    parser.add_argument('--batch-size', type=int, default=10, help='total batch size for all GPUs')

各参数详情如下：

其中，–epochs表示训练次数（如果显卡好可以调大一点，默认为300）， 而–batch-size 是表示一次处理完多少张图片才进行权重更新，显卡不好时可以调小。

执行如下命令，开始训练：

python train.py

训练过程如下：

训练完成后，结果保存在runs/train/exp13目录下，如下：

其中，训练好的权重文件在weights文件夹下：

best.pt文件，表示训练过程中得到的效果最佳的模型；
last.pt文件，表示训练结束后最终得到的模型；

八、测试

修改detect.py文件，主要修改内容，如下：

def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'runs/train/exp13/weights/best.pt', help='model path(s)')
    #parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model path(s)')
    #默认 图片
    parser.add_argument('--source', type=str, default=ROOT / 'data/smokeImages', help='file/dir/URL/glob, 0 for webcam')
    #摄像头
    #parser.add_argument('--source', type=str, default=0, help='file/dir/URL/glob, 0 for webcam')

各参数详情如下：

在py文件中，可以看到，我们将待测试的图片存放在data/smokeImages文件夹下，修改完成后执行如下命令：

python detect.py

所有图片处理完成后，识别结果保存在runs/detect/exp28目录下，如下：