使用YOLOv8 Pose模型推理图像，生成可视化结果和LabelMe格式的标注文件

见代码

点击查看代码

# 导入必要的库
import cv2  # OpenCV库，用于图像读取、保存和处理
import json  # 用于JSON文件的读写，生成LabelMe格式标注文件
from pathlib import Path  # 更优雅的路径处理，兼容Windows/Linux系统
from ultralytics import YOLO  # 导入Ultralytics的YOLOv8库，核心的姿态估计模型
from PIL import Image  # PIL库，用于获取图像原始尺寸（比cv2更准确）
from glob import glob  # 通配符文件查找（代码中未实际使用，保留是为了兼容注释掉的写法）
import numpy as np  # 数值计算库，处理关键点、坐标等数组操作


def pre_pose(MODEL_PATH, IMAGE_DIR, OUTPUT_VIS_DIR, CLASS_NAMES, OUTPUT_LABELME_DIR):
    """
    核心功能函数：使用YOLOv8 Pose模型推理图像，生成可视化结果和LabelMe格式的标注文件
    参数说明：
        MODEL_PATH: str/Path - YOLOv8 Pose模型文件路径（.pt格式，训练好的权重）
        IMAGE_DIR: str/Path - 待推理的图像文件夹路径
        OUTPUT_VIS_DIR: str/Path - 推理结果可视化图像的保存路径
        CLASS_NAMES: list - 类别名称列表（与训练时的类别顺序一致，如["undercut", "line"]）
        OUTPUT_LABELME_DIR: str/Path - LabelMe格式JSON标注文件的保存路径
    """
    # 加载YOLOv8 Pose模型（必须是pose模型，不能用检测/分割模型）
    # 模型加载逻辑：优先加载本地权重，若本地无则自动下载官方预训练模型
    model = YOLO(MODEL_PATH)

    # 遍历指定文件夹下的所有图像文件，支持大小写后缀（jpg/JPG/png/PNG）
    # Path.glob() 是比glob.glob()更推荐的路径遍历方式，兼容性更好
    for img_path in Path(IMAGE_DIR).glob("*.[jJp][pPn][gG]"):
        # 打印当前处理的文件名，方便跟踪进度和排查问题
        print(f"Processing: {img_path.name}")

        # --- 1. 模型推理核心步骤 ---
        # img_path: 图像路径（支持str/Path对象）
        # conf=0.5: 置信度阈值，只保留置信度≥0.5的检测结果，过滤低置信度噪声
        # iou=0.2: NMS交并比阈值，避免重复检测同一目标
        # verbose=False: 关闭推理过程中的详细日志输出，保持控制台整洁
        results = model(img_path, conf=0.5, iou=0.2, verbose=False)
        result = results[0]  # 取第一张图像的推理结果（单图推理，results列表只有1个元素）

        # --- 2. 推理结果可视化并保存 ---
        # result.plot(): 自动绘制检测框（bbox）和关键点（keypoints），返回绘制后的图像数组
        # 绘制规则：不同类别用不同颜色，关键点按顺序连线（YOLOv8 Pose默认逻辑）
        vis_img = result.plot()
        # 拼接可视化图像的保存路径（输出文件夹 + 原文件名）
        vis_save_path = Path(OUTPUT_VIS_DIR) / img_path.name
        # 保存可视化图像：cv2.imwrite要求路径为字符串，Path对象需转str
        cv2.imwrite(str(vis_save_path), vis_img)

        # --- 3. 转换推理结果为LabelMe JSON格式（核心逻辑）---
        # 用PIL打开图像，获取原始宽高（避免cv2读取时的通道/尺寸偏差）
        with Image.open(img_path) as im:
            img_w, img_h = im.size  # img_w: 图像宽度(px)，img_h: 图像高度(px)

        # 初始化LabelMe格式的JSON数据结构（固定模板，需严格遵循）
        labelme_data = {
            "version": "5.0.1",  # LabelMe版本号，固定值即可
            "flags": {},  # 自定义标记，无特殊需求留空
            "shapes": [],  # 核心字段：存储检测框和关键点的标注信息
            "imagePath": img_path.name,  # 图像文件名（LabelMe加载时对应本地文件）
            "imageData": None,  # 图像base64编码，无需填写（LabelMe会自动生成）
            "imageHeight": img_h,  # 图像高度
            "imageWidth": img_w  # 图像宽度
        }

        # 检查是否有有效的姿态估计结果（避免空结果导致的报错）
        # result.keypoints: 关键点结果对象，None表示无检测结果
        if result.keypoints is not None and len(result.keypoints) > 0:
            # 提取检测结果并转换为numpy数组（CPU张量转numpy，避免GPU张量操作）
            boxes = result.boxes.xyxy.cpu().numpy()  # 检测框坐标 [N, 4]，格式：x1,y1,x2,y2（左上角/右下角）
            classes = result.boxes.cls.cpu().numpy()  # 检测类别ID [N,]，N为检测到的目标数量
            confs = result.boxes.conf.cpu().numpy()  # 检测置信度 [N,]，0~1之间
            keypoints = result.keypoints.xy.cpu().numpy()  # 关键点坐标 [N, K, 2]，K为关键点数量（如COCO 17点/自定义4点）

            num_dets = len(boxes)  # 获取检测到的目标总数

            # 遍历每个检测目标，逐个生成标注信息
            for i in range(num_dets):
                # 提取当前目标的检测框坐标，并转换为float类型（JSON要求数值类型）
                x1, y1, x2, y2 = map(float, boxes[i])
                # 提取当前目标的类别ID（int）和置信度（float）
                cls_id = int(classes[i])
                conf = float(confs[i])
                # 提取当前目标的所有关键点坐标 [K, 2]
                kpts = keypoints[i]

                # 根据类别ID获取类别名称，防止索引越界（越界则标为unknown）
                label = CLASS_NAMES[cls_id] if cls_id < len(CLASS_NAMES) else "unknown"

                # === 第一步：添加检测框（bbox）到LabelMe标注 ===
                bbox_shape = {
                    "label": label,  # 检测框类别名称
                    "points": [[x1, y1], [x2, y2]],  # 检测框坐标（左上角、右下角）
                    "group_id": i,  # 分组ID：同一目标的框和关键点用相同ID，方便关联
                    "description": f"bbox_conf={conf:.2f}",  # 备注：检测框置信度（保留2位小数）
                    "shape_type": "rectangle",  # 形状类型：矩形（LabelMe固定值）
                    "flags": {}  # 自定义标记，留空
                }
                # 将检测框添加到LabelMe的shapes列表
                labelme_data["shapes"].append(bbox_shape)

                # === 第二步：添加关键点到LabelMe标注 ===
                # 遍历当前目标的所有关键点（ipt为索引，pt为坐标）
                for ipt, pt in enumerate(kpts):
                    # 提取关键点坐标并转换为float类型
                    x, y = float(pt[0]), float(pt[1])
                    # 过滤无效关键点（YOLOv8中无效关键点坐标为(0,0)，需跳过）
                    if (np.fabs(x) < 1e-6) and (np.fabs(y) < 1e-6):
                        continue
                    # 构建关键点的LabelMe标注结构
                    polygon_shape = {
                        "label": f"{ipt + 2}",  # 关键点标签：用索引+2命名（可根据需求修改，如"key_point_1"）
                        "points": [[x, y]],  # 关键点坐标（LabelMe的point类型要求二维列表）
                        "group_id": i,  # 与对应检测框同组ID，关联框和关键点
                        "description": f"kpts_conf={conf:.2f}",  # 备注：关键点所属目标的置信度
                        "shape_type": "point",  # 形状类型：点（LabelMe固定值）
                        "flags": {}  # 自定义标记，留空
                    }
                    # 将关键点添加到LabelMe的shapes列表
                    labelme_data["shapes"].append(polygon_shape)

        # --- 4. 保存LabelMe JSON文件 ---
        # 拼接JSON文件路径：输出文件夹 + 原文件名（去掉后缀） + .json
        json_path = Path(OUTPUT_LABELME_DIR) / f"{img_path.stem}.json"
        # 写入JSON文件：ensure_ascii=False支持中文，indent=2格式化输出（可读性更好）
        with open(json_path, 'w', encoding='utf-8') as f:
            json.dump(labelme_data, f, indent=2, ensure_ascii=False)

    # 所有图像处理完成后，打印完成提示
    print("✅ Pose 推理、可视化、LabelMe 转换完成！")


