yolo模型训练

训练

创建config目录
将yaml放进

klbq.yaml

train: /run/media/boqi/存储/AI/YOLO/dataset/klbq/train/images
val: /run/media/boqi/存储/AI/YOLO/dataset/klbq/valid/images
test: /run/media/boqi/存储/AI/YOLO/dataset/klbq/test/images
nc: 1
names:
  0: enemy

创建一个 train.py

from ultralytics import YOLO
import torch
import sys
import logging
from pathlib import Path
import argparse
import json
import time

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class TrainingConfig:
    """训练配置类"""
    def __init__(self, project_name='wllt'):
        self.project_name = project_name
        self.model_path = '/opt/ultralytics/config/yolo12s.pt'
        self.data_path = f'config/{project_name}.yaml'
        self.epochs = 300
        self.imgsz = 224
        self.batch_size = -1
        self.resume = False

    def update_project_name(self, new_name):
        """更新项目名称并相应更新相关路径"""
        self.project_name = new_name
        self.data_path = f'config/{new_name}.yaml'

    def to_dict(self):
        """转换为字典格式"""
        return {
            'project_name': self.project_name,
            'model_path': self.model_path,
            'data_path': self.data_path,
            'epochs': self.epochs,
            'imgsz': self.imgsz,
            'batch_size': self.batch_size,
            'resume': self.resume
        }

def setup_device():
    """设置并验证GPU设备"""
    if not torch.cuda.is_available():
        raise RuntimeError("当前没有可用的GPU设备")

    device_count = torch.cuda.device_count()
    if device_count == 0:
        raise RuntimeError("未检测到任何GPU设备")

    gpu_props = torch.cuda.get_device_properties(0)
    logger.info(f"使用GPU: {gpu_props.name}, 显存: {gpu_props.total_memory / 1024**3:.1f}GB")

    return 'cuda:0'

def validate_files(model_path, data_path):
    """验证文件是否存在"""
    if not Path(model_path).exists():
        raise FileNotFoundError(f"模型文件不存在: {model_path}")
    if not Path(data_path).exists():
        raise FileNotFoundError(f"数据配置文件不存在: {data_path}")

