yolov5ds训练步骤

参考链接

• 🍅yolov5ds：Yolov5同时进行目标检测和分割分割（yolov5ds作者的博客介绍）
  • github地址：👉yolov5ds
  • 训练yolov5ds案例：用YOLOv5ds训练自己的数据集——同时检测和分割
  • 对上训练yolov5ds案例的补充：用YOLOv5ds训练自己的数据集，注意点！
• yolov5ds-断点训练、继续训练（yolov5同样使用）

以下步骤是参考：用YOLOv5ds训练自己的数据集——同时检测和分割

⭐2024年2月20日更新：分享一下我用过的代码，不保证一定有用哈：百度网盘链接，包含unet、yolov5ds

0、配置环境

不在赘述，跟YOLOv5差不多

1、下载预训练模型——推荐

在yolov5ds-main根目录新建weights文件夹

下载yolov5预训练模型Releases · ultralytics/yolov5 · GitHub放到weights文件夹中

我下载的是yolov5s.pt，下面均以yolov5s.pt为例

2、准备数据集——非常关键

在yolov5ds-main根目录新建paper_data文件夹

paper_data文件夹下新建det和seg两个文件夹

det文件夹存放检测数据集

seg文件夹存放分割数据集

---

🍀对于数据集，我的整体步骤是：

（1）运行：paper_data/det/json2txt.py，生成json对应的txt文件

（2）运行：paper_data/det/split.py，以9:1划分train、val（没有划分test），存放到paper_data/det/ImageSet/Main

（3）运行：paper_data/det/voc_labels.py，通过（2）划分的数据集将文件绝对路径存放到相应txt中，存放到paper_data/det

（4）运行：paper_data/seg/getmask.py，将raw_data中所有的.png图片复制到paper_data/seg/labels中

（5）运行：paper_data/seg/segsplit.py，按照train.txt和val.txt划分到paper_data/seg/images或labels下的train、val文件夹下

❀❀❀❀❀❀❀❀❀❀❀❀❀❀❀❀❀❀❀❀❀❀❀❀❀❀❀

关于以下我制作数据集的代码，我都分享出来了，有需要的好兄弟自取：👉百度网盘：paper_data

---

2-1、det文件夹下

images文件夹下存放.jpg图像（或者.png格式）

labels文件夹下存放.txt标签文件

Annotations文件夹下存放xml标签文件（这个在训练过程其实用不到，只是如果标签文件是xml格式的话，就暂存在这个文件夹，然后会通过3.voc_labels.py转换成txt格式且划分到labels中的train或val子文件下）

注：xml、txt文件应和对应图像名称相同

再注：

• 因为，我的数据集格式只有实例分割数据集：jpg图像+json格式标签，所以这里还多了json2txt.py文件

• 有xml检测标签的，就直接执行2、3步

• 直接有txt，也只执行2、3步，但是第3步中要注释掉调用convert_annotation(image_id)函数

1. json2txt.py

（1）paper_data/det/json2txt.py，生成json对应的txt文件

import json
import os
import os.path
import re


# 'Bolt hole1':螺栓孔 'Grouting hole':灌浆孔 'cable':电线电缆 'pipe':管 'crack':裂缝

# Support \ Bolt_hole \ Grouting_hole \ Cable \ Pipe \ Sign \ Signal_light \ Railway \ PJB \ Instrument_box \ Crack \ Falling_block
#    0          1             2           3       4      5          6           7       8         9             10         11

# 遍历获得所有标注类别
def getclass(rootdir):
    classes = []
    for file in os.listdir(rootdir):
        load_f = open(os.path.join(rootdir, file), 'r')
        load_dict = json.load(load_f)
        objects = load_dict['shapes']

        for i in range(0, len(objects)):
            label = objects[i]['label']
            if label not in classes:
                classes.append(label)

    return classes


# 获取图片名称
def image_id(rootdir):
    a = []

    # root 表示当前正在访问的文件夹路径
    # dirnames 表示该文件夹下的子目录名list
    # filenames 表示该文件夹下的文件list
    for root, dirnames, filenames in os.walk(rootdir):
        for filename in filenames:
            filename = filename.strip('.json')
            a.append(filename)
    return a


