Spaces:
Runtime error
Runtime error
File size: 1,722 Bytes
b9cc655 |
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 |
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()
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_pose(img, img_path):
save_image(img, img_path)
result = mmpose_coco(img_path)
keypoints = result[0].pred_instances['keypoints'][0]
# Create a dictionary to store keypoints and their names
keypoints_data = {
'keypoints': 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 (img, keypoints_data)
def mmpose_coco(img_path,
config_file = 'mmpose/td-hm_hrnet-w48_8xb32-210e_coco-256x192.py',
checkpoint_file = 'mmpose/td-hm_hrnet-w48_8xb32-210e_coco-256x192-0e67c616_20220913.pth'):
device = torch.cuda.current_device() if torch.cuda.is_available() else 'cpu'
# coco keypoints:
# https://github.com/open-mmlab/mmpose/blob/master/mmpose/datasets/datasets/top_down/topdown_coco_dataset.py#L28
model = init_model(config_file, checkpoint_file, device=device)
results = inference_topdown(model, img_path)
return results
|