yuzl6
1753cafaab
- 新增 README.md 文档，包含完整的 YOLOv8 训练环境搭建指南 - 添加 docker-compose.yml 配置文件，支持 Mac M系列环境下的容器化部署 - 实现 train_and_export.py 核心脚本，集成模型训练、验证和ONNX格式导出功能 - 创建 data.yaml 数据集描述文件，定义刀具分类(普通刀具和管制刀具) - 开发 auto_label.py 自动标注脚本，支持AI辅助的数据集标注处理 - 集成 YOLOv8 中型模型，优化标注准确率和识别能力 - 添加数据集验证和可视化功能，便于标注质量检查
45 lines
1.4 KiB
Python
45 lines
1.4 KiB
Python
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() |