def extract_points(self, heatmaps, nms_heatmaps):
batch_size, channels_num, h, w = heatmaps.shape
assert batch_size == 1, 'Batch size of 1 only supported'
assert channels_num >= self.num_joints
xs, ys, scores = self.top_k(nms_heatmaps)
masks = scores > self.score_threshold
all_keypoints = []
keypoint_id = 0
for k in range(self.num_joints):
# Filter low-score points.
mask = masks[0, k]
x = xs[0, k][mask].ravel()
y = ys[0, k][mask].ravel()
score = scores[0, k][mask].ravel()
n = len(x)
if n == 0:
all_keypoints.append(np.empty((0, 4), dtype=np.float32))
continue
# Apply quarter offset to improve localization accuracy.
x, y = self.refine(heatmaps[0, k], x, y)
clip(x, 0, w - 1, out=x)
clip(y, 0, h - 1, out=y)
# Pack resulting points.
keypoints = np.empty((n, 4), dtype=np.float32)
keypoints[:, 0] = x
keypoints[:, 1] = y
keypoints[:, 2] = score
keypoints[:, 3] = np.arange(keypoint_id, keypoint_id + n)
keypoint_id += n
all_keypoints.append(keypoints)
return all_keypoints
def __call__(self, heatmaps, nms_heatmaps, pafs):
batch_size, _, h, w = heatmaps.shape
assert batch_size == 1, 'Batch size of 1 only supported'
keypoints = self.extract_points(heatmaps, nms_heatmaps)
pafs = np.transpose(pafs, (0, 2, 3, 1))
if self.delta > 0:
for kpts in keypoints:
kpts[:, :2] += self.delta
clip(kpts[:, 0], 0, w - 1, out=kpts[:, 0])
clip(kpts[:, 1], 0, h - 1, out=kpts[:, 1])
pose_entries, keypoints = self.group_keypoints(keypoints, pafs, pose_entry_size=self.num_joints + 2)
poses, scores = self.convert_to_coco_format(pose_entries, keypoints)
if len(poses) > 0:
poses = np.asarray(poses, dtype=np.float32)
poses = poses.reshape((poses.shape[0], -1, 3))
else:
poses = np.empty((0, 17, 3), dtype=np.float32)
scores = np.empty(0, dtype=np.float32)
return poses, scores