show pose body lables

.gitignore

@@ -1,4 +1,4 @@
 .idea/
 __pycache__/
-requirements.txt
-push_model.py
+playground.py
+

app.py

@@ -11,6 +11,7 @@ def predict(image: PIL.Image):
     input_size = 256
     size = (1280, 1280)
     image = PIL.ImageOps.fit(image, size, PIL.Image.LANCZOS)
+    # image = PIL.ImageOps.contain(image, size)
     image_tf = tf.keras.preprocessing.image.img_to_array(image)
     # Resize and pad the image to keep the aspect ratio and fit the expected size.
     input_image = tf.expand_dims(image_tf, axis=0)
@@ -18,8 +19,9 @@ def predict(image: PIL.Image):
     keypoints = movenet(input_image)
     keypoints = np.array(keypoints)
     image = tf.keras.preprocessing.image.array_to_img(image_tf)
-    draw_bones(image, keypoints)
-    return image
+    joints = draw_bones(image, keypoints)
+    points = [f"{x}#{y}" for p, x, y in joints]
+    return image, joints, points
 
 
 footer = r"""
@@ -39,9 +41,18 @@ with gr.Blocks(title="MoveNet") as app:
             run_btn = gr.Button(variant="primary")
         with gr.Column():
             output_img = gr.Image(type="numpy", label="Output image")
-            gr.ClearButton(components=[input_img, output_img], variant="stop")
+            with gr.Accordion("See Positions", open=False):
+                positions = gr.Dataframe(
+                    interactive=True,
+                    headers=["x", "y", "label"],
+                    datatype=["str", "number", "number"],
+                    row_count=16,
+                    col_count=(3, "fixed"),
+                )
+                data = gr.Textbox(label="Positions", lines=17)
+            gr.ClearButton(components=[input_img, output_img, positions, data], variant="stop")
 
-    run_btn.click(predict, [input_img], [output_img])
+    run_btn.click(predict, [input_img], [output_img, positions, data])
 
     with gr.Row():
         blobs = [[f"examples/{x:02d}.jpg"] for x in range(1, 4)]

poser.py

@@ -10,23 +10,31 @@ import PIL.ImageOps
 import numpy as np
 import tensorflow as tf
 from PIL import ImageDraw
+from PIL import ImageFont
 from huggingface_hub import snapshot_download
 
-
 # Dictionary that maps from joint names to keypoint indices.
 KEYPOINT_DICT = {
     '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
+    '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
 }
 
-COLOR_DICT = {
+KEYPOINT_EDGE_INDS_TO_COLOR = {
     (0, 1): 'Magenta',
     (0, 2): 'Cyan',
     (1, 3): 'Magenta',
@@ -48,11 +56,11 @@ COLOR_DICT = {
 }
 
 
-def process_keypoints(keypoints, height, width, threshold=0.22):
+def process_keypoints(keypoints_with_scores, height, width, threshold=0.11):
     """Returns high confidence keypoints and edges for visualization.
 
       Args:
-        keypoints: A numpy array with shape [1, 1, 17, 3] representing
+        keypoints_with_scores: A numpy array with shape [1, 1, 17, 3] representing
           the keypoint coordinates and scores returned from the MoveNet model.
         height: height of the image in pixels.
         width: width of the image in pixels.
