前言
首先需要创建适合yolov8的数据模式,需要将xml文件转成txt文件。修改yolov8的配置文件实现模型的训练
提示:以下是本篇文章正文内容,下面案例可供参考
一、使用步骤
1.xml文件转成txt文件
代码如下(示例):
import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
classes = [自己的数据集标签]
def convert(size, box):
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
if x >= 1:
x = 0.99
if y >= 1:
y = 0.99
if w>=1:
w=0.99
if h>=1:
h=0.99
return (x,y,w,h)
def convert_annotation(rootpath,xmlname):
xmlpath = rootpath + '/xmls'
xmlfile = os.path.join(xmlpath,xmlname)
with open(xmlfile, "r", encoding='UTF-8') as in_file:
txtname = xmlname[:-4]+'.txt'
print(txtname)
txtpath = rootpath + '/worktxt'
if not os.path.exists(txtpath):
os.makedirs(txtpath)
txtfile = os.path.join(txtpath,txtname)
with open(txtfile, "w+" ,encoding='UTF-8') as out_file:
tree=ET.parse(in_file)
root = tree.getroot()
size = root.find('size')
w = int(size.find('width').text)
h = int(size.find('height').text)
out_file.truncate()
for obj in root.iter('object'):
#difficult = obj.find('difficult').text
cls = obj.find('name').text
#if cls not in classes or int(difficult)==1:
#continue
if cls not in classes:
print("image dont excistence")
continue
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))
bb = convert((w,h), b)
out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
if __name__ == "__main__":
rootpath='自己数据集路径'
xmlpath=rootpath+'/xmls'
list=os.listdir(xmlpath)
for i in range(0,len(list)) :
path = os.path.join(xmlpath,list[i])
if ('.xml' in path)or('.XML' in path):
convert_annotation(rootpath,list[i])
print('done', i)
else:
print('not xml file',i)
2.加入数据中
如下(示例):
2.加入数据中
1:先修改ultralytics/yolo/cfg/default.yaml中的default.yaml文件,如果是任务是目标检测则tast:改成detect 。如是语义分割则改成segment
2:如果是训练模式,mode改成train模式。
3:model改成初始模型的绝对路径,
4:data改成自己创建的数据集。这样设置完成后就可以了。接下来需要运行训练脚本
3:模型训练
1:找到项目文件下的train.py文件 路径为:ultralytics/yolo/v8/detect/train.py。
2:直接运行train文件既可实现训练自己的数据集。
如果想自己写训练脚本,如下脚本为训练脚本
from ultralytics import YOLO
import torch
if __name__ == '__main__':
# 加载模型
model = YOLO("C:/Users/12057/Desktop/yolov8/ultralytics/ultralytics/models/v8/yolov8.yaml") # 从头开始构建新模型
# model = YOLO("yolov8s.pt") # 加载预训练模型(推荐用于训练)
# Use the model
results = model.train(data="C:/Users/12057/Desktop/yolov8/ultralytics/ultralytics/datasets/mydata2.yaml", epochs=300, batch=16, workers=8, close_mosaic=0, name='cfg') # 训练模型
# results = model.val() # 在验证集上评估模型性能
# results = model("https://ultralytics.com/images/bus.jpg") # 预测图像
# success = model.export(format="onnx") # 将模型导出为 ONNX 格式
4:训练过程
总结
以上就是今天要讲的内容,本文仅仅简单介绍了利用yolov8训练自己的数据集,而后期换成onnx格式,将在下期进行讲解。