Update gradio_utils/utils.py

gradio_utils/utils.py

@@ -30,7 +30,7 @@ def process_img(support_image, global_state):
     return support_image, global_state
 
 
-def adj_mx_from_edges(num_pts, skeleton, device='cuda', normalization_fix=True):
+def adj_mx_from_edges(num_pts, skeleton, device='cpu', normalization_fix=True):
     adj_mx = torch.empty(0, device=device)
     batch_size = len(skeleton)
     for b in range(batch_size):
@@ -43,16 +43,15 @@ def adj_mx_from_edges(num_pts, skeleton, device='cuda', normalization_fix=True):
     adj = adj_mx + trans_adj_mx * cond - adj_mx * cond
     return adj
 
-@spaces.GPU(duration=30)
 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):
     h, w, c = support_img.shape
     prediction = prediction[-1] * h
     if isinstance(prediction, torch.Tensor):
-        prediction = prediction.cpu().numpy()
+        prediction = prediction.numpy()
     if isinstance(original_skeleton, list):
-        original_skeleton = adj_mx_from_edges(num_pts=100, skeleton=[skeleton]).cpu().numpy()[0]
+        original_skeleton = adj_mx_from_edges(num_pts=100, skeleton=[skeleton]).numpy()[0]
     query_img = (query_img - np.min(query_img)) / (np.max(query_img) - np.min(query_img))
     img = query_img
     w = query_w
@@ -103,8 +102,8 @@ def plot_results(support_img, query_img, support_kp, support_w, query_kp, query_
                                        zorder=1)
                     axes.add_artist(patch)
 
-        plt.axis('off')  # command for hiding the axis.
-        return plt
+    plt.axis('off')  # command for hiding the axis.
+    return plt
 
 
 
@@ -200,7 +199,6 @@ def process(query_img, state,
     vis_s_image = support_img[0].detach().cpu().numpy().transpose(1, 2, 0)
     vis_q_image = q_img[0].detach().cpu().numpy().transpose(1, 2, 0)
     support_kp = kp_src_3d
-    print("Try to Plot")
     out = plot_results(vis_s_image,
                        vis_q_image,
                        support_kp,
@@ -208,11 +206,10 @@ def process(query_img, state,
                        None,
                        vis_s_weight,
                        outputs['skeleton'][1],
-                       torch.tensor(outputs['points']).squeeze(),
+                       torch.tensor(outputs['points']).squeeze().cpu(),
                        original_skeleton=state['skeleton'],
                        img_alpha=1.0,
                        )
-    print("Plot Success!")
     return out