Update app.py

app.py

@@ -1,4 +1,120 @@
 import gradio as gr
+from PIL import Image
+from torchvision import models, transforms
+import gradio as gr
+import subprocess
+import torch
+import cv2
+import numpy as np
+from models.with_mobilenet import PoseEstimationWithMobileNet
+from modules.keypoints import extract_keypoints, group_keypoints
+from modules.load_state import load_state
+from modules.pose import Pose
+import demo
+from recon import reconWrapper
+
+
+def get_rect(net, image, height_size):
+    net = net.eval()
+
+    stride = 8
+    upsample_ratio = 4
+    num_keypoints = Pose.num_kpts
+    previous_poses = []
+    delay = 33
+    image = image[0]
+    rect_path = image.replace('.%s' % (image.split('.')[-1]), '_rect.txt')
+    print('Processing', image)
+    img = cv2.imread(image, cv2.IMREAD_COLOR)
+    orig_img = img.copy()
+    heatmaps, pafs, scale, pad = demo.infer_fast(
+        net, img, height_size, stride, upsample_ratio, cpu=False)
+
+    total_keypoints_num = 0
+    all_keypoints_by_type = []
+    for kpt_idx in range(num_keypoints):  # 19th for bg
+        total_keypoints_num += extract_keypoints(
+            heatmaps[:, :, kpt_idx], all_keypoints_by_type, total_keypoints_num)
+
+    pose_entries, all_keypoints = group_keypoints(
+        all_keypoints_by_type, pafs)
+    for kpt_id in range(all_keypoints.shape[0]):
+        all_keypoints[kpt_id, 0] = (
+            all_keypoints[kpt_id, 0] * stride / upsample_ratio - pad[1]) / scale
+        all_keypoints[kpt_id, 1] = (
+            all_keypoints[kpt_id, 1] * stride / upsample_ratio - pad[0]) / scale
+    current_poses = []
+
+    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('int')
+            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('int')
+            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])
+
+    np.savetxt(rect_path, np.array(rects), fmt='%d')
+    print('Cropping boxes are saved at', rect_path)
+    print(rect_path[0:7] +
+          'apps/'+rect_path[7:])
+
+
+def run_simple_test():
+    resolution = str(256)
+    start_id = -1
+    end_id = -1
+    cmd = ['--dataroot', 'pifuhd/sample_images', '--results_path', './results',
+           '--loadSize', '1024', '--resolution', resolution, '--load_netMR_checkpoint_path',
+           './checkpoints/pifuhd.pt',
+           '--start_id', '%d' % start_id, '--end_id', '%d' % end_id]
+    mesh_path = reconWrapper(cmd, True)
+    print('Mesh is saved at', mesh_path)
+    return mesh_path
+
+
+def predict(image):
+    # Save the input image to a file
+    image_path = 'input_image.png'
+    cv2.imwrite(image_path, cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
+
+    net = PoseEstimationWithMobileNet()
+    checkpoint = torch.load(
+        'pifuhd/checkpoint_iter_370000.pth', map_location='cpu')
+    load_state(net, checkpoint)
+
+    get_rect(net.cuda(), [image_path], 512)
+    mesh_path = run_simple_test()
+    return mesh_path
 
 # Create the Gradio interface
 iface = gr.Interface(