# 获取目标检测Bounding_box
def position(pos):
    # 该函数用来找出xmin,ymin,xmax,ymax即bbox包围框
    x = []
    y = []
    nums = len(pos)
    for i in range(nums):
        x.append(pos[i][0])
        y.append(pos[i][1])
    x_max = max(x)
    x_min = min(x)
    y_max = max(y)
    y_min = min(y)
    b = (float(x_min), float(x_max), float(y_min), float(y_max))
    return b


# 归一化处理
def convert(size, box):
    # 该函数将xmin,ymin,xmax,ymax转为x,y,w,h中心点坐标和宽高
    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(rootdir, out_file, cls_list, image_id):
    # json标签的地址
    load_f = open(rootdir + '/{0}.json'.format(image_id), 'r')
    load_dict = json.load(load_f)

    # 输出标签的地址
    out_file = open(out_file + '/{0}.txt'.format(image_id), 'w')

    #  获取原图的高、宽  用于归一化
    w = load_dict['imageWidth']
    h = load_dict['imageHeight']

    objects = load_dict['shapes']
    nums = len(objects)

    for i in range(0, nums):
        print(image_id + '第{}个'.format(i))
        pos = objects[i]['points']
        box = position(pos)
        bb = convert([w, h], box)

        cls = objects[i]['label']

        if re.match('Support', cls) is not None:
            cls_id = 0
        elif re.match('Bolt_hole', cls) is not None:
            cls_id = 1
        elif re.match('Grouting_hole', cls) is not None:
            cls_id = 2
        elif re.match('Cable', cls) is not None:
            cls_id = 3
        elif re.match('Pipe', cls) is not None:
            cls_id = 4
        elif re.match('Signal_light', cls) is not None:
            cls_id = 6
        elif re.match('Sign', cls) is not None:
            cls_id = 5
        elif re.match('Railway', cls) is not None:
            cls_id = 7
        elif re.match('PJB', cls) is not None:
            cls_id = 8
        elif re.match('Instrument_box', cls) is not None:
            cls_id = 9
        elif re.match('Crack', cls) is not None:
            cls_id = 10
        elif re.match('Falling_block', cls) is not None:
            cls_id = 11
        else:
            cls_id = -1

        out_file.write(str(cls_id) + " " + " ".join([str(round(a, 3)) for a in bb]) + '\n')


if __name__ == '__main__':
    #   json文件输入目录
    rootdir = '../../raw_data/labels-json'

    # 获得所有标注类别
    cls_list = getclass(rootdir)
    print(cls_list)

    #   txt文件输出目录
    out_file = 'labels'
    if not os.path.exists(out_file):
        os.makedirs(out_file)

    #  获取json目录下所有文件名
    ids = image_id(rootdir)

    for id in ids:
        convert_annotation(rootdir, out_file, cls_list, id)
        print(id + '.json' + '已转换')

    print('over!')

2. split.py

det文件夹下新建一个split.py文件，使用下面代码生成ImageSets，里面有一个Main文件夹，Main文件夹里包括test.txt、train.txt、trainval.txt、val.txt四个文本文档

• trainval.txt包含你数据集里所有图像名称

• train.txt为数据集的训练集，为总数据集的90%

• val.txt为数据集的验证集，为总数据集的10%

• test.txt文件里是空的不用担心，因为没有划分测试集

# coding:utf-8

import os
import random
import argparse

parser = argparse.ArgumentParser()
# xml文件的地址，根据自己的数据进行修改xml一般存放在Annotations下
parser.add_argument('--xml_path', default='./images', type=str,
                    help='input xml label path')
# 数据集的划分，地址选择自己数据下的ImageSets/Main
parser.add_argument('--txt_path', default='./ImageSets/Main', type=str,
                    help='output txt label path')
opt = parser.parse_args()

trainval_percent = 1.0
train_percent = 0.9
xmlfilepath = opt.xml_path
txtsavepath = opt.txt_path
total_xml = os.listdir(xmlfilepath)
if not os.path.exists(txtsavepath):
    os.makedirs(txtsavepath)

