init space

.gitignore

@@ -0,0 +1,3 @@
+*.pyc
+__pycache__
+.idea/

app.py

@@ -0,0 +1,59 @@
+import cv2
+import gradio as gr
+import numpy as np
+from human_pose_estimator import PoseEstimator
+
+from pose_estimator import rect
+
+pose_estimator = PoseEstimator("cpu")
+
+
+def get_box(image):
+    image_box, _ = rect(pose_estimator, image)
+    return image_box
+
+
+def predict(img: np.ndarray):
+
+    poses, _, _ = pose_estimator.get_poses(img, 512)
+
+    for pose in poses:
+        pose.draw(img)
+        cv2.rectangle(img, (pose.bbox[0], pose.bbox[1]),
+                      (pose.bbox[0] + pose.bbox[2], pose.bbox[1] + pose.bbox[3]), (0, 255, 0))
+
+    return img
+
+
+footer = r"""
+<center>
+<b>
+Demo for <a href='https://github.com/Daniil-Osokin/lightweight-human-pose-estimation.pytorch'>Lightweight OpenPose</a>
+</b>
+</center>
+"""
+
+with gr.Blocks(title="OpenPose") as app:
+    gr.HTML("<center><h1>Human Pose Estimation Pytorch</h1></center>")
+    gr.HTML("<center><h3>Real-time 2D Multi-Person Pose Estimation on CPU: Lightweight OpenPose</h3></center>")
+    with gr.Row(equal_height=False):
+        with gr.Column():
+            input_img = gr.Image(type="numpy", label="Input image")
+            # input_img = gr.Video(source="webcam")
+            run_btn = gr.Button(variant="primary")
+        with gr.Column():
+            output_img = gr.Image(type="pil", label="Output image")
+            gr.ClearButton(components=[input_img, output_img], variant="stop")
+
+    run_btn.click(predict, [input_img], [output_img])
+
+    with gr.Row():
+        blobs = [[f"examples/{x:02d}.jpg"] for x in range(1, 4)]
+        examples = gr.Dataset(components=[input_img], samples=blobs)
+        examples.click(lambda x: x[0], [examples], [input_img])
+
+    with gr.Row():
+        gr.HTML(footer)
+
+app.launch(share=False, debug=True, show_error=True)
+app.queue()

pose_estimator.py

@@ -0,0 +1,55 @@
+import cv2
+import numpy as np
+from human_pose_estimator.modules.pose import Pose
+
+
+def rect(estimator, image, height_size=512):
+
+    num_keypoints = Pose.num_kpts
+
+    img = cv2.imread(image, cv2.IMREAD_COLOR)
+
+    _, pose_entries, all_keypoints = estimator.get_poses(img, height_size)
+
+    rects = []
+    for n in range(len(pose_entries)):
+        if len(pose_entries[n]) == 0:
+            continue
+        pose_keypoints = np.ones((num_keypoints, 2), dtype=np.int32) * -1
+
+        valid_keypoints = []
+        for kpt_id in range(num_keypoints):
+            if pose_entries[n][kpt_id] != -1.0:  # keypoint was found
+                pose_keypoints[kpt_id, 0] = int(all_keypoints[int(pose_entries[n][kpt_id]), 0])
+                pose_keypoints[kpt_id, 1] = int(all_keypoints[int(pose_entries[n][kpt_id]), 1])
+                valid_keypoints.append([pose_keypoints[kpt_id, 0], pose_keypoints[kpt_id, 1]])
+        valid_keypoints = np.array(valid_keypoints)
+
+        if pose_entries[n][10] != -1.0 or pose_entries[n][13] != -1.0:
+            pmin = valid_keypoints.min(0)
+            pmax = valid_keypoints.max(0)
+
+            center = (0.5 * (pmax[:2] + pmin[:2])).astype(np.int32)
+            radius = int(0.65 * max(pmax[0] - pmin[0], pmax[1] - pmin[1]))
+        elif pose_entries[n][10] == -1.0 and pose_entries[n][13] == -1.0 and pose_entries[n][8] != -1.0 and \
+                pose_entries[n][11] != -1.0:
+            # if leg is missing, use pelvis to get cropping
+            center = (0.5 * (pose_keypoints[8] + pose_keypoints[11])).astype(np.int32)
+            radius = int(1.45 * np.sqrt(((center[None, :] - valid_keypoints) ** 2).sum(1)).max(0))
+            center[1] += int(0.05 * radius)
+        else:
+            center = np.array([img.shape[1] // 2, img.shape[0] // 2])
+            radius = max(img.shape[1] // 2, img.shape[0] // 2)
+
+        x1 = center[0] - radius
+        y1 = center[1] - radius
+
+        rects.append([x1, y1, 2 * radius, 2 * radius])
+
+    for (x, y, w, h) in rects:
+        cv2.rectangle(img, (x, y),
+                      (x + w, y + h),
+                      (0, 0, 255), 2)
+
+    return img, rects
+

requerimets.txt

@@ -0,0 +1,4 @@
+torch>=1.4.0
+numpy
+opencv-python
+human-pose-estimator