本文讲述了如何准备YOLO数据集，包括xml标注文件的转换和数据集的划分，以及YOLOV9模型的预训练、配置和验证过程。
一、YOLO数据集格式分布
在YOLO中，数据集的分布如图，在dataset文件夹下有imags（图片）和labels（标签）。在images和labels文件夹下又分别存放三个文件夹，分别对应测试集、训练集、验证集：
在images文件夹下面存放的是图片，在label文件夹下面存放的是TXT文件，每一个txt文件里面都是标注的信息。以下图为例，文件里面有三行，代表这张图片在标注的时候画了三个框。其中每一行数值代表意思是：类别、物体中心点坐标（x1,y1）,所画框的宽高（x2,y2）
二、数据集制作
首先安装labelimg库，然后使用labelimg进行标注，标注好之后会生成xml文件（不详细叙述）。然后就是读取xml文件来获取标注的信息，再将这些信息存入txt文件中。
2024/12/11 更新数据集制作代码：
把标注得到的xml文件放到一个文件夹中，修改代码对应的路径，然后直接运行就可以生成dataset文件夹了： 太多人评论把所有图片都当做训练集不好，于是我就改回来了，80%做训练集。自己根据需求改吧。
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 # 把所有的数据都当做训练集
train_txt_list = txt_lists[:int(len(txt_lists) * 0.8)]
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文件了
扫一扫
3万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
