TOLOv5训练自己的数据集--漫画人物检测

TOLOv5训练自己的数据集

这是我的第一篇文章，记录我的学习进展。
任务：用YOLOv5检测漫画人物头像，其中训练数据集为自己标注的256张《19天》漫画截图，在官方预训练权重下进行训练得到所需检测器。

YOLOv5模型下载

YOLOv5源码 github:https://github.com/ultralytics/yolov5
先在unbuntu中使用Anaconda创建了个新的虚拟环境：

conda activate # 打开Ananconda Prompt
conda create -n env_name python=3.7 # 创建虚拟环境[env_name]
source activate env_name # 激活虚拟环境

之后就是遵循官网的安装步骤：

# cd至想安放它的地方[parent]
git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txt
# 防止下载至全局环境中可执行：python -m pip install -r requirements.txt

安装结束之后试了一下检测器：

python detect.py --source ./img.jpg[你想检测的图片或文件夹路径] --device 0

如果是第一次运行检测器，会自动下载权值文件yolov5s.pt，要多等一会儿…
检测结果图像： [会保存在yolov5/runs/detector/exp文件夹中]

使用LabelImage标注图片

图片标注工具LabelImage下载链接：https://github.com/Programming99119/labelImg-master
使用LabelImage标注图片：

Open：一次只能打开一张图片[不推荐]；
Open Dir：打开待标注图片所处文件夹；
Change Save Dir：保存标注文件的文件夹；
“Use defualt label”：使用默认标签。
操作快捷键：
A键：上一张图片；
D键：下一张图片；
W键：画框操作；
Ctrl+S：保存操作。
标注结束后，可以在“Change Save Dir”文件夹位置看到与图片相对应的同名同数目的xml文件。
这里我居然把展正希的名字简写敲成了zhanwwzx，全部标注完我才发现标错了 5555
我只标注了‘见一’和‘展正希’的头像，也就是两类标签，因为不想肝啦
我标注的256张《19天》漫画人物训练数据集在百度网盘：
https://pan.baidu.com/s/15UcZjTQLC7J7RfJ9dlTyyA
提取码：dzk5

生成所需数据集

yolov5训练所需数据集格式详见官网说明，
我参考学习的链接：https://blog.csdn.net/qq_36756866/article/details/109111065
(里面的内容更丰富，还有各种情况下的报错与纠正)
1.首先需要创建一个文件夹datasets，它需要与yolov5文件夹平行放置，
满足如下分布：

2.编写一个数据批处理小程序xml_to_txt.py，把xml文件转化成yolov5所需的txt文件，并保存在labels文件夹下。

# -*- coding: utf-8 -*-
import xml.etree.ElementTree as ET
import os
import random
from os import getcwd

sets = ['train', 'val', 'test']
classes = ['jianyi', 'zhanwwzx']   # 改成自己的数据集类别
Path = '/data/datasets/oneday' # 改成自己创建的数据集所在绝对路径
trainval_percent = 1.0  
train_percent = 0.9

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):
    # print(image_id)
    in_file = open(Path+f'/Annotations/{image_id}.xml', encoding='UTF-8')
    out_file = open(Path+f'/Annotations/{image_id}.txt', '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'):
        cls = obj.find('name').text
        if cls not in classes:
            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')

total_xml = os.listdir(Path+'/Annotations')
print(total_xml)
num = len(total_xml)
print('total_xml: ', num)
list_index = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
print('random-sample...')
tv = random.sample(list_index, tv)
tr = random.sample(tv, tr)
val = []
train = []
for index in list_index:
    if index in tr:
        train.append(total_xml[index][:-4])
    else:
        val.append(total_xml[index][:-4])

sets2 = {'train': train, 'val': val, 'test': []}
print('convert-annotation...')
for image_set in sets:
    if not os.path.exists(Path+'/labels'):
        os.makedirs(Path+'/labels')
    list_file = open(Path+f'/{image_set}.txt', 'w')
    for image_id in sets2[image_set]:
        list_file.write(Path+f'/images/{image_id}.jpg\n')
        convert_annotation(image_id)
    list_file.close()
print('end.')

运行完之后，labels文件夹下会有同名同数量的txt文件；
同时，oneday文件夹下会多出三个txt文件，train.txt 、val.txt 、test.txt，它们分别记录有训练数据集、验证数据集、测试数据集对应的绝对路径。
打开查看会发现其中的test.txt文件为空的，这是因为程序中 trainval_percent = 1.0 表示全部数据用于训练和验证了，没有准备测试数据集； train_percent = 0.9 表示其中9/10的数据用于训练、1/10的数据用于验证。

更改配置文件

编写一个yaml配置文件：
记录train.txt和val.txt的绝对路径；
记录训练集类的数目；
按顺序记录类的名称。

train: /home/xxx/PycharmProjects/datasets/oneday/train.txt
val: /home/xxx/PycharmProjects/datasets/oneday/val.txt
# number of classes 
nc: 2
# class names 
names: ['jianyi','zhanzx']

因为此次训练的数据集较为简单，因此使用yolov5s模型，只需更改一个参数，
将yolov5/models/yolov5s.yaml中类的数目改为2即可：

# YOLOv5 🚀 by Ultralytics, GPL-3.0 license

# Parameters
nc: 2  # number of classes 【更改此处】
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.50  # layer channel multiple
...

训练

执行：(cd至yolov5文件夹下)

python train.py --img 640 --batch 16 --epoch 300 --data [yaml配置文件路径] --cfg models/yolov5s.yaml --weight yolov5s.pt --device ‘0’

结果：

训练得到的权值文件保存至runs/train/exp/weights文件夹下，有两个pt文件，last.pt和best.pt，它们都可以用于测试。

测试

执行:

python detect.py --weights runs/train/exp/weights/best.pt --source [你想检测的图片或含图片的文件夹] --device 0

若需要预测框的数据保存至txt文件中，则添加 --save-txt
我得到的预测结果：



😃😃😃
以上是我做的全部过程，
如有错误，感谢指正。