如何解析姿态估计 tflite 模型的热图输出?

我开始使用姿态估计 tflite 模型来获取人类的关键点。

https://www.tensorflow.org/lite/models/pose_estimation/overview

我已经开始使用单张图片或一个人并调用模型:

img = cv.imread('photos\standing\\3.jpg')img = tf.reshape(tf.image.resize(img, [257,257]), [1,257,257,3])model = tf.lite.Interpreter('models\posenet_mobilenet_v1_100_257x257_multi_kpt_stripped.tflite')model.allocate_tensors()input_details = model.get_input_details()output_details = model.get_output_details()floating_model = input_details[0]['dtype'] == np.float32if floating_model:    img = (np.float32(img) - 127.5) / 127.5model.set_tensor(input_details[0]['index'], img)model.invoke()output_data =  model.get_tensor(output_details[0]['index'])# o()offset_data = model.get_tensor(output_details[1]['index'])results = np.squeeze(output_data)offsets_results = np.squeeze(offset_data)print("output shape: {}".format(output_data.shape))np.savez('sample3.npz', results, offsets_results)

但是我在正确解析输出以获取每个身体部位的坐标/置信度方面遇到了困难。是否有人有解释这个模型结果的 Python 示例?(例如:使用它们将关键点映射回原始图像)

我的代码(来自一个类的一个片段,该类基本上直接从模型输出中获取 np 数组):

def get_keypoints(self, data):        height, width, num_keypoints = data.shape        keypoints = []        for keypoint in range(0, num_keypoints):            maxval = data[0][0][keypoint]            maxrow = 0            maxcol = 0            for row in range(0, width):                for col in range(0,height):                    if data[row][col][keypoint] > maxval:                        maxrow = row                        maxcol = col                        maxval = data[row][col][keypoint]            keypoints.append(KeyPoint(keypoint, maxrow, maxcol, maxval))            # keypoints = [Keypoint(x,y,z) for x,y,z in ]        return keypointsdef get_image_coordinates_from_keypoints(self, offsets):        height, width, depth = (257,257,3)        # [(x,y,confidence)]        coords = [{ 'point': k.body_part,                    'location': (k.x / (width - 1)*width + offsets[k.y][k.x][k.index],                   k.y / (height - 1)*height + offsets[k.y][k.x][k.index]),                    'confidence': k.confidence}                 for k in self.keypoints]        return coords

在将索引与部位匹配后,我的输出是: enter image description here

这里的一些坐标是负数,这显然是不正确的。我的错误在哪里?


回答:

import numpy as np

对于输出热图和偏移量的姿态估计模型,可以通过以下步骤获得所需的点:

  1. 对热图执行 Sigmoid 操作:

    scores = sigmoid(heatmaps)

  2. 每个姿态的关键点通常由一个二维矩阵表示,该矩阵中的最大值与模型认为该点在输入图像中的位置有关。使用 argmax2D 获得每个矩阵中该值的 x 和 y 索引,该值本身代表置信值:

    x,y = np.unravel_index(np.argmax(scores[:,:,keypointindex]), scores[:,:,keypointindex].shape)confidences = scores[x,y,keypointindex]

  3. 使用该 x,y 来查找用于计算关键点最终位置的相应偏移向量:

    offset_vector = (offsets[y,x,keypointindex], offsets[y,x,num_keypoints+keypointindex])

  4. 在获得关键点坐标和偏移量后,可以通过以下方式计算关键点的最终位置:

    image_positions = np.add(np.array(heatmap_positions) * output_stride, offset_vectors)

参见 这里 以了解如何获取输出步长,如果您还没有的话。tflite 姿态估计的输出步长为 32。

一个从姿态估计模型输出中提取关键点的函数。不包括 KeyPoint

def get_keypoints(self, heatmaps, offsets, output_stride=32):        scores = sigmoid(heatmaps)        num_keypoints = scores.shape[2]        heatmap_positions = []        offset_vectors = []        confidences = []        for ki in range(0, num_keypoints ):            x,y = np.unravel_index(np.argmax(scores[:,:,ki]), scores[:,:,ki].shape)            confidences.append(scores[x,y,ki])            offset_vector = (offsets[y,x,ki], offsets[y,x,num_keypoints+ki])            heatmap_positions.append((x,y))            offset_vectors.append(offset_vector)        image_positions = np.add(np.array(heatmap_positions) * output_stride, offset_vectors)        keypoints = [KeyPoint(i, pos, confidences[i]) for i, pos in enumerate(image_positions)]        return keypoints

关键点类:

PARTS = {    0: 'NOSE',    1: 'LEFT_EYE',    2: 'RIGHT_EYE',    3: 'LEFT_EAR',    4: 'RIGHT_EAR',    5: 'LEFT_SHOULDER',    6: 'RIGHT_SHOULDER',    7: 'LEFT_ELBOW',    8: 'RIGHT_ELBOW',    9: 'LEFT_WRIST',    10: 'RIGHT_WRIST',    11: 'LEFT_HIP',    12: 'RIGHT_HIP',    13: 'LEFT_KNEE',    14: 'RIGHT_KNEE',    15: 'LEFT_ANKLE',    16: 'RIGHT_ANKLE'}class KeyPoint():    def __init__(self, index, pos, v):        x, y = pos        self.x = x        self.y = y        self.index = index        self.body_part = PARTS.get(index)        self.confidence = v    def point(self):        return int(self.y), int(self.x)    def to_string(self):        return 'part: {} location: {} confidence: {}'.format(            self.body_part, (self.x, self.y), self.confidence)

Related Posts

使用LSTM在Python中预测未来值

这段代码可以预测指定股票的当前日期之前的值,但不能预测…

如何在gensim的word2vec模型中查找双词组的相似性

我有一个word2vec模型,假设我使用的是googl…

dask_xgboost.predict 可以工作但无法显示 – 数据必须是一维的

我试图使用 XGBoost 创建模型。 看起来我成功地…

ML Tuning – Cross Validation in Spark

我在https://spark.apache.org/…

如何在React JS中使用fetch从REST API获取预测

我正在开发一个应用程序,其中Flask REST AP…

如何分析ML.NET中多类分类预测得分数组?

我在ML.NET中创建了一个多类分类项目。该项目可以对…

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注