yuzl6
1753cafaab
- 新增 README.md 文档，包含完整的 YOLOv8 训练环境搭建指南 - 添加 docker-compose.yml 配置文件，支持 Mac M系列环境下的容器化部署 - 实现 train_and_export.py 核心脚本，集成模型训练、验证和ONNX格式导出功能 - 创建 data.yaml 数据集描述文件，定义刀具分类(普通刀具和管制刀具) - 开发 auto_label.py 自动标注脚本，支持AI辅助的数据集标注处理 - 集成 YOLOv8 中型模型，优化标注准确率和识别能力 - 添加数据集验证和可视化功能，便于标注质量检查
174 lines
5.1 KiB
Markdown
174 lines
5.1 KiB
Markdown
你好！我是 Gemini。为了确保你能顺利地在 Mac 环境下利用 Docker 完成 YOLOv8 模型的训练、验证以及最终导出（ONNX 格式），我为你整理了这份完整的文档和代码集。
|
||
|
||
这份文档将指导你完成从零构建训练环境到获得生产可用模型的全过程。
|
||
|
||
---
|
||
|
||
## 刀具认定 AI 模型训练与导出指南 (YOLOv8)
|
||
|
||
### 1. 目录结构规划
|
||
|
||
在你的 Mac 上创建一个工作目录（例如 `knife-ai-training`），并按以下结构组织文件：
|
||
|
||
```text
|
||
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 环境。
|
||
|
||
**返回的是完整代码：**
|
||
|
||
```yaml
|
||
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 你的数据在哪里，以及有哪些分类。
|
||
|
||
**返回的是完整代码：**
|
||
|
||
```yaml
|
||
# 这里的路径必须是容器内部的绝对路径
|
||
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 格式。
|
||
|
||
**返回的是完整代码：**
|
||
|
||
```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()
|
||
|
||
```
|
||
|
||
---
|
||
|
||
### 5. 执行流程操作步骤
|
||
|
||
#### 步骤 A：启动容器
|
||
|
||
在 `knife-ai-training` 目录下打开 Mac 终端执行：
|
||
|
||
```bash
|
||
docker-compose up -d
|
||
|
||
```
|
||
|
||
#### 步骤 B：执行训练与导出
|
||
|
||
容器启动后，直接通过以下命令触发 Python 脚本：
|
||
|
||
```bash
|
||
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) 的内部数据结构**？ |