yolov5从0基础搭建以及训练自己的数据集

1.下载代码
yolov5
2.搭建环境
首先安装cuda 推荐安装10.2 以及对应的cudnn（可选择8.3.2）版本（30系显卡要安装11以上）
cuda官网链接
这部分就不赘述了 网上很多教程
然后使用 anaconda 创建虚拟环境 python =3.8 或者3.7都行

conda create -n yolov5 python=3.8

然后激活 环境

conda activate yolov5

pytorch老版本网址
这里选择的1.8 只要大于1.7就可以，下面命令安装的是cuda=10.2的GPU版本，选择合适自己cuda版本的命令

conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cudatoolkit=10.2 -c pytorch

安装完成后
也可安装cpu版（cpu和gpu 二选一）

conda install pytorch==1.8.1 torchvision==0.9.1 torchaudio==0.8.1 cpuonly -c pytorch

然后安装其他的包：
在anaconda prompt 下切换到代码路径（包含yolov5中requirements.txt文件的路径）然后输入如下代码

pip install -r requirements.txt

最后安装一个labelimg 打标用的

pip install labelimg

打完标后就是xml转换txt 代码如下 只需修改 classes以及路径即可

import xml.etree.ElementTree as ET
import os
import glob
import cv2

classes = ['qiao', 'zonglvshu', 'renxinghengdao', 'xiaofangshuang', 'honglvdeng', 'zixingche', 'xiaojiaoche', 'yancong', 'louti', 'jidongche', 'tuolaji', 'bus', 'school_bus', 'motuo', 'chuzuche', 'tingchejishiqi']



def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x1 = box[0] * dw
    x2 = box[1] * dw
    x3 = box[1] * dw
    x4 = box[0] * dw
    y1 = box[2] * dh
    y2 = box[2] * dh
    y3 = box[3] * dh
    y4 = box[3] * dh

    return x1, x2, x3, x4, y1, y2, y3, y4

def convert2(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = box[0] + w / 2
    y = box[2] + h / 2
    x = x * dw
    y = y * dh
    w = w * dw
    h = h * dh
    return x, y, w, h


# 对于单个xml的处理
def convert_annotation(label_xml, label_txt_path, label_xml_path):
    basename = os.path.splitext(label_xml)[0]
    # print(basename)
    in_file = open(os.path.join(label_xml_path, label_xml), encoding='utf-8')  # 图片对应的xml地址
    out_file = open(os.path.join(label_txt_path, '{}.txt'.format(basename)), 'w')
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    # if size is None:
    # img = cv2.imread(os.path.join(r'D:\download\baidu\123\posetive_img', "{}.jpg".format(basename)))
    # w, h = img.shape[1::-1]
    #
    # else:
    #
    #
    w = int(size.find('width').text)
    h = int(size.find('height').text)

    # 在一个XML中每个Object的迭代
    b = root.iter('object')
    for obj in root.iter('object'):
        # iter()方法可以递归遍历元素/树的所有子元素
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        # 如果训练标签中的品种不在程序预定品种，或者difficult = 1，跳过此object
        if cls not in classes or int(difficult) == 1:
            print(cls)
            print(basename)
            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))
        bb = convert2((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')


def generate_txt(label_xml_path, label_txt_path):
    if not os.path.exists(label_txt_path):  # 不存在文件夹
        os.makedirs(label_txt_path)
    label_xmls = os.listdir(label_xml_path)
    for label_xml in label_xmls:
        convert_annotation(label_xml, label_txt_path, label_xml_path)


if __name__ == '__main__':

    generate_txt(label_xml_path=r'D:\xianyu\gugeyanzheng\22222\label',
                 label_txt_path=r'D:\xianyu\gugeyanzheng\22222\txt')

然后就是随机分配训练集以及验证集
首先在yolov5里面手动创建dataset文件夹 文件目录如下：

然后运行随机分配代码：

import os
import shutil
import random
path = r'D:\xianyu\gugeyanzheng\img2'  # 图片路径
files = os.listdir(path)
random.shuffle(files)
random.shuffle(files)
random.shuffle(files)
save_path = r'D:\document\github\yolov5-master\dataset'  # 保存数据的路径
ra = 0.9 # 随机分配的比例
form = 'jpg' # 文件的格式
tihuan_name = 'img2'  # 图片的最后文件路径
train = files[:int(ra*len(files))]
val = files[int(ra*len(files)):]
for file in train:
    fp = os.path.join(path, file)
    shutil.copyfile(fp, os.path.join(save_path, 'images', 'train', file))
    shutil.copyfile(fp.replace(tihuan_name, 'txt').replace('.{}'.format(form), '.txt'), os.path.join(save_path, 'labels', 'train', file.replace('.{}'.format(form), '.txt')))
for file in val:
    fp = os.path.join(path, file)
    shutil.copyfile(fp, os.path.join(save_path, 'images', 'val', file))
    shutil.copyfile(fp.replace(tihuan_name, 'txt').replace('.{}'.format(form), '.txt'), os.path.join(save_path, 'labels', 'val', file.replace('.{}'.format(form), '.txt')))

分配完成后，就在data路径下产检自己的data.yaml ,只需修改nc（类别数） names（类别的名字）train以及val路径

# Custom data for safety helmet
# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: /home/data/yolo/images/train
val: /home/data/yolo/images/val
# number of classes
nc: 2
# class names
names: ['phone', 'person']

然后去data/hype/hyp.scratch.yaml ，里面修改超参数，重点看下数据增强的方式 是否与数据集匹配
然后就是改下train.py 文件 主要修改就是下面这些，第一个是预训练权重，第二个是网络的配置文件（权重和配置文件二选一就可），
然后改下数据的yaml文件，然后就是epochs 根据需要改，最后就是batch-size 根据电脑的性能改。最后就是修改device，使用cpu或gpu

parser.add_argument('--weights', type=str, default='./pretrained/yolov5s.pt', help='initial weights path')
parser.add_argument('--cfg', type=str, default='', help='model.yaml path')
parser.add_argument('--data', type=str, default='./data/head_data.yaml', help='dataset.yaml path')
parser.add_argument('--hyp', type=str, default='data/hyps/hyp.scratch.yaml', help='hyperparameters path')
parser.add_argument('--epochs', type=int, default=150)
parser.add_argument('--batch-size', type=int, default=2, help='total batch size for all GPUs, -1 for autobatch')
parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='train, val image size (pixels)')
parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')

改完后就可以开始训练了。

下面开始测试：
打开detect文件 第一个是权重 修改训练好的权重，刚训练好的默认保存在run/train路径下。第二个就是测试的文件 可以是单张图片 也可以使路径，也可以使摄像头0，后面是图像尺寸、置信度以及device 根据需要修改。

parser.add_argument('--weights', nargs='+', type=str, default='./pretrained/yolov3-tiny.pt', help='model path(s)')
parser.add_argument('--source', type=str, default=r'D:\download\coco\val\img', help='file/dir/URL/glob, 0 for webcam')
parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w')
parser.add_argument('--conf-thres', type=float, default=0.25, help='confidence threshold')
parser.add_argument('--iou-thres', type=float, default=0.45, help='NMS IoU threshold')
parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')

最后测试的结果会被保存在run/detect路径下