val_yolov11.py

import os
from pathlib import Path
from ultralytics import YOLO


def main():
    # ==========================================
    # 1. 路径配置 (请根据您的实际目录结构确认)
    # ==========================================
    PROJECT_ROOT = Path(__file__).resolve().parent

    # 替换为您实际训练产出的 best.pt 路径
    # 例如：PROJECT_ROOT / "save_models" / "yolov11_ppe_neg_ft" / "weights" / "best.pt"
    model_path = PROJECT_ROOT / "runs" / "train" / "yolov11_ppe_v1002" / "weights" / "best.pt"
    data_yaml = PROJECT_ROOT / "datasets" / "helmet-vest" / "rebuild" / "data.yaml"

    if not model_path.exists():
        print(f"❌ 找不到模型文件: {model_path}")
        return
    if not data_yaml.exists():
        print(f"❌ 找不到数据配置: {data_yaml}")
        return

    # ==========================================
    # 2. 加载模型
    # ==========================================
    print(f"🚀 正在加载模型: {model_path}")
    model = YOLO(str(model_path))

    # ==========================================
    # 3. 运行验证 (Validation)
    # ==========================================
    print("📊 开始在验证集上评估模型指标...")

    # val() 函数会自动计算 mAP、Precision、Recall 并生成图表
    metrics = model.val(
        data=str(data_yaml),
        split="val",  # "val" 为验证集，"test" 为测试集
        imgsz=640,  # 与训练时保持一致
        batch=16,  # 验证时不需要计算梯度，batch 可以适度调大
        conf=0.25,  # 评估默认的置信度阈值 (0.25 是 COCO 标准)
        iou=0.60,  # NMS 的 IOU 阈值
        device=0,  # 使用 GPU (0)
        plots=True,  # 【关键】必须为 True，才会生成混淆矩阵等可视化图表
        verbose=True
    )

    # ==========================================
    # 4. 提取并打印核心指标
    # ==========================================
    print("\n" + "=" * 50)
    print("🎉 验证任务完成！核心指标提取如下：")

    # 提取 mAP
    print(f"🏆 整体 mAP50:     {metrics.box.map50:.4f}")
    print(f"📏 整体 mAP50-95:  {metrics.box.map:.4f}")

    # 提取 Precision 和 Recall (取所有类别的平均值)
    # metrics.box.p 和 metrics.box.r 返回的是数组，通常取平均值(mean)来代表整体
    mean_precision = metrics.box.p.mean()
    mean_recall = metrics.box.r.mean()
    print(f"🎯 整体 Precision (精确率 - 越高说明误报越少): {mean_precision:.4f}")
    print(f"🔍 整体 Recall    (召回率 - 越高说明漏检越少): {mean_recall:.4f}")

    # 打印各类别详细指标
    print("\n📌 各类别详细指标 (P / R / mAP50):")
    class_names = model.names

    # 获取每个类别的独立指标
    for i, cls_idx in enumerate(metrics.box.ap_class_index):
        cls_name = class_names[cls_idx]
        p = metrics.box.p[i]
        r = metrics.box.r[i]
        map50 = metrics.box.ap50[i]
        print(f"  - {cls_name:<10}: Precision={p:.4f}, Recall={r:.4f}, mAP50={map50:.4f}")

    print("\n📂 详细评估报告与图表已保存至:")
    print(f"👉 {metrics.save_dir}")
    print("=" * 50)
    print("💡 重点检查指南：")
    print("请打开上述目录中的 'confusion_matrix.png'。")
    print("查看最后一行（背景行），验证背景被误认为头盔/反光衣的数量是否显著下降！")


if __name__ == "__main__":
    main()