yolov5使用笔记

前期准备

首先，将yolov5源代码下载至本地，并配置好运行yolov5代码所需环境。
其次，准备好自己已标注好的数据集。标注文件格式为YOLO格式。YOLO格式的标注文件一般以txt文件形式保存，否则按照以下格式对标注文件进行格式转换：

• 每一行存放一个标注对象
• 每行格式：类别 x坐标中心 y坐标中心 宽 高
• 框坐标必须是标准化的xywh格式，如果框坐标是以像素为单位的，将框的x中心坐标x_center和宽w除以图像的宽，将框的y中心坐标y_center和高h除以图像的高
• 类别的索引号要从0开始

转换标注文件(标注文件正确的可跳过此步)

由于我标注的文件格式是VOC格式(即，XML文件格式)，所以要进行标注文件格式的转换，我的存放数据文件夹目录如下：

• Annotations：存放标注文件
• JPEGImages：存放图片文件

具体代码如下：

"""
将VOC格式的标注文件(XML)转化为YOLO格式标注文件(TXT)
"""

import xml.etree.ElementTree as ET
from os import listdir

classes = ["cigarette","cigarette_hand","cigarette_mouth","smoky","smoking"]

def convert_annotation(image_id):
    in_file = open('.\\smoking\\Annotations\\%s.xml' % image_id)
    tree=ET.parse(in_file)
    root = tree.getroot()
    in_file.close()
    
    imgsize = root.find("size")
    imgh = int(imgsize.find("height").text)
    imgw = int(imgsize.find("width").text)

    for obj in root.iter('object'):
        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')
        x_center = (int(xmlbox.find('xmin').text)+int(xmlbox.find('xmax').text))/(2*imgw)
        y_center = (int(xmlbox.find('ymin').text)+int(xmlbox.find('ymax').text))/(2*imgh)
        w = (int(xmlbox.find('xmax').text)-int(xmlbox.find('xmin').text))/imgw
        h = (int(xmlbox.find('ymax').text)-int(xmlbox.find('ymin').text))/imgh
        b = (round(x_center,6), round(y_center,6), round(w,6), round(h,6))
        
        with open(".\\smoking\\Annotations\\%s.txt" % image_id, "a") as f:
            f.write(str(cls_id)+" " + " ".join([str(a) for a in b]))


for image_id in listdir(".\\smoking\\Annotations"):
    convert_annotation(image_id.split(".")[0])

训练集，验证集，测试集划分

将准备好的数据集按照以下格式存放：

mydata
  |--images
       |--train
       |--val
  |--labels
       |--train
       |--val

构建数据配置文件

根据coco128.yaml配置文件构建自己的数据集配置文件(对应设置进行替换即可)：

# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: ../coco128/images/train2017/  # 这里最好使用绝对路径
val: ../coco128/images/train2017/  # 这里最好使用绝对路径

# number of classes修改此处
nc: 80

# class names修改此处
names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
        'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
        'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
        'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
        'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
        'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
        'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
        'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
        'hair drier', 'toothbrush']

修改模型配置文件

我用的预训练权重是yolov5s，因此将models文件夹下的yolov5s.yaml中的类别数修改为自己的类别数。

# parameters
nc: 80  # number of classes修改此处
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]
  [[-1, 1, Focus, [64, 3]],  # 0-P1/2
   [-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)
  ]

训练模型

先定个epochs等于5，测试模型是否可以顺利运行：

python train.py --img 640 --batch 16 --epochs 5 --data coco128.yaml --weights yolov5s.pt

若是能够顺利运行，将epochs的值调大进行正式训练，训练完成后，权重文件保存在run文件夹下。

预测

python detect.py --source 0  # webcam
                            file.jpg  # image 
                            file.mp4  # video
                            path/  # directory
                            path/*.jpg  # glob
                            rtsp://170.93.143.139/rtplive/470011e600ef003a004ee33696235daa  # rtsp stream
                            rtmp://192.168.1.105/live/test  # rtmp stream
                            http://112.50.243.8/PLTV/88888888/224/3221225900/1.m3u8  # http stream

python detect.py --source inference/images --weights yolov5m.pt --conf 0.25 --view-img