from ultralytics import YOLO

# 将所有执行代码都放入这个 if 语句块中
if __name__ == '__main__':
    # 1. 加载一个预训练模型
    # 这行代码可以在 if 语句内外，但为了代码清晰和规范，建议放入
    model = YOLO('yolov8n.pt')

    # 2. 训练模型
    # 核心的训练启动代码必须在 if 语句内
    results = model.train(data='train_data_up/weld.yaml',
                          epochs=50,
                          imgsz=640,
                          device=0,
                          workers=8)  # 可以显式指定 workers 数量

    # 3. (可选) 打印训练结果保存的路径
    print("训练完成！")
    # 注意: 在新版ultralytics中，results可能不直接包含save_dir，但训练日志会打印出来
    # 训练结果通常保存在 runs_bottom/detect/trainX 目录下