YOLOV9训练集制作+Train+Val记录

一、YOLO数据集格式分布

        在YOLO中，数据集的分布如图，在dataset文件夹下有imags（图片）和labels（标签）。在images和labels文件夹下又分别存放三个文件夹，分别对应测试集、训练集、验证集：

        在images文件夹下面存放的是图片，在label文件夹下面存放的是TXT文件，每一个txt文件里面都是标注的信息。以下图为例，文件里面有三行，代表这张图片在标注的时候画了三个框。其中每一行数值代表意思是：类别、物体中心点坐标（x1,y1）,所画框的宽高（x2,y2）

二、数据集制作

        首先安装labelimg库，然后使用labelimg进行标注，标注好之后会生成xml文件（不详细叙述）。然后就是读取xml文件来获取标注的信息，再将这些信息存入txt文件中。

2024/05/27更新数据集制作代码：

        把标注得到的xml文件放到一个文件夹中，修改代码对应的路径，然后直接运行就可以生成dataset文件夹了：

import os
import random
import shutil
import numpy as np
import xml.etree.ElementTree as ET
import cv2


def conver_annotation(xml_folder, xml_name, class_list, labels_folder, img_path):
    xml_path = os.path.join(xml_folder, xml_name)
    print("xml_path:", xml_path)
    in_file = open(xml_path, encoding='utf-8')
    tree = ET.parse(in_file)
    root = tree.getroot()
    # for SIZE in root.iter('size'):
    #     width = int(SIZE.find('width').text)
    #     height = int(SIZE.find('height').text)
    img = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), -1)
    height, width, _ = img.shape
    label_name = xml_name.split(".")[0] + ".txt"
    label_path = os.path.join(labels_folder, label_name)
    if os.path.exists(label_path):
        os.remove(label_path)
    with open(label_path, "a", encoding="utf-8") as f:
        for obj in root.iter('object'):
            difficult = 0
            if obj.find('difficult') != None:
                difficult = obj.find('difficult').text
            cls = obj.find('name').text
            if cls not in class_list or int(difficult) == 1:
                continue
            cls_id = class_list.index(cls)
            xmlbox = obj.find('bndbox')
            # x1, y1, x2, y2 = (float(xmlbox.find('xmin').text)/width, float(xmlbox.find('ymin').text)/height,
            #      float(xmlbox.find('xmax').text)/width, float(xmlbox.find('ymax').text)/height)

            x1, y1, x2, y2 = (float(xmlbox.find('xmin').text), float(xmlbox.find('ymin').text),
                              float(xmlbox.find('xmax').text), float(xmlbox.find('ymax').text))
            # print("cls:", cls_id)
            # print("box:", (x1 + x2) / 2 / width, (y1 + y2) / 2 / height, (x2 - x1) / width, (y2 - y1) / height)
            x_c, y_c, w, h = (x1 + x2) / 2 / width, (y1 + y2) / 2 / height, (x2 - x1) / width, (y2 - y1) / height
            if x_c >1 or y_c > 1 or w > 1 or h > 1:
                print("坐标信息错误")
                break

            # f.write("%s %s %s %s %s\n" % (cls_id, x1, y1, x2, y2))
            f.write("%s %s %s %s %s\n" % (cls_id, x_c, y_c, w, h))


def gen_txtfile(xml_folder, image_folder, save_label_folder, class_name):
    for xml_name in os.listdir(xml_folder):
        if xml_name.endswith(".xml"):
            img_name = xml_name.split(".")[0] + ".jpg"
            img_path = os.path.join(image_folder, img_name)
            conver_annotation(xml_folder, xml_name, class_name, save_label_folder, img_path)


