当前位置：首页 > news >正文

yolov5自定义模型训练二

news 2025/7/23 14:14:28

前期准备好了用于训练识别是否有火灾的数据集后就可以开始修改yolo相关文件来进行训练

数据集放到yolov5目录里

在data目录下新建yaml文件设置数据集信息如下

在model文件夹下新增新的model文件

开始训练

训练出错

确认后是对训练数据集文件夹里的文件名字有要求，原先

改为如下后

仍然出错如下，看log应该时转换lable时的问题

确认了时label没做归一化处理，新增归一化处理的转换脚本代码，

import os
import xml.etree.ElementTree as ETdef convert_xml_to_txt(xml_path, txt_path):# 读取XML文件tree = ET.parse(xml_path)root = tree.getroot()# 获取图像的宽度和高度width = int(root.find('size/width').text)height = int(root.find('size/height').text)with open(txt_path, 'w') as f:# 遍历所有的目标for obj in root.findall('object'):# 获取类别名称class_name = obj.find('name').text# 获取边界框坐标xmin = int(obj.find('bndbox/xmin').text)ymin = int(obj.find('bndbox/ymin').text)xmax = int(obj.find('bndbox/xmax').text)ymax = int(obj.find('bndbox/ymax').text)# 归一化处理x_center = (xmin + xmax) / (2 * width)y_center = (ymin + ymax) / (2 * height)box_width = (xmax - xmin) / widthbox_height = (ymax - ymin) / height# 写入到TXT文件class_index = class_names.index(class_name)f.write(f"{class_index} {x_center} {y_center} {box_width} {box_height}\n")# 读取类别名称列表
class_names = ['fire']# 遍历xml文件夹中的所有xml文件
xml_folder = 'D:/fire-dataset/fire-dataset/validation/annotations'
txt_folder = 'D:/fire-dataset/fire-dataset/validation/txtannotations'
for xml_file in os.listdir(xml_folder):if xml_file.endswith('.xml'):xml_path = os.path.join(xml_folder, xml_file)txt_file = os.path.splitext(xml_file)[0] + '.txt'txt_path = os.path.join(txt_folder, txt_file)convert_xml_to_txt(xml_path, txt_path)

重新转换后重新训练