【目标识别检测】Win10下训练YOLOv5人脸检测模型

【目标识别检测】Win10下训练YOLOv5人脸检测模型

更新时间：2021.1.9
操作系统：win10

环境搭建：https://blog.csdn.net/qq_45389690/article/details/111306902

准备数据集

下载Wider Face数据集
百度云：https://pan.baidu.com/s/1HyTz9beaCdXl26qMzofz2A
提取码：AFVH

在./VOC2012下创建classes.name，里面只写类别，每个类别一行，如

face

再在./VOC2012下创建voc2yolo.py文件，内容如下：

from __future__ import print_function
 
import os
import random
import glob
import xml.etree.ElementTree as ET
 
def xml_reader(filename):
    """ Parse a PASCAL VOC xml file """
    tree = ET.parse(filename)
    size = tree.find('size')
    width = int(size.find('width').text)
    height = int(size.find('height').text)
    objects = []
    for obj in tree.findall('object'):
        if(obj.find('name').text=='armor_blue' or obj.find('name').text=='armor_red'):
            obj_struct = {
     }
            obj_struct['name'] = obj.find('name').text
            bbox = obj.find('bndbox')
            obj_struct['bbox'] = [round(float(bbox.find('xmin').text)),
                                  round(float(bbox.find('ymin').text)),
                                  round(float(bbox.find('xmax').text)),
                                  round(float(bbox.find('ymax').text))]
            objects.append(obj_struct)
    return width, height, objects
 
 
def voc2yolo(filename):
    classes_dict = {
     }
    with open("classes.names") as f:
        for idx, line in enumerate(f.readlines()):
            class_name = line.strip()
            classes_dict[class_name] = idx
    
    width, height, objects = xml_reader(filename)
 
    lines = []
    for obj in objects:
        x, y, x2, y2 = obj['bbox']
        class_name = obj['name']
        label = classes_dict[class_name]
        cx = (x2+x)*0.5 / width
        cy = (y2+y)*0.5 / height
        w = (x2-x)*1. / width
        h = (y2-y)*1. / height
        line = "%s %.6f %.6f %.6f %.6f\n" % (label, cx, cy, w, h)
        lines.append(line)
 
    txt_name = filename.replace(".xml", ".txt").replace("Annotations", "labels")
    with open(txt_name, "w") as f:
        f.writelines(lines)
 
 
def get_image_list(image_dir, suffix=['jpg', 'jpeg', 'JPG', 'JPEG','png']):
    '''get all image path ends with suffix'''
    if not os.path.exists(image_dir):
        print("PATH:%s not exists" % image_dir)
        return []
    imglist = []
    for root, sdirs, files in os.walk(image_dir):
        if not files:
            continue
        for filename in files:
            filepath = "data/custom/" + os.path.join(root, filename) + "\n"
            if filename.split('.')[-1] in suffix:
                imglist.append(filepath)
    return imglist
 
 
def imglist2file(imglist):
    random.shuffle(imglist)
    train_list = imglist[:-100]
    valid_list = imglist[-100:]
    with open("train.txt", "w") as f:
        f.writelines(train_list)
    with open("valid.txt", "w") as f:
        f.writelines(valid_list)
 
 
if __name__ == "__main__":
    xml_path_list = glob.glob("Annotations/*.xml")
    for xml_path in xml_path_list:
        voc2yolo(xml_path)
 
 
    imglist = get_image_list("JPEGImages")
    imglist2file(imglist)

并运行该文件，生成train.txt和valid.txt。

在yolov5-master/data创建一个新的文件夹，目录结构如下：

-wider_face.yaml
-wider_face
--images #图片，把JPEGimages下所有图片复制到这里
---000001.jpg……
--labels #.txt标签
---000001.txt……
--classes.names
--train.txt
--valid.txt

在./data下创建.yaml文件,冒号后有空格

train: ./data/wider_face/
val: ./data/wider_face/
nc: 1
names: ['face']

训练模型

打开train.py，根据自己的需求和设备情况配置如下信息。

	parser.add_argument('--weights', type=str, default='./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='data/wider_face.yaml', help='data.yaml path')
    parser.add_argument('--hyp', type=str, default='data/hyp.scratch.yaml', help='hyperparameters path')
    parser.add_argument('--epochs', type=int, default=100)
    parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs')
    parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[train, test] image sizes')

运行train.py