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)
相关推荐
_一两风2 分钟前
用 n8n 自动化生成「每日科技热点速递」:从 RSS 到 AI 写作全流程实战(小白必看)
人工智能·rss·deepseek
极昆仑智慧8 分钟前
OpenAI推出了支持人工智能的浏览器ChatGPT Atlas
人工智能·chatgpt
hunteritself11 分钟前
阿里千问上线记忆,Manus 1.5 全栈升级,ChatGPT 将推成人模式!| AI Weekly 10.13-10.19
大数据·人工智能·深度学习·机器学习·chatgpt
姓刘的哦15 分钟前
基于线程池的配电房图像检测
人工智能·计算机视觉·目标跟踪
机器之心19 分钟前
谷歌最强AI,被港科大开源超了?让海外创作者喊出「King Bomb」的P图大杀器来了
人工智能·openai
DKunYu26 分钟前
PyTorch入门
人工智能·pytorch·python·深度学习
机器之心31 分钟前
谷歌重大突破!量子计算首次可验证,登《Nature》封面
人工智能·openai
gustt37 分钟前
每天一杯科技资讯咖啡,用 n8n 喂给 AI
人工智能·node.js
技术闲聊DD37 分钟前
深度学习(8)- PyTorch 数据处理与加载
人工智能·pytorch·深度学习
Juchecar39 分钟前
从LLM训练和推理角度,区分 “知识”与“智慧”
人工智能