Spaces:
Running
Running
File size: 2,647 Bytes
d3e8abf |
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 |
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()
|