from ultralytics import YOLO
import torch

def apply_data_augmentation():
    """应用数据增强技术提升模型泛化能力"""
    # 检查是否支持CUDA
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f'Using device: {device}')

    # 加载预训练模型
    model = YOLO('yolo12s.yaml')

    # 开始训练
    model.train(
        data='datasets/data.yaml',
        epochs=600,
        imgsz=640,
        batch=32,
        device=0,
        optimizer='SGD',
        lr0=0.001,
        project='run',
        name='train',
        workers=4,
    # 设置模型保存周期（每50个epoch保存一次）
        save_period=50,
        # 数据增强参数
        hsv_h=0.015,  # 色调增强
        hsv_s=0.7,   # 饱和度增强
        hsv_v=0.4,   # 明度增强
        degrees=30,  # 随机旋转角度
        translate=0.1,  # 平移比例
        scale=0.5,   # 缩放比例
        shear=15,    # 剪切变换角度
        perspective=0.0001,  # 透视变换概率
        flipud=0.5,  # 上下翻转概率
        fliplr=0.5,  # 左右翻转概率
        mosaic=1.0,  # 马赛克增强概率
        mixup=0.2,   # MixUp增强概率
        copy_paste=0.3,  # 复制粘贴增强概率
        auto_augment='randaugment',  # 自动增强策略
        # erasing=0.4,  # 随机擦除概率
        # fraction=0.9,  # 训练验证集划分比例
        cache= 'disk'  # 启用缓存加速训练
    )

if __name__ == '__main__':
    apply_data_augmentation()