Spaces:
Running
Running
| import numpy as np | |
| from human_pose_estimator.modules.pose import Pose | |
| def rect(estimator, img: np.ndarray, height_size=512): | |
| num_keypoints = Pose.num_kpts | |
| _, pose_entries, all_keypoints = estimator.get_poses(img, height_size) | |
| rects = [] | |
| for n in range(len(pose_entries)): | |
| if len(pose_entries[n]) == 0: | |
| continue | |
| pose_keypoints = np.ones((num_keypoints, 2), dtype=np.int32) * -1 | |
| valid_keypoints = [] | |
| for kpt_id in range(num_keypoints): | |
| if pose_entries[n][kpt_id] != -1.0: # keypoint was found | |
| pose_keypoints[kpt_id, 0] = int(all_keypoints[int(pose_entries[n][kpt_id]), 0]) | |
| pose_keypoints[kpt_id, 1] = int(all_keypoints[int(pose_entries[n][kpt_id]), 1]) | |
| valid_keypoints.append([pose_keypoints[kpt_id, 0], pose_keypoints[kpt_id, 1]]) | |
| valid_keypoints = np.array(valid_keypoints) | |
| if pose_entries[n][10] != -1.0 or pose_entries[n][13] != -1.0: | |
| pmin = valid_keypoints.min(0) | |
| pmax = valid_keypoints.max(0) | |
| center = (0.5 * (pmax[:2] + pmin[:2])).astype(np.int32) | |
| radius = int(0.65 * max(pmax[0] - pmin[0], pmax[1] - pmin[1])) | |
| elif pose_entries[n][10] == -1.0 and pose_entries[n][13] == -1.0 and pose_entries[n][8] != -1.0 and \ | |
| pose_entries[n][11] != -1.0: | |
| # if leg is missing, use pelvis to get cropping | |
| center = (0.5 * (pose_keypoints[8] + pose_keypoints[11])).astype(np.int32) | |
| radius = int(1.45 * np.sqrt(((center[None, :] - valid_keypoints) ** 2).sum(1)).max(0)) | |
| center[1] += int(0.05 * radius) | |
| else: | |
| center = np.array([img.shape[1] // 2, img.shape[0] // 2]) | |
| radius = max(img.shape[1] // 2, img.shape[0] // 2) | |
| x1 = center[0] - radius | |
| y1 = center[1] - radius | |
| rects.append([x1, y1, 2 * radius, 2 * radius]) | |
| return img, rects | |