Update gradio_utils/utils.py

gradio_utils/utils.py

@@ -44,7 +44,7 @@ def adj_mx_from_edges(num_pts, skeleton, device='cpu', normalization_fix=True):
     adj = adj_mx + trans_adj_mx * cond - adj_mx * cond
     return adj
 
-def plot_results(support_img, query_img, support_kp, support_w, query_kp, query_w,
+def plot_results(support_img, query_img, 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
@@ -52,7 +52,7 @@ def plot_results(support_img, query_img, support_kp, support_w, query_kp, query_
     if isinstance(prediction, torch.Tensor):
         prediction = prediction.numpy()
     if isinstance(original_skeleton, list):
-        original_skeleton = adj_mx_from_edges(num_pts=100, skeleton=[skeleton]).numpy()[0]
+        original_skeleton = adj_mx_from_edges(num_pts=prediction.shape[0], skeleton=[original_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
@@ -204,12 +204,9 @@ def process(query_img, state,
     support_kp = kp_src_3d
     out = plot_results(vis_s_image,
                        vis_q_image,
-                       support_kp,
                        vis_s_weight,
-                       None,
-                       vis_s_weight,
-                       outputs['skeleton'][1],
-                       torch.tensor(outputs['points']).squeeze().cpu(),
+                       skeleton=outputs['skeleton'][1],
+                       prediction=torch.tensor(outputs['points']).squeeze().cpu(),
                        original_skeleton=state['skeleton'],
                        img_alpha=1.0,
                        )