import os
from ultralytics import YOLO

def main():
    # 1. 选择基础模型 (yolov8n 是最轻量级，适合嵌入式和实时推理)
    print(">>> 正在加载预训练模型...")
    model = YOLO('yolov8n.pt') 

    # 2. 开始训练
    # data: 指向 data.yaml 路径
    # epochs: 训练轮数 (模拟测试可设为 1-5，正式环境建议 100+)
    # imgsz: 图片输入尺寸 (标准为 640)
    # device: 'cpu' (Docker 模拟环境建议用 cpu)
    print(">>> 启动模型训练...")
    results = model.train(
        data='/usr/src/app/datasets/knife_data/data.yaml',
        epochs=10, 
        imgsz=640,
        device='cpu',
        project='models',
        name='knife_task'
    )

    # 3. 验证模型 (在验证集上跑一遍结果)
    print(">>> 正在执行模型验证...")
    metrics = model.val()
    print(f"模型精度 (mAP@50-95): {metrics.box.map}")

    # 4. 导出为 ONNX 格式 (Java / Spring Boot 推理必选)
    # format: 'onnx'
    # opset: 推荐 12-17，确保 Java 库的兼容性
    # slim: 压缩模型，减少体积
    print(">>> 正在导出为 ONNX 格式...")
    onnx_path = model.export(
        format='onnx', 
        imgsz=640, 
        opset=12,
        simplify=True
    )
    
    print(f"\n[成功] 模型训练、验证与导出已完成！")
    print(f"最终推理模型位于: {onnx_path}")

if __name__ == "__main__":
    main()