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)
相关推荐
gptplusplus19 分钟前
AI智能体(Agent):从“辅助决策”到“自主行动”,重新定义下一个商业时代
人工智能
别忘了微笑啊20 分钟前
hCaptcha 图像识别 API 对接说明
人工智能
大千AI助手42 分钟前
灾难性遗忘:神经网络持续学习的核心挑战与解决方案
人工智能·深度学习·神经网络·大模型·llm·持续学习·灾难性遗忘
gotouniverse43 分钟前
之前自学RAG时做的调研
人工智能
新智元1 小时前
刚刚,英伟达祭出下一代 GPU!狂飙百万 token 巨兽,投 1 亿爆赚 50 亿
人工智能·openai
霍格沃兹_测试1 小时前
从零开始搭建Qwen智能体:新手也能轻松上手指南
人工智能
SmartJavaAI1 小时前
Java调用Whisper和Vosk语音识别(ASR)模型,实现高效实时语音识别(附源码)
java·人工智能·whisper·语音识别
山东小木1 小时前
JBoltAI需求分析大师:基于SpringBoot的大模型智能需求文档生成解决方案
人工智能·spring boot·后端·需求分析·jboltai·javaai·aigs
君名余曰正则1 小时前
【竞赛系列】机器学习实操项目08——全球城市计算AI挑战赛(数据可视化分析)
人工智能·机器学习·信息可视化
算家计算1 小时前
一张图+一段音频=电影级视频!阿里Wan2.2-S2V-14B本地部署教程:实现丝滑口型同步
人工智能·开源·aigc