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