feat(yolo): 添加刀具识别AI模型训练与自动标注功能

- 新增 README.md 文档,包含完整的 YOLOv8 训练环境搭建指南
- 添加 docker-compose.yml 配置文件,支持 Mac M系列环境下的容器化部署
- 实现 train_and_export.py 核心脚本,集成模型训练、验证和ONNX格式导出功能
- 创建 data.yaml 数据集描述文件,定义刀具分类(普通刀具和管制刀具)
- 开发 auto_label.py 自动标注脚本,支持AI辅助的数据集标注处理
- 集成 YOLOv8 中型模型,优化标注准确率和识别能力
- 添加数据集验证和可视化功能,便于标注质量检查
This commit is contained in:
yuzl6
2026-03-10 17:15:05 +08:00
commit 1753cafaab
49 changed files with 331 additions and 0 deletions
Vendored
BIN
View File
Binary file not shown.
+174
View File
@@ -0,0 +1,174 @@
你好!我是 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) 的内部数据结构**
+91
View File
@@ -0,0 +1,91 @@
import os
import shutil
from ultralytics import YOLO
# ================= 配置区 =================
DIR_NORMAL = 'datasets/raw_normal'
DIR_PROHIBITED = 'datasets/raw_prohibited'
IMG_TARGET = 'datasets/knife_data/images/train'
LBL_TARGET = 'datasets/knife_data/labels/train'
VERIFY_DIR = 'runs/verify_labels'
# 【关键优化 1】:换用中型模型 (yolov8m.pt),识别能力比 n 强很多
# 如果你的 Mac 内存足够(8G以上),可以直接写 yolov8x.pt 效果最好
model = YOLO('yolov8m.pt')
# ==========================================
def safe_makedirs(path):
if not os.path.exists(path):
os.makedirs(path, exist_ok=True)
def process_folder(src_dir, target_class_id, prefix):
if not os.path.exists(src_dir) or not any(f.lower().endswith(('.jpg', '.png', '.jpeg')) for f in os.listdir(src_dir)):
return
print(f"\n🚀 正在对【{prefix}】进行深度扫描,ID -> {target_class_id}")
# 【关键优化 2】:调低 conf (置信度阈值)
# 0.1 代表只要有 10% 的把握是刀,就把它标出来。
results = model.predict(
source=src_dir,
conf=0.1,
save_txt=True,
save=True,
classes=[43],
project='runs/detect',
name=f'tmp_{prefix}',
exist_ok=True
)
actual_save_dir = results[0].save_dir
tmp_label_dir = os.path.join(actual_save_dir, 'labels')
# 1. 搬运并修正标签
label_count = 0
if os.path.exists(tmp_label_dir):
for f in os.listdir(tmp_label_dir):
if f.endswith('.txt'):
src_path = os.path.join(tmp_label_dir, f)
dst_path = os.path.join(LBL_TARGET, f)
with open(src_path, 'r') as file:
lines = file.readlines()
corrected_lines = [f"{target_class_id} " + " ".join(l.split()[1:]) + "\n" for l in lines]
with open(dst_path, 'w') as file:
file.writelines(corrected_lines)
label_count += 1
# 2. 同步图片
img_count = 0
for f in os.listdir(src_dir):
if f.lower().endswith(('.jpg', '.jpeg', '.png')):
shutil.copy(os.path.join(src_dir, f), os.path.join(IMG_TARGET, f))
img_count += 1
# 3. 整理验证图
for f in os.listdir(actual_save_dir):
if f.lower().endswith(('.jpg', '.jpeg', '.png')):
shutil.move(os.path.join(actual_save_dir, f), os.path.join(VERIFY_DIR, f"{prefix}_{f}"))
print(f" ✅ 完成:图片 {img_count} 张,AI 捕获标注 {label_count}")
def run_workflow():
# 初始化环境
for d in [IMG_TARGET, LBL_TARGET, VERIFY_DIR]:
if os.path.exists(d): shutil.rmtree(d)
safe_makedirs(d)
process_folder(DIR_NORMAL, 0, "normal")
process_folder(DIR_PROHIBITED, 1, "prohibited")
print("\n" + ""*40)
print(f"✅ 自动标注执行完毕")
print(f"📊 识别统计:共生成 {len(os.listdir(LBL_TARGET))} 个标签文件")
print(f"⚠️ 警告:如果标签文件少于图片总数,说明 AI 还是没认全。")
print(f"🛠️ 下一步建议:使用 LabelImg 打开 {IMG_TARGET} 手动补齐剩下的框。")
print(""*40)
if __name__ == "__main__":
run_workflow()
BIN
View File
Binary file not shown.
BIN
View File
Binary file not shown.
+9
View File
@@ -0,0 +1,9 @@
# 这里的路径必须是容器内部的绝对路径
path: /usr/src/app/datasets/knife_data
train: images/train
val: images/val
# 类别定义 (根据你的业务字典自定义)
names:
0: normal_knife # 普通民用刀具
1: prohibited_knife # 管制刀具
Binary file not shown.
Binary file not shown.
Binary file not shown.

After

Width:  |  Height:  |  Size: 358 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 34 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 31 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 24 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 101 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 113 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 298 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 346 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 84 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 854 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 102 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 17 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 56 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 328 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 28 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 337 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 52 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 21 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 174 KiB

+12
View File
@@ -0,0 +1,12 @@
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
BIN
View File
Binary file not shown.
Binary file not shown.

After

Width:  |  Height:  |  Size: 461 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 70 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 54 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 33 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 154 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 389 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 516 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 126 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 135 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 32 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 102 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 428 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 35 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 486 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 77 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 32 KiB

+45
View File
@@ -0,0 +1,45 @@
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()