from ultralytics import YOLO

if __name__ == '__main__':
    # ✅ 推荐：加载官方预训练模型 或 自己的 best.pt
    # model = YOLO('yolo11m-obb.pt')  # 官方预训练（如果有）
    model = YOLO(r'/home/hx/yolo/ultralytics_yolo11-main/runs/train/exp_obb_new/weights/best.pt')  # 使用 best 而非 last

    results = model.train(
        data='obb_data1.yaml',
        epochs=500,                    # 减少 epochs，配合早停
        patience=0,                   # 50 轮无提升则停止
        imgsz=640,
        batch=4,
        workers=10,
        device='0',
        project='runs/train',
        name='exp_obb5',               # 建议递增实验编号
        exist_ok=False,

        # 优化器
        optimizer='AdamW',
        lr0=0.0005,                    # 更稳定的学习率
        weight_decay=0.01,
        momentum=0.937,

        # 数据增强（OBB 关键）
        degrees=5.0,                  # 随机旋转 ±10°
        translate=0.1,                 # 平移
        scale=0.5,                     # 缩放比例
        shear=1.0,                     # 剪切
        flipud=0.0,                    # 不推荐上下翻转（角度易错）
        fliplr=0.5,                    # ✅ 水平翻转，OBB 支持良好
        hsv_h=0.015,                   # 色调扰动
        hsv_s=0.7,                     # 饱和度
        hsv_v=0.4,                     # 明度

        # 其他
        close_mosaic=50,               # 最后10轮关闭 Mosaic 增强，提升稳定性
        val=True,                      # 每轮验证
    )