@@ -67,52 +75,57 @@ def process_keypoints(keypoints, height, width, threshold=0.22):
       """
     keypoints_all = []
     keypoint_edges_all = []
-    colors = []
-    num_instances, _, _, _ = keypoints.shape
+    num_instances, _, _, _ = keypoints_with_scores.shape
     for idx in range(num_instances):
-        kpts_x = keypoints[0, idx, :, 1]
-        kpts_y = keypoints[0, idx, :, 0]
-        kpts_scores = keypoints[0, idx, :, 2]
-        kpts_absolute_xy = np.stack(
-            [width * np.array(kpts_x), height * np.array(kpts_y)], axis=-1)
-        kpts_above_thresh_absolute = kpts_absolute_xy[
-                                     kpts_scores > threshold, :]
+        kpts_x = keypoints_with_scores[0, idx, :, 1]
+        kpts_y = keypoints_with_scores[0, idx, :, 0]
+        kpts_scores = keypoints_with_scores[0, idx, :, 2]
+        kpts_dict = list(KEYPOINT_DICT.keys())
+        kpts_absolute_xy = np.stack([kpts_dict, width * np.array(kpts_x), height * np.array(kpts_y)], axis=-1)
+        kpts_above_thresh_absolute = kpts_absolute_xy[kpts_scores > threshold, :]
         keypoints_all.append(kpts_above_thresh_absolute)
 
-        for edge_pair, color in COLOR_DICT.items():
-            if (kpts_scores[edge_pair[0]] > threshold and
-                    kpts_scores[edge_pair[1]] > threshold):
-                x_start = kpts_absolute_xy[edge_pair[0], 0]
-                y_start = kpts_absolute_xy[edge_pair[0], 1]
-                x_end = kpts_absolute_xy[edge_pair[1], 0]
-                y_end = kpts_absolute_xy[edge_pair[1], 1]
+        for edge_pair, color in KEYPOINT_EDGE_INDS_TO_COLOR.items():
+            if kpts_scores[edge_pair[0]] > threshold and kpts_scores[edge_pair[1]] > threshold:
+                x_start = kpts_absolute_xy[edge_pair[0], 1]
+                y_start = kpts_absolute_xy[edge_pair[0], 2]
+                x_end = kpts_absolute_xy[edge_pair[1], 1]
+                y_end = kpts_absolute_xy[edge_pair[1], 2]
                 line_seg = np.array([[x_start, y_start], [x_end, y_end]])
-                keypoint_edges_all.append(line_seg)
-                colors.append(color)
+                keypoint_edges_all.append([line_seg, color])
     if keypoints_all:
-        joints = np.concatenate(keypoints_all, axis=0)
+        keypoints_xy = np.concatenate(keypoints_all, axis=0)
     else:
-        joints = np.zeros((0, 17, 2))
+        keypoints_xy = np.zeros((0, 17, 2))
 
     if keypoint_edges_all:
-        bones = np.stack(keypoint_edges_all, axis=0)
+        edges_xy = np.stack(keypoint_edges_all, axis=0)
     else:
-        bones = np.zeros((0, 2, 2))
-    return joints, bones, colors
+        edges_xy = np.zeros((0, 2, 2))
+    return keypoints_xy, edges_xy
 
 
 def draw_bones(pixmap: PIL.Image, keypoints):
     draw = ImageDraw.Draw(pixmap)
-    joints, bones, colors = process_keypoints(keypoints, pixmap.height, pixmap.width)
+    joints, bones = process_keypoints(keypoints, pixmap.height, pixmap.width)
+
+    font = ImageFont.truetype("./Arial.ttf", 22)
+    print(joints)
 
-    for bone, color in zip(bones.tolist(), colors):
+    for bone, color in bones:
+        bone = bone.astype(np.float32)
         draw.line((*bone[0], *bone[1]), fill=color, width=4)
 
     radio = 3
 
-    for c_x, c_y in joints:
+    for label, c_x, c_y in joints:
+        c_x = float(c_x)
+        c_y = float(c_y)
         shape = [(c_x - radio, c_y - radio), (c_x + radio, c_y + radio)]
         draw.ellipse(shape, fill="red", outline="red")
+        draw.text((c_x, c_y), label, font=font, align="left", fill="blue")
+
+    return joints
 
 
 def movenet(image):
@@ -136,6 +149,3 @@ def movenet(image):
     outputs = model(image)
     # Output is a [1, 1, 17, 3] tensor.
     return outputs['output_0'].numpy()
-
-
-