前言

  上一篇《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目录下,如下: