纪录自己YOLOv5 训练自己数据集的踩坑记录

自己数据集的格式

我自己的数据集格式是类voc形式的
—annotation
···.xml
···.xml
—ImageSets
Mains
—JPEGImages
···.jpg

我的格式大致如图。

修改yolov5中的data/VOC.yaml

首先不多说直接

git clone https://github.com.cnpmjs.org/ultralytics/yolov5.git

导入自己服务器中之后,修改data/VOC.yaml

# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ../datasets/VOC  
train: # train images (relative to 'path')  16551 images
  - images/train2012
  - images/train2007
  - images/val2012
  - images/val2007
val: # val images (relative to 'path')  4952 images
  - images/test2007
test: # test images (optional)
  - images/test2007

# Classes
nc: 20  # number of classes
names: ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog',
        'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor']  # class names

依次修改:
path文件路径,(这里要注意自己的绝对路径问题,容易报错,尤其是自己按照自己的命名方式,一定注意)
train
val
test
依次进行修改 如果是VOC格式类型的就直接在train val test后面➕上txt的后缀。
如果是文件夹类型的 就无须修改。

nc:有多少类写多少类

name:种类名称照葫芦画瓢就可以。

下面我给出Bubbliiiing大佬的代码,生成train test val.txt文件
基于我自己需求进行了更改:

import os
import random 
random.seed(0)

xmlfilepath='VOCdevkit/VOC2007/Annotations' #依照自己的数据集位置进行更改
saveBasePath="VOCdevkit/VOC2007/ImageSets/Main/"#依照自己的数据集位置进行更改
 
#----------------------------------------------------------------------#
#   想要增加测试集修改trainval_percent
#   train_percent不需要修改
#----------------------------------------------------------------------#
trainval_percent=0.9
train_percent=1

temp_xml = os.listdir(xmlfilepath)
total_xml = []
for xml in temp_xml:
    if xml.endswith(".xml"):
        total_xml.append(xml)

num=len(total_xml)  
list=range(num)  
tv=int(num*trainval_percent)  
tr=int(tv*train_percent)  
trainval= random.sample(list,tv)  
train=random.sample(trainval,tr)  
 
print("train and val size",tv)
print("traub suze",tr)
ftrainval = open(os.path.join(saveBasePath,'trainval.txt'), 'w')  
ftest = open(os.path.join(saveBasePath,'test.txt'), 'w')  
ftrain = open(os.path.join(saveBasePath,'train.txt'), 'w')  
fval = open(os.path.join(saveBasePath,'val.txt'), 'w')  
 
for i  in list:  
    name='VOCdevkit/VOC2007/JPEGImages/'+total_xml[i][:-4]+'.jpg'+'\n'
    #这里前面的VOCdevkit/VOC2007···这些是依照自己数据集位置进行更改
    #我试过不加前缀 以及 .jpg 显示找不到文件 所以就自己进行修改。
    if i in trainval:  
        ftrainval.write(name)  
        if i in train:  
            ftrain.write(name)  
        else:  
            fval.write(name)  
    else:  
        ftest.write(name)  
  
ftrainval.close()  
ftrain.close()  
fval.close()  
ftest .close()

进行label的生成

yolov5有一点注意的是 需要带有label文件,这个是voc数据集没有的,所以还需要自己生成自己文件的label文件。
我下面也给出label文件的代码,以及一个可以查询自己voc数据集的代码。

查询VOC数据集种类以及数量的代码

import os
import xml.dom.minidom

xml_path = r'E:\\pycharm\\dataset\\Yolov3&Yolov4&Yolov5\\Yolov5\\VOCdevkit\\VOC2007\\Annotations\\'
#xml_path 需要更改为自己的路径 Windows系统就用\\ ,liunx系统就使用/,实在不会直接绝对路径就行
files = os.listdir(xml_path)

gt_dict = {}

