pycharm 中使用YOLOV8训练自己的数据集

YOLOV8训练自己的数据集

1.下载 yolov8 代码到本地

命令： git clone https://github.com/ultralytics/ultralytics.git

2.数据集准备

自己的数据集是VOC2007格式的，怕麻烦还是写了个代码转为yolo标签格式。
这是VOC格式数据集格式

首先是生成Main文件夹下的txt文件

import os
import random

# 随机数种子
random.seed(0)

#   指向VOC数据集所在的文件夹
VOCdevkit_path = './VOCdevkit'
VOCdevkit_sets = [('2007', 'train'), ('2007', 'val')]

#   trainval_percent用于指定(训练集+验证集)与测试集的比例，默认情况下 (训练集+验证集):测试集 = 9:1
#   train_percent用于指定(训练集+验证集)中训练集与验证集的比例，默认情况下 训练集:验证集 = 9:1
trainval_percent = 0.9
train_percent = 0.9

print("Generate txt in ImageSets.")
xmlfilepath = os.path.join(VOCdevkit_path, 'VOC2007/Annotations')
saveBasePath = os.path.join(VOCdevkit_path, 'VOC2007/ImageSets/Main')
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("train size", 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 = total_xml[i][:-4] + '\n'
    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()
print("Generate txt in ImageSets done.")

然后使用以下代码转换为在yolov8中跑的数据集格式

'''
将XML格式的数据集转换成YOLO TXT格式
'''

import os
import shutil
import xml.etree.ElementTree as ET

def convert_annotation(xml_paths, classes):

    for xml_path in xml_paths:
        txt_path = xml_path.replace("Annotations", "labels").replace(".xml", ".txt")
        txtfile = open(txt_path, 'w')

        in_file = open(xml_path, encoding="utf-8")
        tree = ET.parse(in_file)
        root = tree.getroot()
        size = root.find('size')
        width = int(size.find('width').text)
        height = int(size.find('height').text)

        for obj in root.iter("object"):
            classname = obj.find("name").text
            if classname not in classes:
                continue
            cls_id = classes.index(classname)

            xmin = (float(obj.find("bndbox").find("xmin").text))
            ymin = (float(obj.find("bndbox").find("ymin").text))
            xmax = (float(obj.find("bndbox").find("xmax").text))
            ymax = (float(obj.find("bndbox").find("ymax").text))

            x_center = (xmin + xmax) / 2.0
            y_center = (ymin + ymax) / 2.0
            x = x_center / width
            y = y_center / height
            w = (xmax - xmin) / width
            h = (ymax - ymin) / height
            datas = (x, y, w, h)
            txtfile.write(str(cls_id) + " " + " ".join([str(data) for data in datas]) + '\n')
        in_file.close()
        txtfile.close()

if __name__ == "__main__":
    # VOC数据集根目录
    vocDir = "./VOCdevkit"

    # 创建YOLO数据集标签文件夹
    yololabel = os.path.join(vocDir, "VOC2007", "labels")
    yoloimage = os.path.join(vocDir, "VOC2007", "images")
    if not os.path.exists(yololabel):
        os.mkdir(yololabel)
    if not os.path.exists(yoloimage):
        os.mkdir(yoloimage)

    # 数据集类别 我这里是单类别目标检测
    classes = ['wheat']

    # 按照训练集验证集测试集将标签分别转换成yolo txt格式
    VOCdevkit_sets = [('2007', 'train'), ('2007', 'val'), ('2007', 'test')]
    for year, image_set in VOCdevkit_sets:
        image_ids = open(os.path.join(vocDir, 'VOC%s/ImageSets/Main/%s.txt' % (year, image_set)),
                         encoding='utf-8').read().strip().split()
        yolofile = open(os.path.join(vocDir, 'VOC%s/%s.txt' % (year, image_set)),"w", encoding='utf-8')
        xml_list = []
        for image_id in image_ids:
            xmlpath = os.path.join(vocDir, 'VOC%s/Annotations/%s.xml' % (year, image_id))
            xml_list.append(xmlpath)

            imagepath = os.path.join(vocDir, 'VOC%s/JPEGImages/%s.jpg' % (year, image_id))
            newimagepath = imagepath.replace("JPEGImages", "images")
            shutil.copy(imagepath, newimagepath)
            path = os.path.abspath(newimagepath)
            yolofile.write(path)
            yolofile.write("\n")
        yolofile.close()

        # 将xml标签转为txt
        convert_annotation(xml_list, classes)
        print("-------------------------------")
        print("process file number:", len(xml_list))

运行完会生成以下文件，组成了一个可以在yolov8里使用的数据集

训练

yolov8训练参数都集成到一个文件中
修改ultralytics/yolo/v8/detect/train.py文件，修改data路径

修改default.yaml文件参数，先修改模型文件和数据集文件

修改ultralytics/models/v8/yolov8.yaml 的模型文件nc为自己数据集的类别数

在ultralytics/datasets/mydataset.yaml，新建一个mydataset.yaml文件，参考coco.yaml文件，path填写数据集train.txt文件绝对路径，类别名称修改为自己的数据集

然后修改default.yaml的其他训练参数，自己看着想改啥改啥，我就修改了这几项
训练文件保存在ultralytics/yolo/v8/detect/rundetect/log文件夹下

验证

修改ultralytics/yolo/v8/detect/val.py文件下data路径

在default.yaml文件下修改mode为val，修改model为训练获得的best.pt路径，这里写绝对路径，然后修改下val的ion和conf参数

因为不是命令行运行代码要修改ultralytics/yolo/engine/validator.py文件，解析参数，然后直接运行val.py文件

推理

修改default.yaml文件mode为predict，修改model为训练的权重文件路径，修改iou和conf值，修改source路径，然后直接运行predict.py文件

错误处理

直接 pip uninstall wandb 卸载wandb库，不使用这个库，对我来说没什么用。

##结束
这里全部都是直接在pycharm中运行的，和命令行运行代码不一样，可以参考，但不一定完全适合。