num = len(total_xml)
list_index = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)

random.seed(0)  # 设置随机种子
trainval = random.sample(list_index, tv)
train = random.sample(trainval, tr)

file_trainval = open(txtsavepath + '/trainval.txt', 'w')
file_test = open(txtsavepath + '/test.txt', 'w')
file_train = open(txtsavepath + '/train.txt', 'w')
file_val = open(txtsavepath + '/val.txt', 'w')

for i in list_index:
    name = total_xml[i][:-4] + '\n'
    if i in trainval:
        file_trainval.write(name)
        if i in train:
            file_train.write(name)
        else:
            file_val.write(name)
    else:
        file_test.write(name)

file_trainval.close()
file_train.close()
file_val.close()
file_test.close()

3. voc_labels.py

更改记录：

1. 添加了abs()函数，以保证转换为txt时都是正数，不然可能运行代码时加载数据会出错
2. out_dir如果不存在则创建

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

sets = ['train', 'val', 'test']
classes = ['Bolt_hole', 'Grouting_hole', 'Crack']  # 改成自己的类别
abs_path = os.getcwd()
print(abs_path)  # E:\A_new_segdet\yolov5ds-main\paper_data\det


def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    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):
    # in_file = open(os.path.join(abs_path, 'Annotations', f'{image_id}.xml'), encoding='UTF-8')
    in_file = open(os.path.join(r'E:\A_new_dataset\A_tunnel_crack\labels_xml', f'{image_id}.xml'), encoding='UTF-8')
    out_dir = os.path.join(abs_path, 'labels')
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)
    out_file = open(os.path.join(out_dir, f'{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
        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(abs(a)) for a in bb]) + '\n')


wd = getcwd()
for image_set in sets:
    if not os.path.exists('./labels'):
        os.makedirs('./labels')
    image_ids = open('./ImageSets/Main/%s.txt' % (image_set)).read().strip().split()
    list_file = open('./%s.txt' % (image_set), 'w')
    for image_id in image_ids:
        list_file.write(abs_path + '/images/%s.png\n' % (image_id))  # 注意这里的自己的图像后缀是png还是jpg或者其他
        convert_annotation(image_id)  # 因为我没用到这个xml转为txt的，就会把这个注释掉
    list_file.close()

运行之后会在det目录下生成train.txt、test.txt、val.txt三个文件，对应的图像名称前加入了绝对路径

运行完以上代码后，det文件夹结构为：

2-2、seg文件夹下

images存放图像

labels存放分割图像，类似下面这种图

（经过下面的segsplit.py代码后，会在images、labels里面分别创建train、val两个子文件夹，这里的train和val里的图要与上述提到的det/Main/train.txt和val.txt里分好的相对应）

注：labels图像应和images对应图像名称相同，格式为png

这种图是用labelme标注后，使用labelme里自带的一个程序生成

• 可以参考我的这篇博文：labelme的安装、批量转换方法（版本号4.5.13）

• 用自己的数据集训练maskrcnn_沙雅云的博客-CSDN博客_maskrcnn训练自己的数据集

通过网上批量转换labelme生成的json代码，生成了如下的文件夹，然后下面的getmask.py就是实现将这些文件夹中的label.png全部提出来放到一起

1. getmask.py

"""
将所有json转成的label.png拷贝到专门的文件夹
"""

import shutil
import os

json_dir = '../../raw_data/labels-json-labelmeout'
mask_dir = 'labels'

if not os.path.exists(mask_dir):
    os.makedirs(mask_dir)

jsonlist = os.listdir(json_dir)

for x in jsonlist:
    # 复制图像并重命名
    file_new_path = shutil.copy(os.path.join(json_dir, x, 'label.png'), os.path.join(mask_dir, x[:-5] + '.png'))
    print(x + '.png' + ' done.')

2. segsplit.py

import os
import shutil

"""
实现：
1. 将train.txt 和 val.txt中指定的原始图像【复制】到seg/images/train子文件或者val子文件
2. 将train.txt 和 val.txt中指定的原始图像对应的标签png图像，移动到seg/labels/train子文件夹或val子文件夹
"""