if __name__ == '__main__':

    for xm in files:
        xmlfile = xml_path + xm
        dom = xml.dom.minidom.parse(xmlfile)  # 读取xml文档
        root = dom.documentElement  # 得到文档元素对象
        filenamelist = root.getElementsByTagName("filename")
        filename = filenamelist[0].childNodes[0].data
        objectlist = root.getElementsByTagName("object")
        ##
        for objects in objectlist:
            namelist = objects.getElementsByTagName("name")
            objectname = namelist[0].childNodes[0].data
            if objectname == '-':
                print(filename)
            if objectname in gt_dict:
                gt_dict[objectname] += 1
            else:
                gt_dict[objectname] = 1
            # for nl in namelist:
            #     objectname = nl.childNodes[0].data
            #     if objectname in gt_dict:
            #         gt_dict[objectname] += 1
            #     else:
            #         gt_dict[objectname] = 1
    dic = sorted(gt_dict.items(), key=lambda d: d[1], reverse=True)
    print(dic)
    print(len(dic))

VOC数据集生成label的代码

# -*- coding: utf-8 -*-
# @Time    : 2020/7/29 20:29
# @Author  : PeterH
# @Email   : peterhuang0323@outlook.com
# @File    : data_cfg.py
# @Software: PyCharm
# @Brief   : 生成测试、验证、训练的图片和标签

import os
import shutil
from pathlib import Path
from shutil import copyfile

from PIL import Image, ImageDraw
from xml.dom.minidom import parse
import numpy as np
from tqdm import tqdm

FILE_ROOT = Path(r"./VOCdevkit/VOC2007")

# 原始数据集
IMAGE_SET_ROOT = FILE_ROOT.joinpath(r"ImageSets/Main")  # 图片区分文件的路径
print(IMAGE_SET_ROOT)
IMAGE_PATH = FILE_ROOT.joinpath(r"JPEGImages")  # 图片的位置
print(IMAGE_PATH)
ANNOTATIONS_PATH = FILE_ROOT.joinpath(r"Annotations")  # 数据集标签文件的位置
print("=============")
print(ANNOTATIONS_PATH)
LABELS_ROOT = FILE_ROOT.joinpath(r"Labels")  # 进行归一化之后的标签位置
print(LABELS_ROOT)

# YOLO 需要的数据集形式的新数据集
DEST_IMAGES_PATH = Path(r"score/images")  # 区分训练集、测试集、验证集的图片目标路径
DEST_LABELS_PATH = Path(r"score/labels")  # 区分训练集、测试集、验证集的标签文件目标路径


def cord_converter(size, box):
    """
    将标注的 xml 文件标注转换为 darknet 形的坐标
    :param size: 图片的尺寸: [w,h]
    :param box: anchor box 的坐标 [左上角x,左上角y,右下角x,右下角y,]
    :return: 转换后的 [x,y,w,h]
    """

    x1 = int(box[0])
    y1 = int(box[1])
    x2 = int(box[2])
    y2 = int(box[3])

    dw = np.float32(1. / int(size[0]))
    dh = np.float32(1. / int(size[1]))

    w = x2 - x1
    h = y2 - y1
    x = x1 + (w / 2)
    y = y1 + (h / 2)

    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return [x, y, w, h]


def save_label_file(img_jpg_file_name, size, img_box):
    """
    保存标签的解析文件
    :param img_jpg_file_name:
    :param size:
    :param img_box:
    :return:
    """
    save_file_name = LABELS_ROOT.joinpath(img_jpg_file_name).with_suffix('.txt')
    with open(save_file_name, "a+") as f:
        for box in img_box:
            if box[0] == 'Leconte':  #注意这里需要进行更改为自己的数据集种类,此外有多少类依次累加就可以了
                cls_num = 0
            elif box[0] == 'coleoptera':
                cls_num = 1
            elif box[0] == 'armandi':  
                cls_num = 2
            elif box[0] == 'Boerner':
                cls_num = 3
            elif box[0] == 'acuminatus':  
                cls_num = 4
            elif box[0] == 'Linnaeus':
                cls_num = 5
            elif box[0] == 'linnaeus':
                cls_num = 6
            else:
                continue
            new_box = cord_converter(size, box[1:])  # 转换坐标
            f.write(f"{cls_num} {new_box[0]} {new_box[1]} {new_box[2]} {new_box[3]}\n")


def test_dataset_box_feature(file_name, point_array):
    """
    使用样本数据测试数据集的建议框
    :param file_name: 图片文件名
    :param point_array: 全部的点 [建议框sx1,sy1,sx2,sy2]
    :return: None
    """
    im = Image.open(IMAGE_PATH.joinpath(file_name).with_suffix(".jpg"))
    im_draw = ImageDraw.Draw(im)
    for box in point_array:
        x1 = box[1]
        y1 = box[2]
        x2 = box[3]
        y2 = box[4]
        im_draw.rectangle((x1, y1, x2, y2), outline='red')

    im.show()


