from ultralytics import YOLO
import os

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

def main():
    # 初始化 Weights & Biases
    import wandb
    wandb.login(key='7cfbcf76a18a8441b04eb5d7adb988e69a79705e')  # 替换为你的 API 密钥
    wandb.init(project="YOLO-Training", name="YOLOv11_finetune", mode="online")

    # 加载 YOLO 模型
    model = YOLO('yolo11n.pt')

    # 开始训练
    model.train(
        data=r"D:\AIM\lemon\lemon_quality_dataset_YOLO11\data.yaml",  # 数据集配置文件路径
        epochs=2,                              # 训练轮次
        imgsz=640,                             # 图像大小
        batch=16,                              # 批量大小
        lr0=0.005,                              # 初始学习率
        workers=0,                             # 设置为 0，禁用多进程数据加载
        device=0,                              # 设备 (0 = 第一块 GPU, 'cpu' = CPU)
        project="YOLO-Training",               # W&B 项目名称
        name="YOLOv11_finetune",               # W&B 实验名称
        exist_ok=True                          # 如果目录存在是否覆盖
    )

if __name__ == '__main__':
    import torch.multiprocessing as mp
    mp.freeze_support()  # 解决 Windows 下的多进程启动问题
    main()