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