yolov5自定义模型训练一

要拿yolo来用到自己希望使用的场合就需要自己来训练模型权重，要训练模型权重需要足够的训练数据集，数据越多训练后的模型权重越准确。

数据集需要按yolov5的格式进行收集，收集图片后需要用标注工具进行标注，常用的有LabelImg软件或者一些更先进的有部分带ai识别的标注工具

数据集标注工作是个体力活，目前一般有专门做这方面的标记的公司。

 网上找了现成的用于识别火的数据集，不过开源数据集的标注文件默认不是yolo支持的txt文件，因此需要将xml转换到txt文件。

用脚本做个转换，脚本转换代码如下

import os
import xml.etree.ElementTree as ET

def convert_xml_to_txt(xml_path, txt_path):
    tree = ET.parse(xml_path)
    root = tree.getroot()

    with open(txt_path, 'w') as f:
        for obj in root.findall('object'):
            class_name = obj.find('name').text
            bbox = obj.find('bndbox')
            xmin = float(bbox.find('xmin').text)
            ymin = float(bbox.find('ymin').text)
            xmax = float(bbox.find('xmax').text)
            ymax = float(bbox.find('ymax').text)

            width = xmax - xmin
            height = ymax - ymin
            x_center = xmin + width / 2
            y_center = ymin + height / 2

            class_index = class_names.index(class_name)
            f.write(f"{class_index} {x_center} {y_center} {width} {height}\n")


class_names = ['fire']


xml_folder = 'D:/fire-dataset/fire-dataset/train/annotations'
txt_folder = 'D:/fire-dataset/fire-dataset/train/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)

脚本转换后转换为需要的格式文件就可以开始准备进行训练，不过训练先需要按要训练的数据进行修改