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

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