Update gradio_utils/utils.py

gradio_utils/utils.py

@@ -37,6 +37,7 @@ def adj_mx_from_edges(num_pts, skeleton, device='cuda', normalization_fix=True):
     return adj
 
 
+@spaces.GPU
 def plot_results(support_img, query_img, support_kp, support_w, query_kp, query_w,
                  skeleton=None, prediction=None, radius=6, in_color=None,
                  original_skeleton=None, img_alpha=0.6, target_keypoints=None):
@@ -126,13 +127,13 @@ def plot_results(support_img, query_img, support_kp, support_w, query_kp, query_
         plt.subplots_adjust(0, 0, 1, 1, 0, 0)
         return plt
 
-
+@spaces.GPU
 def process(query_img, state,
             cfg_path='configs/test/1shot_split1.py',
             checkpoint_path='ckpt/1shot_split1.pth'):
     cfg = Config.fromfile(cfg_path)
-    width, height, _ = state['original_support_image'].shape
-    kp_src_np = np.array(state['kp_src']).copy().astype(np.float32)
+    width, height, _ = state['images']['image_orig'].shape
+    kp_src_np = np.array(state['points']).copy().astype(np.float32)
     kp_src_np[:, 0] = kp_src_np[:, 0] / (width // 4) * cfg.model.encoder_config.img_size
     kp_src_np[:, 1] = kp_src_np[:, 1] / (height // 4) * cfg.model.encoder_config.img_size
     kp_src_np = np.flip(kp_src_np, 1).copy()
@@ -226,7 +227,7 @@ def process(query_img, state,
 
 def update_examples(support_img, posed_support, query_img, state, r=0.015, width=0.02):
     state['color_idx'] = 0
-    state['original_support_image'] = np.array(support_img)[:, :, ::-1].copy()
+    state['images']['image_orig'] = np.array(support_img)[:, :, ::-1].copy()
     support_img, posed_support, _ = set_query(support_img, state, example=True)
     w, h = support_img.size
     draw_pose = ImageDraw.Draw(support_img)