def get_xml_data(img_xml_file: Path):
    """
    获取 xml 数据
    :param img_xml_file: 图片路径
    :return:
    """
    dom = parse(str(img_xml_file))
    xml_root = dom.documentElement
    img_name = xml_root.getElementsByTagName("filename")[0].childNodes[0].data
    img_size = xml_root.getElementsByTagName("size")[0]
    objects = xml_root.getElementsByTagName("object")
    img_w = img_size.getElementsByTagName("width")[0].childNodes[0].data
    img_h = img_size.getElementsByTagName("height")[0].childNodes[0].data
    img_c = img_size.getElementsByTagName("depth")[0].childNodes[0].data
    img_box = []
    for box in objects:
        cls_name = box.getElementsByTagName("name")[0].childNodes[0].data
        x1 = int(box.getElementsByTagName("xmin")[0].childNodes[0].data)
        y1 = int(box.getElementsByTagName("ymin")[0].childNodes[0].data)
        x2 = int(box.getElementsByTagName("xmax")[0].childNodes[0].data)
        y2 = int(box.getElementsByTagName("ymax")[0].childNodes[0].data)
        img_box.append([cls_name, x1, y1, x2, y2])

    # test_dataset_box_feature(img_xml_file.name, img_box)
    save_label_file(img_xml_file.name, [img_w, img_h], img_box)


def copy_data(img_set_source, img_labels_root, imgs_source, dataset_type):
    """
    将标签文件和图片复制到最终数据集文件夹中
    :param img_set_source: 原数据集图片总路径
    :param img_labels_root: 生成的 txt 标签总路径
    :param imgs_source:
    :param dataset_type: 生成数据集的种类
    :return:
    """
    file_name = img_set_source.joinpath(dataset_type).with_suffix(".txt")  # 获取对应数据集种类的图片

    # 判断目标图片文件夹和标签文件夹是否存在,不存在则创建
    os.makedirs(FILE_ROOT.joinpath(DEST_IMAGES_PATH, dataset_type), exist_ok=True)
    os.makedirs(FILE_ROOT.joinpath(DEST_LABELS_PATH, dataset_type), exist_ok=True)

    with open(file_name, encoding="UTF-8") as f:
        for img_name in tqdm(f.read().splitlines()):
            img_sor_file = imgs_source.joinpath(img_name).with_suffix('.jpg')
            label_sor_file = img_labels_root.joinpath(img_name).with_suffix('.txt')

            # 复制图片
            dict_file = FILE_ROOT.joinpath(DEST_IMAGES_PATH, dataset_type, img_name).with_suffix('.jpg')
            copyfile(img_sor_file, dict_file)

            # 复制 label
            dict_file = FILE_ROOT.joinpath(DEST_LABELS_PATH, dataset_type, img_name).with_suffix('.txt')
            copyfile(label_sor_file, dict_file)


if __name__ == '__main__':
    root = ANNOTATIONS_PATH  # 数据集 xml 标签的位置

    if LABELS_ROOT.exists():
        # 清空标签文件夹
        print("Cleaning Label dir for safety generating label, pls wait...")
        shutil.rmtree(LABELS_ROOT)
        print("Cleaning Label dir done!")
    LABELS_ROOT.mkdir(exist_ok=True)  # 建立 Label 文件夹

    # 生成标签
    print("Generating Label files...")
    with tqdm(total=len(os.listdir(root))) as p_bar:
        for file in root.iterdir():
            p_bar.update(1)
            get_xml_data(file)

    # 将文件进行 train、val、test 的区分
    for dataset_input_type in ["train", "val", "test"]:
        print(f"Copying data {dataset_input_type}, pls wait...")
        copy_data(IMAGE_SET_ROOT, LABELS_ROOT, IMAGE_PATH, dataset_input_type)

此代码引用@Author : PeterH 大佬的

训练

    parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='initial weights path')
    parser.add_argument('--cfg', type=str, default='', help='model.yaml path')
    parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.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=300)
    parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs')

觉得麻烦的同志 可以进行train.py文件上面的默认值进行更改,无非就是–data 这一类
还有 --cfg 这一类更改为自己适合的文件。

然后运行:Python train.py 就OK了。
在这里插入图片描述

评论 19
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值