knife-ai-train/README.md

你好！我是 Gemini。为了确保你能顺利地在 Mac 环境下利用 Docker 完成 YOLOv8 模型的训练、验证以及最终导出（ONNX 格式），我为你整理了这份完整的文档和代码集。

这份文档将指导你完成从零构建训练环境到获得生产可用模型的全过程。

---

刀具认定 AI 模型训练与导出指南 (YOLOv8)

1. 目录结构规划

在你的 Mac 上创建一个工作目录（例如 knife-ai-training），并按以下结构组织文件：

knife-ai-training/
├── docker-compose.yml        # Docker 容器配置
├── train_and_export.py       # 核心训练与导出脚本
├── models/                   # 存放生成的模型权重 (.pt, .onnx)
└── datasets/                 # 存放训练数据
    └── knife_data/
        ├── data.yaml         # YOLO 数据集描述文件
        ├── images/           # 存放图片
        │   ├── train/        # 训练集
        │   └── val/          # 验证集
        └── labels/           # 存放标注 (.txt)
            ├── train/
            └── val/

---

2. 环境配置文件 (docker-compose.yml)

该文件定义了一个运行 YOLOv8 所需的 Python 环境。

返回的是完整代码：

services:
  yolov8-train:
    image: ultralytics/ultralytics:latest-cpu
    container_name: yolov8_knife
    volumes:
      # 将当前工作目录挂载到容器内的 /usr/src/app
      - .:/usr/src/app
    # 适配 Mac M系列(ARM64) 运行 Linux amd64 镜像
    platform: linux/amd64 
    working_dir: /usr/src/app
    # 保持容器持续运行，方便执行 exec 命令
    command: sleep infinity

---

3. 数据集描述文件 (datasets/knife_data/data.yaml)

此文件告诉 YOLOv8 你的数据在哪里，以及有哪些分类。

返回的是完整代码：

# 这里的路径必须是容器内部的绝对路径
path: /usr/src/app/datasets/knife_data 
train: images/train
val: images/val

# 类别定义 (根据你的业务字典自定义)
names:
  0: normal_knife      # 普通民用刀具
  1: prohibited_knife  # 管制刀具

---

4. 核心训练与导出脚本 (train_and_export.py)

该脚本集成了加载预训练权重、正式训练、模型验证、以及导出为 Java 推理所需的 ONNX 格式。

返回的是完整代码：

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()

---

5. 执行流程操作步骤

步骤 A：启动容器

在 knife-ai-training 目录下打开 Mac 终端执行：

docker-compose up -d

步骤 B：执行训练与导出

容器启动后，直接通过以下命令触发 Python 脚本：

docker exec -it yolov8_knife python3 train_and_export.py #训练模型
docker exec -it yolov8_knife python3 auto_label.py #标注模型

步骤 C：获取产出物

脚本运行结束后，检查你的 Mac 本地目录：

• 训练权重：models/knife_task/weights/best.pt
• 生产模型：models/knife_task/weights/best.onnx （将此文件重命名为 knife_v8.onnx 放入 Java 项目的 resources 下）。

---

6. 给后续数据集增加的建议

1. 数据清洗：确保 images 下的每个图片都有对应的同名 .txt 标签文件。
2. 标注格式：YOLO 格式为 class_id x_center y_center width height（归一化到 0-1 之间）。
3. 样本平衡：尽量保证管制刀具与普通刀具的图片数量比例接近 1:1，否则模型会产生“偏见”。

这份代码和配置为你搭建了标准化的生产线。只要你往 datasets 里丢新数据并重新跑一遍脚本，就能获得精度更高的模型。

返回的是完整代码。 你是否需要我继续为你详细讲解 YOLO 标注文件 (.txt) 的内部数据结构？