from ultralytics import YOLO
import os


def train_segmentation_model(
    dataset_path="datasets/xianshiping/data.yaml",
    model_name="yolov8n-seg.pt",
    epochs=100,
    batch=8,
    imgsz=640,
    device="cpu"
):
    """
    训练YOLO分割模型（实例分割，支持多边形标注）
    
    Args:
        dataset_path: 数据集配置文件路径
        model_name: 预训练模型名称
                   可选: "yolov8n-seg.pt", "yolov8s-seg.pt", "yolov8m-seg.pt", 
                        "yolov8l-seg.pt", "yolov8x-seg.pt"
                        或 "yolo11n-seg.pt", "yolo11s-seg.pt" 等
        epochs: 训练轮数
        batch: 批次大小
        imgsz: 图片尺寸
        device: 设备，"cpu" 或 "cuda" 或 "0" (GPU编号)
    """
    # 检查数据集配置文件是否存在
    if not os.path.exists(dataset_path):
        print(f"错误: 数据集配置文件不存在: {dataset_path}")
        return
    
    # 检查模型文件是否存在，如果不存在会自动下载
    print(f"加载模型: {model_name}")
    model = YOLO(model_name)
    
    print(f"开始训练分割模型...")
    print(f"数据集: {dataset_path}")
    print(f"训练轮数: {epochs}")
    print(f"批次大小: {batch}")
    print(f"图片尺寸: {imgsz}")
    print(f"设备: {device}")
    
    # 自动检测并使用最佳设备
    import torch
    if device == "cpu":
        # 检查是否有Apple Silicon GPU (MPS)
        if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
            device = "mps"
            print("✓ 检测到Apple Silicon GPU，使用MPS加速训练")
        # 检查是否有CUDA GPU
        elif torch.cuda.is_available():
            device = "cuda"
            print("✓ 检测到CUDA GPU，使用GPU加速训练")
        else:
            print("⚠ 使用CPU训练（较慢，建议使用GPU）")
    
    print("-" * 50)
    
    # 开始训练
    results = model.train(
        data=dataset_path,
        epochs=epochs,
        batch=batch,
        imgsz=imgsz,
        device=device,
        task="segment",  # 明确指定分割任务
        project="runs/segment",  # 保存到segment目录
        name="train",  # 训练名称
    )
    
    print("-" * 50)
    print("训练完成！")
    print(f"模型保存在: {results.save_dir}")
    
    return results


def main():
    """主函数 - 训练检测仪数据集"""
    # 数据集信息：
    # - 训练集: 978张图片
    # - 验证集: 245张图片
    # - 总计: 1223张图片
    # - 类别: 6个（搭电线、拖车上有车、仪表盘、轮胎、平安马甲、检测仪）
    
    train_segmentation_model(
        dataset_path="datasets/xianshiping/data.yaml",
        model_name="yolov8n-seg.pt",  # 使用YOLOv8 nano分割模型
        epochs=150,      # 1223张照片，建议150轮（配合早停）
        batch=8,         # CPU/MPS建议8，GPU可以16-32
        imgsz=640,       # 图片尺寸
        device="cpu"     # 自动检测并使用最佳设备（MPS/GPU/CPU）
    )


if __name__ == '__main__':
    main()