def main():
    """
    主函数：配置参数 + 调用核心函数
    所有可配置参数都集中在这里，方便修改和维护
    """
    # ==================== 核心配置区（根据实际需求修改）====================
    # 1. 模型路径：替换为你的YOLOv8 Pose训练权重（.pt文件）
    MODEL_PATH = R"D:\work\ultralytics-main-v8_\runs\pose\exp_15undercut_yolov8n_pose_1280_iter5_260311\weights\best.pt"

    # 2. 待推理图像文件夹：替换为你的图像路径
    IMAGE_DIR = R"D:\work\0_task\15_MWmangKongUndercut\0_undercut_pose_1280\1_anno_done"

    # 3. 可视化结果输出路径：自动在原图像路径后拼接后缀，无需手动创建（代码会自动建）
    OUTPUT_VIS_DIR = IMAGE_DIR + r"_pyresult_view_pose"

    # 4. LabelMe JSON输出路径：同上
    OUTPUT_LABELME_DIR = IMAGE_DIR + r"_pyresult_anno_pose"

    # 5. 类别名称列表：与训练时的.yaml配置文件中的names一致，顺序不能错
    CLASS_NAMES = ["1"]

    # 创建输出目录：exist_ok=True 表示如果目录已存在则不报错
    Path(OUTPUT_VIS_DIR).mkdir(exist_ok=True)
    Path(OUTPUT_LABELME_DIR).mkdir(exist_ok=True)

    # 调用核心推理函数，传入配置参数
    pre_pose(MODEL_PATH, IMAGE_DIR, OUTPUT_VIS_DIR, CLASS_NAMES, OUTPUT_LABELME_DIR)


# 程序入口：只有直接运行该脚本时，才执行main函数（避免被导入时自动执行）
if __name__ == '__main__':
    main()