hrnet训练的pt模型进行关键点识别的简单前向推理

本篇展示hrnet的人体关键点识别的简单前向推理,假设每张图只有一个人的情况下做出的检测推理。可以保存其关键点信息或直接在图中画出关键点可视化查看检测结果。

我是使用openmmlab框架来训练hrnet的,所以以下推理代码中也会调用相关库。若创建环境有问腿可以直接使用我配好的环境直接创建,下载我上传的环境配置文件,直接使用命令conda env create -f environment.yml即可使用这个文件来创建相同的环境。

有的可能包含特定的通道(如conda-forge),确保在你的电脑中也配置了这些通道,或者单独pip install某个包。

python 复制代码
import os
import json
import numpy as np
from mmpose.apis import init_model, inference_topdown
import cv2

# 配置文件路径和模型路径
config_file = '/home/.../pose_td-hm_hrnet-w48_8xb32-210e_Squat.py'
checkpoint_file = '/home/.../best_coco_AP_epoch_290.pth'

# 初始化姿态估计模型
pose_model = init_model(config_file, checkpoint_file, device='cuda:0')

# 图像文件夹路径
image_folder = ''
output_path = ''

# 处理每一张图像
for image_file in os.listdir(image_folder):
    if os.path.splitext(image_file)[-1] != ".jpg":
        continue
    result = []
    image_path = os.path.join(image_folder, image_file)

    image = cv2.imread(image_path)

    # 确保图像加载成功
    if image is None:
        print(f"Failed to load image {image_file}")
        continue

    # 获取图像宽高
    height, width = image.shape[:2]

    # 假设每张图像中只有一个人,设置边界框为整张图像大小
    person_results = np.array([[0, 0, width, height]])

    # 打印 person_results 确认内容
    print(f"Processing image: {image_file}")
    print(f"Person results: {person_results}")

    # 推理得到关键点坐标
    try:
        pose_results = inference_topdown(pose_model, image, person_results, bbox_format='xyxy')
    except KeyError as e:
        print(f"Error processing image {image_file}: {e}")
        print(f"Person results: {person_results}")
        continue

    # # 打印 pose_results 确认内容
    # print(f"Pose results: {pose_results}")

    # 提取关键点坐标并检查是否检测出17个关键点
    if len(pose_results) > 0 and pose_results[0].pred_instances.keypoints.shape[1] == 17:
        keypoints = pose_results[0].pred_instances.keypoints
        result = {
            'image': image_file,
            'scaleratio': 1,
            'keypoints': keypoints.tolist()
        }
        # results.append(result)
    else:
        print(f"No valid pose results for image {image_file}")

    for key in range(len(result["keypoints"][0])):
        result["keypoints"][0][key].append(1.0)
    # 将结果保存为 JSON 文件
    output_file = os.path.join(output_path, os.path.splitext(image_file)[0] + ".json")
    with open(output_file, 'w') as f:
        json.dump(result, f, indent=4)
    print(f"Results saved to {output_file}")

以上是将检测结果保存为关键点信息,json内容如图:

分为3部分:imagepath为图片的绝对路径;scaleratio为标注时在标注工具的放大倍数(标注工具后面会写);keypoints则为17个关键点的信息,每3个值为一个点,分别是点的xy值及是否被遮挡,若为1则该点没被遮挡,若为2则被遮挡(比如手放膝盖上则膝盖的点被遮挡)。关键点的顺序从鼻子开始到脚踝,从左到右,如图:

如果想要直接可视化检测结果,可以直接在图上画点,将上面代码的最后一段保存为json的代码替换为以下代码:

python 复制代码
for point in result["keypoints"][0]:
    cv2.circle(image, (int(point[0]), int(point[1])), 5, (0, 0, 255), -1)
save_path = os.path.join(output_path, os.path.splitext(image_file)[0] + "_point.jpg")
cv2.imwrite(save_path, image)
print(save_path)
相关推荐
粟悟饭&龟波功8 分钟前
【论文精读】DeepSeek-OCR:探索视觉 - 文本压缩的新范式
人工智能
机器之心11 分钟前
刚刚,Kimi开源新架构,开始押注线性注意力
人工智能·openai
IT_陈寒12 分钟前
JavaScript 性能优化:3个V8引擎隐藏技巧让你的代码提速50%
前端·人工智能·后端
故事挺秃然36 分钟前
NLP模型优化
人工智能·自然语言处理·nlp
文火冰糖的硅基工坊44 分钟前
[人工智能-大模型-78]:模型层技术 - 深度神经网络的网络架构的演进,这不仅是一条技术路线图,更是一部 “机器如何逐步逼近人类认知方式” 的进化史诗。
人工智能·架构·dnn
IT古董1 小时前
【第六章:项目实战之推荐/广告系统】3.精排算法-(2)精排算法模型精讲: DNN、deepFM、ESMM、PLE、MMOE算法精讲与实现- DNN 精排模型
人工智能·神经网络·dnn
IT古董1 小时前
【第六章:项目实战之推荐/广告系统】3.精排算法-(2)精排算法模型精讲: DNN、deepFM、ESMM、PLE、MMOE算法精讲与实现- PLE 模型
人工智能·神经网络·dnn
Serverless 社区1 小时前
吉利汽车携手阿里云函数计算,打造新一代 AI 座舱推理引擎
人工智能·阿里云·汽车
网上邻居YY1 小时前
Arcgis表格数据导入+可视化--小白教程(以景观多样性指数为例)
图像处理·计算机视觉·arcgis·excel转表
没有梦想的咸鱼185-1037-16631 小时前
AI大模型支持下的:ArcGIS数据处理、空间分析、可视化及多案例综合应用
人工智能·arcgis·chatgpt·数据分析