YOLOv5车牌识别实战教程(二)理论基础

dataset/
    images/
        train/
        val/
        test/
    labels/
        train/
        val/
        test/

import os
import random
import shutil
 
random.seed(42)
 
def split_data(dataset_path, train_ratio, val_ratio):
    images_path = os.path.join(dataset_path, 'images')
    labels_path = os.path.join(dataset_path, 'labels')
    
    train_path = os.path.join(images_path, 'train')
    val_path = os.path.join(images_path, 'val')
    test_path = os.path.join(images_path, 'test')
 
    os.makedirs(train_path, exist_ok=True)
    os.makedirs(val_path, exist_ok=True)
    os.makedirs(test_path, exist_ok=True)
 
    image_files = [f for f in os.listdir(images_path) if f.endswith('.jpg')]
 
    random.shuffle(image_files)
 
    num_train = int(len(image_files) * train_ratio)
    num_val = int(len(image_files) * val_ratio)
 
    train_files = image_files[:num_train]
    val_files = image_files[num_train:num_train + num_val]
    test_files = image_files[num_train + num_val:]
 
    for file in train_files:
        shutil.move(os.path.join(images_path, file), os.path.join(train_path, file))
        shutil.move(os.path.join(labels_path, file.replace('.jpg', '.txt')), os.path.join(labels_path, 'train', file.replace('.jpg', '.txt')))
 
    for file in val_files:
        shutil.move(os.path.join(images_path, file), os.path.join(val_path, file))
        shutil.move(os.path.join(labels_path, file.replace('.jpg', '.txt')), os.path.join(labels_path, 'val', file.replace('.jpg', '.txt')))
 
    for file in test_files:
        shutil.move(os.path.join(images_path, file), os.path.join(test_path, file))
        shutil.move(os.path.join(labels_path, file.replace('.jpg', '.txt')), os.path.join(labels_path, 'test', file.replace('.jpg', '.txt')))
 
dataset_path = 'dataset'
train_ratio = 0.8
val_ratio = 0.1
split_data(dataset_path, train_ratio, val_ratio)

import os
import xml.etree.ElementTree as ET
 
def voc_to_yolo(xml_file, img_width, img_height):
    tree = ET.parse(xml_file)
    root = tree.getroot()
 
    yolo_annots = []
 
    for obj in root.findall('object'):
        class_name = obj.find('name').text
        class_id = class_name_to_id(class_name)  # 自定义函数，将类别名称转换为对应的ID
        bbox = obj.find('bndbox')
 
        xmin = int(bbox.find('xmin').text)
        ymin = int(bbox.find('ymin').text)
        xmax = int(bbox.find('xmax').text)
        ymax = int(bbox.find('ymax').text)
 
        x_center = (xmin + xmax) / 2 / img_width
        y_center = (ymin + ymax) / 2 / img_height
        width = (xmax - xmin) / img_width
        height = (ymax - ymin) / img_height
 
        yolo_annots.append(f"{class_id} {x_center} {y_center} {width} {height}")
 
    return yolo_annots
 
# 示例：转换一个XML文件，并保存为YOLO格式的TXT文件
xml_file = 'example.xml'
img_width = 640
img_height = 480
yolo_annots = voc_to_yolo(xml_file, img_width, img_height)
 
with open('example.txt', 'w') as f:
    for annot in yolo_annots:
        f.write(annot + '\n')

# 数据增强设置
train:
  ...
  mosaic: 1.0  # Mosaic数据增强的概率
  mixup: 0.0  # MixUp数据增强的概率
  ...
  hflip: 0.5 
水平翻转的概率
 
    vflip: 0.0 # 垂直翻转的概率
...
hsv_h: 0.015 # 色相变换系数
hsv_s: 0.7 # 饱和度变换系数
hsv_v: 0.4 # 亮度变换系数
...

#训练集设置
 
train:
path: dataset/images/train # 训练集图像路径
 
img_size: [640, 640] # 输入图像大小
 
batch_size: 16 # 批次大小
 
 
#验证集设置
 
val:
path: dataset/images/val # 验证集图像路径
 
img_size: [640, 640] # 输入图像大小
 
batch_size: 16 # 批次大小