def split_dataset(txt_folder, img_folder, dataset_folder):

    txt_lists = os.listdir(txt_folder)
    random.shuffle(txt_lists)

    train_txt_list = txt_lists      #  把所有的数据都当做训练集
    val_txt_list = txt_lists[int(len(txt_lists) * 0.8): int(len(txt_lists) * 0.9)]   # 10%作为验证集
    test_txt_list = txt_lists[int(len(txt_lists) * 0.9):]   # 10%作为测试集

    # 创建文件夹
    train_img_folder = dataset_folder + "/images/train"
    val_img_folder = dataset_folder + "/images/val"
    test_img_folder = dataset_folder + "/images/test"
    train_label_folder = dataset_folder + "/labels/train"
    val_label_folder = dataset_folder + "/labels/val"
    test_label_folder = dataset_folder + "/labels/test"
    for folder in [train_img_folder, val_img_folder, test_img_folder, train_label_folder, val_label_folder, test_label_folder]:
        if not os.path.exists(folder):
            os.makedirs(folder)

    def copy_files(one_txt_list, output_img_folder, output_label_folder):
        for txt_name in one_txt_list:
            if txt_name.endswith(".txt"):
                # 原始文件路径
                txt_path = os.path.join(txt_folder, txt_name)
                img_name = txt_name.split(".")[0] + ".jpg"
                img_path = os.path.join(img_folder, img_name)

                # 新保存文件路径
                new_txt_path = os.path.join(output_label_folder, txt_name)
                new_img_path = os.path.join(output_img_folder, img_name)
                shutil.copy(txt_path, new_txt_path)
                shutil.copy(img_path, new_img_path)
                print("img_path:", img_path)

    copy_files(train_txt_list, train_img_folder, train_label_folder)
    copy_files(val_txt_list, val_img_folder, val_label_folder)
    copy_files(test_txt_list, test_img_folder, test_label_folder)

if __name__ == '__main__':
    # xml所在的文件夹
    xml_folder = r""
    # 图片所在的文件夹
    image_folder = r""
    # 保存txt文件的文件夹路径
    save_label_folder = r""
    # 如果文件夹不存在，则创建文件夹
    if not os.path.exists(save_label_folder):
        os.makedirs(save_label_folder)
    # 类别名称
    class_name = [
        "num_0", "num_1", "num_2", "num_3", "num_4", "num_5", "num_6",
        "num_7", "num_8", "num_9", "charB_A", "charB_B", "charB_C", "charB_D", "charB_E", "charB_F", "charB_G",
        "charB_H", "charB_I", "charB_J", "charB_K", "charB_L", "charB_M", "charB_N", "charB_P", "charB_Q", "charB_R", "charB_S", "charB_T",
        "charB_U", "charB_V", "charB_W", "charB_X", "charB_Y", "charB_Z", "word_he", "word_ge"
    ]
    # 生成所有文件的txt
    gen_txtfile(xml_folder, image_folder, save_label_folder, class_name)

    # 训练数据的保存路劲，dataset
    dataset_save_folder = r"XXXXXXXXXXXX\dataset"
    if not os.path.exists(dataset_save_folder):
        os.makedirs(dataset_save_folder)
    # 分割数据，并生成训练的dataset格式
    split_dataset(save_label_folder, image_folder, dataset_save_folder)

三、训练部分

YOLOV9代码地址：https://github.com/WongKinYiu/yolov9/tree/main

将代码克隆到本地之后，在github下滑页面会看到下载预训练模型，目前只有yolov9-c何yolov9-e可以下载，以yolov9-c为例：

        下载完预训练模型之后将预训练模型放到工程目录下（也就是yolov9下）。然后再data文件夹下创建一个data.yaml文件。yaml文件里面分别是训练数据路径、类别数量、类别名称（类别名称要和标注时的顺序一致）

弄完之后，打开train.py文件，在train.py文件下面加上一行，防止报错。

os.environ['KMP_DUPLICATE_LIB_OK']='TRUE'

然后下滑找到参数配置文件，修改配置参数。参数介绍：           

                weights：下载的预训练模型

                cfg:对应模型的参数配置文件

                data:前面创建的data.yaml文件

显存比较小的可以把batchsize设小一点，这里设置1，电脑CPU不好的可以把wokers设为0

配置参数搞完之后，再修改一下代码。打开utils文件夹下面的loss_tal.py文件，修改如下参数，不修改运行会报错。修改完之后就可以直接运行train.py文件了。

0530更新：在训练完训练集之后，会进行验证。在val的时候会报错。

解决方案：看官方项目的评论区里也提到相关的问题，试了一下发现可以解决。但具体原因还没去深入了解， 解决方案如下：

四、验证部分

训练完之后模型以及训练的其他输出结果文件会保存在runs/weights里。打开val.py文件，

在import os后面加上

os.environ['KMP_DUPLICATE_LIB_OK']='TRUE'

然后开始修改配置参数：

修改完之后直接运行会报错，代码有个小bug需要修改。打开utils文件夹下的general.py文件第903行修改。然后就可以直接运行val.py文件了