def train_yolo(config: TrainingConfig):
    """YOLO训练函数"""
    try:
        # 验证文件和设备
        validate_files(config.model_path, config.data_path)
        device = setup_device()

        # 加载模型
        model = YOLO(config.model_path)

        # 保存训练配置
        config_save_path = f"runs/{config.project_name}/config.json"
        Path(config_save_path).parent.mkdir(parents=True, exist_ok=True)

        with open(config_save_path, 'w', encoding='utf-8') as f:
            json.dump(config.to_dict(), f, indent=2, ensure_ascii=False)

        logger.info(f"开始训练项目: {config.project_name}")
        start_time = time.time()

        # 训练
        results = model.train(
            data=config.data_path,
            epochs=config.epochs,
            batch=config.batch_size,
            imgsz=config.imgsz,
            device=device,
            amp=True,
            cache=True,
            patience=max(10, config.epochs // 30),
            augment=True,
            mosaic=False,
            save_period=max(10, config.epochs // 30),
            project='runs',
            name=config.project_name,
            exist_ok=True,
            resume=config.resume,
            hsv_h=0.015,   #(色相 不同地图的整体色调、滤镜、以及技能（如蝰蛇的毒）会影响全局色彩。可以设置一个较小的值
            hsv_s=0.1,   #饱和度扰动
            hsv_v=0.05,   #亮度 这是最有用的！角色会出现在阴影或高光区域，亮度变化极大
            fliplr=0.0,  #左右翻转
            flipud=0.0,  #上下翻转 关
            degrees=0.0, #旋转 关
            scale=0.1,      # 缩放 可开小幅度 0.9~1.1，如需微调大小
            shear=0.0,     # 斜切 关
            translate=0.01, # 平移 可视抖动用，太大会造成错位
            copy_paste=0.0, #它会随机将一个目标的“贴片”复制粘贴到另一张图上，能有效模拟目标重叠和遮挡的情况
            mixup=0.0, #轻微开启有助于正则化，防止过拟合，如 mixup=0.1
        )

        training_time = time.time() - start_time
        logger.info(f"项目 {config.project_name} 训练完成! 耗时: {training_time/3600:.2f} 小时")

        return results

    except Exception as e:
        logger.error(f"训练失败: {e}")
        sys.exit(1)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='YOLO训练脚本')
    parser.add_argument('--project-name', type=str, default='wllt', help='项目名称')
    parser.add_argument('--epochs', type=int, default=300, help='训练轮数')
    parser.add_argument('--batch-size', type=int, default=-1, help='批次大小')
    parser.add_argument('--resume', action='store_true', help='恢复训练')

    args = parser.parse_args()

    # 创建配置
    config = TrainingConfig(args.project_name)
    config.epochs = args.epochs
    config.batch_size = args.batch_size
    config.resume = args.resume

    # 开始训练
    train_yolo(config)

导出ONN

创建一个 onnx.py

from ultralytics import YOLO

# 加载训练好的YOLO12模型权重（官方或自定义）
model = YOLO("/opt/ultralytics/runs/klbq/weights/best.pt")  # 或者你的自定义权重路径

# 导出为ONNX格式，默认会生成yolov12n.onnx文件
model.export(format="onnx", imgsz=224, half=False, dynamic=False, simplify=True, device=0)

数据集分配

import os
import random
import shutil
from tqdm import tqdm

# 清空文件夹
def clear_dir(path):
    if os.path.exists(path):
        shutil.rmtree(path)
    os.makedirs(path)

def ensure_dirs(dir_list):
    for d in dir_list:
        os.makedirs(d, exist_ok=True)

def get_valid_filenames(image_dir, label_dir, img_ext='.jpg', label_ext='.txt'):
    image_names = {os.path.splitext(f)[0] for f in os.listdir(image_dir) if f.endswith(img_ext)}
    label_names = {os.path.splitext(f)[0] for f in os.listdir(label_dir) if f.endswith(label_ext)}
    return sorted(list(image_names & label_names))

def split_dataset(filenames, train_ratio, valid_ratio):
    total = len(filenames)
    train_cnt = int(total * train_ratio)
    valid_cnt = int(total * valid_ratio)
    return (
        filenames[:train_cnt],
        filenames[train_cnt:train_cnt+valid_cnt],
        filenames[train_cnt+valid_cnt:]
    )

def copy_file_with_progress(src, dst):
    size = os.path.getsize(src)
    with open(src, 'rb') as fsrc, open(dst, 'wb') as fdst:
        for chunk in iter(lambda: fsrc.read(1024*1024), b''):
            fdst.write(chunk)

def batch_copy(filenames, image_dir, label_dir, img_ext, label_ext, out_img_dir, out_label_dir, pbar):
    for fname in filenames:
        src_img = os.path.join(image_dir, fname + img_ext)
        dst_img = os.path.join(out_img_dir, fname + img_ext)
        src_lbl = os.path.join(label_dir, fname + label_ext)
        dst_lbl = os.path.join(out_label_dir, fname + label_ext)
        copy_file_with_progress(src_img, dst_img)
        copy_file_with_progress(src_lbl, dst_lbl)
        pbar.update(1)

def main():
    # random.seed(123)

    # 配置参数 - 只需要在这里修改路径
    output_dir = '/run/media/boqi/存储/AI/YOLO/dataset/klbq/'
    root_dir = '/run/media/boqi/存储/AI/YOLO/out/klbq/'

    image_dir = os.path.join(root_dir, 'Images')
    label_dir = os.path.join(root_dir, 'labels')
    img_ext, label_ext = '.jpg', '.txt'
    train_ratio, valid_ratio, test_ratio = 0.8, 0.1, 0.1

    # 检查目录是否存在
    if not os.path.isdir(image_dir) or not os.path.isdir(label_dir):
        raise FileNotFoundError("图像或标签目录不存在，请检查路径！")

    # 先清空输出目录
    clear_dir(output_dir)

    filenames = get_valid_filenames(image_dir, label_dir, img_ext, label_ext)
    random.shuffle(filenames)
    train_files, valid_files, test_files = split_dataset(filenames, train_ratio, valid_ratio)

    # 输出各部分图片数量
    print(f"train: {len(train_files)} 张")
    print(f"valid: {len(valid_files)} 张")
    print(f"test: {len(test_files)} 张")

    # 输出目录
    dirs = {
        'train_img': os.path.join(output_dir, 'train', 'images'),
        'train_lbl': os.path.join(output_dir, 'train', 'labels'),
        'valid_img': os.path.join(output_dir, 'valid', 'images'),
        'valid_lbl': os.path.join(output_dir, 'valid', 'labels'),
        'test_img': os.path.join(output_dir, 'test', 'images'),
        'test_lbl': os.path.join(output_dir, 'test', 'labels'),
    }

    ensure_dirs(dirs.values())
    total_files = len(train_files) + len(valid_files) + len(test_files)
    with tqdm(total=total_files, desc='总进度', unit='文件') as pbar:
        batch_copy(train_files, image_dir, label_dir, img_ext, label_ext, dirs['train_img'], dirs['train_lbl'], pbar)
        batch_copy(valid_files, image_dir, label_dir, img_ext, label_ext, dirs['valid_img'], dirs['valid_lbl'], pbar)
        batch_copy(test_files, image_dir, label_dir, img_ext, label_ext, dirs['test_img'], dirs['test_lbl'], pbar)

if __name__ == '__main__':
    main()

voc转Yolo

import os
import xml.etree.ElementTree as ET
from tqdm import tqdm  # 进度条库

images_dir = './Images'
annotations_dir = './Annotations'
labels_dir = './labels'

if not os.path.exists(labels_dir):
    os.makedirs(labels_dir)

# 获取图片和标注文件名（不带扩展名）
image_files = [os.path.splitext(f)[0] for f in os.listdir(images_dir) if f.lower().endswith('.jpg')]
annotation_files = [os.path.splitext(f)[0] for f in os.listdir(annotations_dir) if f.lower().endswith('.xml')]

# 1. 生成空标注（带进度条）
print("正在为没有标注的图片生成空标注文件...")
no_annotation_images = set(image_files) - set(annotation_files)
for img_name in tqdm(no_annotation_images, desc="生成空标注", ncols=80):
    with open(os.path.join(labels_dir, img_name + '.txt'), 'w') as f:
        pass  # 空文件

# 2. 转换XML为YOLO格式（带进度条）
def voc_to_yolo(xml_path):
    tree = ET.parse(xml_path)
    root = tree.getroot()
    size = root.find('size')
    width = int(size.find('width').text)
    height = int(size.find('height').text)
    lines = []
    for obj in root.findall('object'):
        cls_id = 0  # 假设只有一个类别
        xmlbox = obj.find('bndbox')
        xmin = int(xmlbox.find('xmin').text)
        ymin = int(xmlbox.find('ymin').text)
        xmax = int(xmlbox.find('xmax').text)
        ymax = int(xmlbox.find('ymax').text)
        x_center = (xmin + xmax) / 2.0 / width
        y_center = (ymin + ymax) / 2.0 / height
        w = (xmax - xmin) / width
        h = (ymax - ymin) / height
        lines.append(f"{cls_id} {x_center:.6f} {y_center:.6f} {w:.6f} {h:.6f}")
    return '\n'.join(lines)

print("正在将XML标注转换为YOLO格式...")
for xml_name in tqdm(annotation_files, desc="转换标注", ncols=80):
    xml_path = os.path.join(annotations_dir, xml_name + '.xml')
    yolo_txt = voc_to_yolo(xml_path)
    with open(os.path.join(labels_dir, xml_name + '.txt'), 'w') as f:
        f.write(yolo_txt)

print("\n全部处理完成！labels 目录已生成所有YOLO格式标注。")