Spaces:
Running
Running
| import tensorflow as tf | |
| import numpy as np | |
| from PIL import ImageDraw, ImageFont | |
| from huggingface_hub import snapshot_download | |
| # Keypoint labels for MoveNet (17 human body parts) | |
| KEYPOINT_LABELS = { | |
| 'nose': 0, 'left_eye': 1, 'right_eye': 2, 'left_ear': 3, 'right_ear': 4, | |
| 'left_shoulder': 5, 'right_shoulder': 6, 'left_elbow': 7, 'right_elbow': 8, | |
| 'left_wrist': 9, 'right_wrist': 10, 'left_hip': 11, 'right_hip': 12, | |
| 'left_knee': 13, 'right_knee': 14, 'left_ankle': 15, 'right_ankle': 16 | |
| } | |
| # Stylish connection colors for the skeleton bones | |
| SKELETON_EDGES = { | |
| (0, 1): '#FF66CC', (0, 2): '#66FFFF', (1, 3): '#FF66CC', (2, 4): '#66FFFF', | |
| (0, 5): '#FF99CC', (0, 6): '#99FFFF', (5, 7): '#FF6699', (7, 9): '#FF3366', | |
| (6, 8): '#66CCCC', (8, 10): '#33CCCC', (5, 6): '#CCCC00', (5, 11): '#FF9966', | |
| (6, 12): '#66FF99', (11, 12): '#999900', (11, 13): '#FF6600', (13, 15): '#FF3300', | |
| (12, 14): '#00CC99', (14, 16): '#009966' | |
| } | |
| def process_keypoints(prediction, img_height, img_width, confidence=0.12): | |
| all_joints = [] | |
| all_bones = [] | |
| instance_count = prediction.shape[1] | |
| for i in range(instance_count): | |
| x_coords = prediction[0, i, :, 1] * img_width | |
| y_coords = prediction[0, i, :, 0] * img_height | |
| scores = prediction[0, i, :, 2] | |
| labels = list(KEYPOINT_LABELS.keys()) | |
| keypoints = np.stack([labels, x_coords, y_coords], axis=-1) | |
| visible_kpts = keypoints[scores > confidence] | |
| all_joints.append(visible_kpts) | |
| for (a, b), color in SKELETON_EDGES.items(): | |
| if scores[a] > confidence and scores[b] > confidence: | |
| segment = np.array([[x_coords[a], y_coords[a]], [x_coords[b], y_coords[b]]]) | |
| all_bones.append((segment, color)) | |
| return np.concatenate(all_joints, axis=0), all_bones | |
| def draw_bones(image, keypoints): | |
| draw = ImageDraw.Draw(image) | |
| font = ImageFont.load_default() | |
| joints, bones = process_keypoints(keypoints, image.height, image.width) | |
| for bone, color in bones: | |
| draw.line((*bone[0], *bone[1]), fill=color, width=3) | |
| for label, x, y in joints: | |
| cx, cy = float(x), float(y) | |
| radius = 4 | |
| draw.ellipse([(cx - radius, cy - radius), (cx + radius, cy + radius)], fill="#FF0000", outline="#222222") | |
| draw.text((cx + 5, cy - 5), label, font=font, fill="#0000CC") | |
| return joints | |
| def movenet(image): | |
| model_path = snapshot_download("lukassso/movenet-myking") | |
| model = tf.saved_model.load(model_path).signatures["serving_default"] | |
| image = tf.cast(image, tf.int32) | |
| result = model(image) | |
| return result["output_0"].numpy() | |