一、准备前的文件夹目录介绍
bag-images文件夹:用来存放原始数据集所有的.jpg图片
xml文件夹:用来存放原始数据集打过标签的所有xml文件
txt文件夹:用来存放原始数据集,由xml格式转换为txt格式的所有文件
bag文件夹:是我们目标制作的数据集,用于后期跑实验
bag文件夹下有images和labels文件夹,每个文件夹下都有一个train和val文件夹
images文件夹:用来存放目标数据集的所有图片,分为train训练和val验证两部分
labels文件夹:用来存放目标数据集的所有对应的txt文件,分为train训练和val验证两部分
二、首先将自己标注完成的数据集进行格式转换
下面是将文件由xml格式转换为txt格式代码,建一个voc_label.py脚本,将下面代码复制进去,运行,生成(注意该文件路径地址):
import xml.etree.ElementTree as ET
import osclasses = ["bag"] # 类别,改成自己的类别名称CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))def convert(size, box):dw = 1. / size[0]dh = 1. / size[1]x = (box[0] + box[1]) / 2.0y = (box[2] + box[3]) / 2.0w = box[1] - box[0]h = box[3] - box[2]x = x * dww = w * dwy = y * dhh = h * dhreturn (x, y, w, h)def convert_annotation(image_id):in_file = open(r'E:\人包物\bag\xml/%s.xml' % (image_id), encoding='UTF-8')#E:\人包物\bag\xml 为xml文件地址out_file = open(r'E:\人包物\bag\txt/%s.txt' % (image_id), 'w') # 生成txt格式文件 #'E:\人包物\bag\txt 为将要输出生成的txt文件地址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# print(cls)if cls not in classes:continuecls_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))bb = convert((w, h), b)out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')xml_path = os.path.join(CURRENT_DIR, r'E:\人包物\bag\xml/') #E:\人包物\bag\xml 为xml文件地址# xml list
img_xmls = os.listdir(xml_path)
for img_xml in img_xmls:label_name = img_xml.split('.')[0]print(label_name)convert_annotation(label_name)
三、将处理好的图片和txt文本进行train和val数据集划分
import os # 用于处理文件路径、创建目录等操作
import random # 用于生成随机数种子、打乱列表等操作
import shutil # 用于生成随机数种子、打乱列表等操作
# 设置随机数种子
random.seed(123)
# 定义文件夹路径
image_dir = r'E:\人包物\images' # 原始图像所在的子目录
label_dir = r'E:\人包物\labels' # 原始标签所在的子目录
output_dir = r'E:\人包物\bag' # 处理后的数据集输出目录
# 定义训练集、验证集和测试集比例
train_ratio = 0.8 # 训练集比例
valid_ratio = 0.1 # 验证集比例
test_ratio = 0.1 # 测试集比例
# 获取所有图像文件和标签文件的文件名(不包括文件扩展名)
image_filenames = [os.path.splitext(f)[0] for f in os.listdir(image_dir)] # 提取所有图像文件的文件名列表
label_filenames = [os.path.splitext(f)[0] for f in os.listdir(label_dir)] # 提取所有标签文件的文件名列表
# 随机打乱文件名列表
random.shuffle(image_filenames) # 打乱图像文件的文件名列表
# 计算训练集、验证集和测试集的数量
total_count = len(image_filenames) # 总文件数
train_count = int(total_count * train_ratio) # 训练集文件数
valid_count = int(total_count * valid_ratio) # 验证集文件数
test_count = total_count - train_count - valid_count # 测试集文件数
# 定义输出文件夹路径
train_image_dir = os.path.join(output_dir, 'train', 'images') # 训练集图像输出目录
train_label_dir = os.path.join(output_dir, 'train', 'labels') # 训练集标签输出目录
valid_image_dir = os.path.join(output_dir, 'valid', 'images') # 验证集图像输出目录
valid_label_dir = os.path.join(output_dir, 'valid', 'labels') # 验证集标签输出目录
test_image_dir = os.path.join(output_dir, 'test', 'images') # 测试集图像输出目录
test_label_dir = os.path.join(output_dir, 'test', 'labels') # 测试集标签输出目录
# 创建输出文件夹
os.makedirs(train_image_dir, exist_ok=True) # 创建训练集图像输出目录
os.makedirs(train_label_dir, exist_ok=True) # 创建训练集标签输出目录
os.makedirs(valid_image_dir, exist_ok=True) # 创建验证集图像输出目录
os.makedirs(valid_label_dir, exist_ok=True) # 创建验证集标签输出目录
os.makedirs(test_image_dir, exist_ok=True) # 创建测试集图像输出目录
os.makedirs(test_label_dir, exist_ok=True) # 创建测试集标签输出目录
# 将图像和标签文件划分到不同的数据集中
for i, filename in enumerate(image_filenames):# 如果文件数量小于训练数据集大小,则将文件复制到训练数据集目录中if i < train_count:output_image_dir = train_image_diroutput_label_dir = train_label_dir# 如果文件数量小于训练数据集大小+验证数据集大小,则将文件复制到验证数据集目录中elif i < train_count + valid_count:output_image_dir = valid_image_diroutput_label_dir = valid_label_dir# 否则,将文件复制到测试数据集目录中else:output_image_dir = test_image_diroutput_label_dir = test_label_dir# 复制图像文件src_image_path = os.path.join(image_dir, filename + '.jpg') # 获取图像文件的源路径dst_image_path = os.path.join(output_image_dir, filename + '.jpg') # 获取图像文件的目标路径shutil.copy(src_image_path, dst_image_path) # 复制图像文件到目标路径# 复制标签文件src_label_path = os.path.join(label_dir, filename + '.txt') # 获取标签文件的源路径dst_label_path = os.path.join(output_label_dir, filename + '.txt') # 获取标签文件的目标路径shutil.copy(src_label_path, dst_label_path) # 复制标签文件到目标路径
划分好之后就是下面这个样子:
将整个bag文件夹,复制移动到D:\cs\yolov8\ultralytics\datasets文件夹下,(一定要在此文件夹)
在D:\cs\yolov8\ultralytics\datasets文件夹下新建一个bag.yaml文件
# Ultralytics YOLO 🚀, AGPL-3.0 license
# COCO128 dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics
# Example usage: yolo train data=person.yaml
# parent
# ├── ultralytics
# └── datasets
# └── coco128 ← downloads here (7 MB)# 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: D:/cs/yolov8/ultralytics/datasets/bag # dataset root dir
train: images/train # train images (relative to 'path') 128 images
val: images/val # val images (relative to 'path') 128 images
test: # test images (optional)# Classes
names:0: bag
下载与训练模型
https://github.com/ultralytics/assets/releases
我们选择yolov8n.pt,权重下载好放在根目录下
接下来在根目录创建一个train.py脚本文件,将下面代码复制过去,进行训练
from ultralytics import YOLO# 加载模型
# model = YOLO("yolov8n.yaml") # 从头开始构建新模型
model = YOLO("yolov8n.pt") # 加载预训练模型(推荐用于训练)# Use the model
results = model.train(data=r"D:\cs\yolov8\ultralytics\datasets\bag.yaml", epochs=200, batch=16, ) # 训练模型
遇到的问题,若训练过程中,遇到报错:
RuntimeError: CUDA error: CUBLAS_STATUS_NOT_INITIALIZED when calling `cublasCreate(handle)`
找到D:\cs\yolov8\ultralytics\ultralytics\cfg\default.yaml路径下的default.yaml文件,将works改为0,windows下容易报这种错误
成功: