Spaces:
Runtime error
Runtime error
File size: 2,552 Bytes
b9cc655 c205b8f b9cc655 c205b8f b9cc655 c205b8f b9cc655 c205b8f b9cc655 c205b8f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
import cv2
import numpy as np
import torch
from mmpose.apis import inference_topdown, init_model
from mmpose.utils import register_all_modules
register_all_modules()
# coco keypoints:
# https://github.com/open-mmlab/mmpose/blob/master/mmpose/datasets/datasets/top_down/topdown_coco_dataset.py#L28
model = init_model('mmpose/td-hm_hrnet-w48_8xb32-210e_coco-256x192.py', 'mmpose/td-hm_hrnet-w48_8xb32-210e_coco-256x192-0e67c616_20220913.pth', device=torch.cuda.current_device() if torch.cuda.is_available() else 'cpu')
def save_image(img, img_path):
# Convert PIL image to OpenCV image
img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
# Save OpenCV image
cv2.imwrite(img_path, img)
def predict_poses(front_img, front_img_path, side_img, side_img_path):
save_image(front_img, front_img_path)
save_image(side_img, side_img_path)
mmpose_result = mmpose_coco(front_img_path, side_img_path)
front_result = mmpose_result[0]
side_result = mmpose_result[1]
front_keypoints = front_result[0].pred_instances['keypoints'][0]
front_keypoints_data = {
'keypoints': front_keypoints.tolist(),
'keypoint_names': [
'nose',
'left_eye',
'right_eye',
'left_ear',
'right_ear',
'left_shoulder',
'right_shoulder',
'left_elbow',
'right_elbow',
'left_wrist',
'right_wrist',
'left_hip',
'right_hip',
'left_knee',
'right_knee',
'left_ankle',
'right_ankle'
]
}
side_keypoints = side_result[0].pred_instances['keypoints'][0]
side_keypoints_data = {
'keypoints': side_keypoints.tolist(),
'keypoint_names': [
'nose',
'left_eye',
'right_eye',
'left_ear',
'right_ear',
'left_shoulder',
'right_shoulder',
'left_elbow',
'right_elbow',
'left_wrist',
'right_wrist',
'left_hip',
'right_hip',
'left_knee',
'right_knee',
'left_ankle',
'right_ankle'
]
}
return ((front_img, front_keypoints_data), (side_img, side_keypoints_data))
def mmpose_coco(front_img_path,
side_img_path):
front_results = inference_topdown(model, front_img_path)
side_results = inference_topdown(model, side_img_path)
return (front_results, side_results)
|