# 读取txt文本内容
def openreadtxt(file_name):
    data = []
    f = open(file_name, 'r')  # 打开文件
    for row in f.readlines():  # 读取所有行
        tmp_list = row.split('\n')[0]  # 按‘，’切分每行的数据
        data.append(tmp_list)  # 将每行数据插入data中
    f.close()
    return data


if __name__ == '__main__':
    # 指定存放所有的mask文件夹
    labels_dir = r'E:\A_new_dataset\A_tunnel_crack\masks'

    # ---------------在seg/images 和 seg/labels 文件夹下分别创建train、val文件夹---------------- #
    images_train_dir = './images/train'
    images_val_dir = './images.val'
    labels_train_dir = './labels/train'
    labels_val_dir = './labels/val'

    os.makedirs(images_train_dir)
    os.makedirs(images_val_dir)
    os.makedirs(labels_train_dir)
    os.makedirs(labels_val_dir)
    # ------------------------------------------------------------------------------------ #

    train_list = openreadtxt('../det/train.txt')
    val_list = openreadtxt('../det/val.txt')
    print(train_list, len(train_list))
    print(val_list, len(val_list))

    for file_path in train_list:
        file = file_path.split('/')[-1]
        filename, _ = os.path.splitext(file)

        # 复制图像到seg/images/train
        shutil.copy(file_path, os.path.join(images_train_dir, file))
        # 复制labels到seg/labels/train
        labels_name = filename + '.png'
        shutil.copy(os.path.join(labels_dir, labels_name), os.path.join(labels_train_dir, labels_name))

    for file_path in val_list:
        file = file_path.split('/')[-1]
        filename, _ = os.path.splitext(file)

        # 复制图像到seg/images/val
        shutil.copy(file_path, os.path.join(images_val_dir, file))
        # 复制labels到seg/labels/val
        labels_name = filename + '.png'
        shutil.copy(os.path.join(labels_dir, labels_name), os.path.join(labels_val_dir, labels_name))

经过以上两个代码之后的 paper_data/seg 文件夹结构：

到此，检测、分割的数据集都做好了，paper_data数据集总结构为：

3、配置文件参数修改

3-1、models/segheads.yaml

• segnc：改为自己的分割类别数 + 1（这里一定要+1）

3-2、data/voc.yaml

• train：改为自己det文件夹下train.txt路径

• val：改为自己det文件夹下val.txt路径

• road_seg_train：改为自己seg文件夹下images/train文件夹路径

• road_seg_val：改为自己seg文件夹下images/val文件夹路径

• nc：改为自己的检测类别数

• segnc：改为自己的分割类别数（这里一定不要+1 ！！！）

3-3、models/yolov5s.yaml

nc：改为自己的检测类别数

3-4、trainds.py

parse_opt函数下修改对应default里面的内容

关于resume断点训练、继续训练的方法：yolov5ds-断点训练、继续训练（yolov5同样适用）

4、遇到的问题

4-1、运行trainds.py

• 出现报错：RuntimeError: weight tensor should be defined either for all or no class at ...

原因：这个文件中计算分割损失时，没有初始化weight，所以无法设置weight（可能是我自己的torch.nn有问题？）

解决方案：在 trainds.py 中用 Ctrl+F 搜索定位到 SegLoss = nn.CrossEntropyLoss，然后去除掉里面的weight参数，改为以下：

SegLoss = nn.CrossEntropyLoss(ignore_index=255)     # 删掉原本里面的weight
# SegLoss = nn.CrossEntropyLoss(weight=torch.tensor(seg_weights, device=device), ignore_index=255)

4-2、运行detectds.py

• 出现报错：RuntimeError: Input type (torch.cuda.HalfTensor) and weight type (torch.HalfTensor)

原因：输入放在的gpu上，权重却没有放在gpu上，导致数据类型不一致

解决方案：在 detectds.py 中用 Ctrl+F 搜索定位到model = ckpts['model']位置，然后在下方加上一行代码：

model = model.cuda()
# 或者加上下面这行代码
# model.to(device)

• IndexError: index 1 is out of bounds for axis 0 with size 1