本文共 4151 字，大约阅读时间需要 13 分钟。

制作数据集与训练YOLO模型（PASCAL VOC格式）

制作数据集

在开始训练YOLO模型之前，需要先制作数据集。这可以通过简单的Python脚本实现，以下是详细的步骤说明：

1. 确定文件路径

xmlfilepath = 'Annotations'
txtsavepath = 'ImageSets'
total_xml = os.listdir(xmlfilepath)
num = len(total_xml)

2. 划分训练集、验证集和测试集

trainval_percent = 0.9  # 训练+验证所占的比例
train_percent = 0.9  # 训练集所占的比例

3. 随机划分数据集

list = range(num)
tv = int(num * trainval_percent)  # 训练+验证的样本数
tr = int(tv * train_percent)      # 训练集的样本数
trainval = random.sample(list, tv)  # 训练+验证集
train = random.sample(trainval, tr)  # 训练集

4. 保存数据集标签文件

ftrainval = open(txtsavepath + '/trainval.txt', 'w')
ftrain = open(txtsavepath + '/train.txt', 'w')
fval = open(txtsavepath + '/val.txt', 'w')
fetest = open(txtsavepath + '/test.txt', 'w')
for i in list:
    name = total_xml[i][:-4] + '\n' 
    if i in trainval:
        ftrainval.write(name)
        if i in train:
            ftrain.write(name)
        else:
            fval.write(name)
    else:
        ftest.write(name)

5. 转换格式

from xml.etree.ElementTree import ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
sets = ['train', 'test', 'val']
classes = ['crazing', 'inclusion', 'patches', 'pitted_surface', 'rolled-in_scale', 'scratches']
def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)
def convert_annotation(image_id):
    in_file = open(f'Dataset/Annotations/{image_id}.xml', 'r')
    out_file = open(f'Dataset/labels/{image_id}.txt', 'w')
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    width = int(size.find('width').text)
    height = int(size.find('height').text)
    
    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
            
        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((width, height), b)
        out_file.write(f'{cls_id} {bb[0]} {bb[1]} {bb[2]} {bb[3]}\n')
    
    in_file.close()
    out_file.close()
# 创建所需目录
if not os.path.exists('Dataset/labels/'):
    os.makedirs('Dataset/labels/')
for image_set in sets:
    if not os.path.exists(f'Dataset/labels/'):
        os.makedirs('Dataset/labels/')
    image_ids = open(f'Dataset/ImageSets/{image_set}.txt', 'r').read().split()
    list_file = open(f'Dataset/{image_set}.txt', 'w')
    
    for image_id in image_ids:
        list_file.write(f'Dataset/images/{image_id}.jpg\n')
        convert_annotation(image_id)
    
    list_file.close()

6. 运行结果

这一步操作完成后，可以看到 trainval.txt、train.txt、val.txt 和 test.txt 文件分别包含了对应的图像名称和标签。

修改配置文件

在YOLO模型的训练配置文件中，主要需要注明以下信息：

示例配置文件内容

# 模型参数
model_path = ' models/yolov3-tiny.cfg'   # 模型配置文件路径
weights_path = ' models/yolov3-tiny.pth'  # 预训练权重文件路径
batch_size = 8                      # 批量大小
学习率 = 0.001                 # 初始学习率
stu_difficult = 1                # 为 True 表示启用困难度标记（建议值）
# 数据参数
data_path = 'Dataset/'               # 数据集路径
split_file = 'train.txt'           # 训练集标签文件
label_classes = 5                   # 类别数量
batch_empty = False               # True 表示在空数据集上训练
# 训练参数
num_epochs = 200                  # 训练 epoch 数量
optimizer = 'ADEAGrad'           # 优化器选择
learning_momentum = 0.9          # 优化器的动量参数
损失项 = ['DiceLoss']            # 损失函数
损失项权重 = [1.0]             # 损失函数权重
批次归一化 = True               # True 表示使用批次归一化
# 输出参数
model_output = 'yolov3-tiny'      # 输出模型名
tensorboard_log = True           # True 表示启用 tensorboard 日志
结果保存路径 = 'logs/'          # 日志和检查点保存路径
检查点间隔 = 1000              # 间隔保存检查点
结果显 Print()间隔 = 100       # 每隔多少步打印训练损失

超参数设置

在YOLO模型的训练中，超参数的选择对训练效果至关重要。建议从以下配置开始：

• epoch：200-300次训练
• batch-size：8-16
• 学习率：0.001-0.01

可以根据训练效果逐步调整学习率和批量大小，确保模型收敛。

以上就是完整的数据集制作与YOLO模型训练配置说明。如果需要进一步的具体实施步骤，欢迎在评论区留言。