[英]How to parse the heatmap output for the pose estimation tflite model?
我從姿勢估計 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.float32
if floating_model:
img = (np.float32(img) - 127.5) / 127.5
model.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 keypoints
def 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
這里的一些坐標是負數,這不可能是正確的。 我的錯誤在哪里?
import numpy as np
對於輸出熱圖和偏移量的姿勢估計模型。 可以通過以下方式獲得所需的點數:
對熱圖執行 sigmoid 操作:
scores = sigmoid(heatmaps)
這些姿勢的每個關鍵點通常由一個二維矩陣表示,該矩陣中的最大值與模型認為該點在輸入圖像中的位置有關。 使用 argmax2D 獲取每個矩陣中該值的 x 和 y 索引,該值本身表示置信度值:
x,y = np.unravel_index(np.argmax(scores[:,:,keypointindex]), scores[:,:,keypointindex].shape)
confidences = scores[x,y,keypointindex]
x,y = np.unravel_index(np.argmax(scores[:,:,keypointindex]), scores[:,:,keypointindex].shape)
confidences = scores[x,y,keypointindex]
那個 x,y 用於找到對應的偏移向量,用於計算關鍵點的最終位置:
offset_vector = (offsets[y,x,keypointindex], offsets[y,x,num_keypoints+keypointindex])
獲得關鍵點坐標和偏移后,您可以使用 () 計算關鍵點的最終位置:
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)
聲明:本站的技術帖子網頁,遵循CC BY-SA 4.0協議,如果您需要轉載,請注明本站網址或者原文地址。任何問題請咨詢:yoyou2525@163.com.