Update gradio_utils/utils.py

gradio_utils/utils.py

@@ -1,4 +1,6 @@
+import copy
 import spaces
+import json
 import random
 import collections
 import gradio as gr
@@ -13,14 +15,18 @@ from mmpose.core import wrap_fp16_model
 from mmpose.models import build_posenet
 from torchvision import transforms
 import matplotlib.patheffects as mpe
+
+from EdgeCape import TopDownGenerateTargetFewShot
 from demo import Resize_Pad
 from EdgeCape.models import *
 
 
 def process_img(support_image, global_state):
     global_state['images']['image_orig'] = support_image
-    global_state['images']['image_kp'] = support_image
-    reset_kp(global_state)
+    if global_state["load_example"]:
+        global_state["load_example"] = False
+        return global_state['images']['image_kp'], global_state
+    _, _ = reset_kp(global_state)
     return support_image, global_state
 
 
@@ -28,7 +34,7 @@ def adj_mx_from_edges(num_pts, skeleton, device='cuda', normalization_fix=True):
     adj_mx = torch.empty(0, device=device)
     batch_size = len(skeleton)
     for b in range(batch_size):
-        edges = torch.tensor(skeleton[b])
+        edges = torch.tensor(skeleton[b]).long()
         adj = torch.zeros(num_pts, num_pts, device=device)
         adj[edges[:, 0], edges[:, 1]] = 1
         adj_mx = torch.concatenate((adj_mx, adj.unsqueeze(0)), dim=0)
@@ -37,7 +43,6 @@ 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,
@@ -46,104 +51,79 @@ def plot_results(support_img, query_img, support_kp, support_w, query_kp, query_
     prediction = prediction[-1] * h
     if isinstance(prediction, torch.Tensor):
         prediction = prediction.cpu().numpy()
-    if isinstance(skeleton, list):
-        skeleton = adj_mx_from_edges(num_pts=100, skeleton=[skeleton]).cpu().numpy()[0]
-        original_skeleton = skeleton
-    support_img = (support_img - np.min(support_img)) / (np.max(support_img) - np.min(support_img))
+    if isinstance(original_skeleton, list):
+        original_skeleton = adj_mx_from_edges(num_pts=100, skeleton=[skeleton]).cpu().numpy()[0]
     query_img = (query_img - np.min(query_img)) / (np.max(query_img) - np.min(query_img))
-    error_mask = None
-    for id, (img, w, keypoint, adj) in enumerate(zip([support_img, support_img, query_img],
-                                                     [support_w, support_w, query_w],
-                                                     # [support_kp, query_kp])):
-                                                     [support_kp, support_kp, prediction],
-                                                     [original_skeleton, skeleton, skeleton])):
-        color = in_color
-        f, axes = plt.subplots()
-        plt.imshow(img, alpha=img_alpha)
-
-        # On qeury image plot
-        if id == 2 and target_keypoints is not None:
-            error = np.linalg.norm(keypoint - target_keypoints, axis=-1)
-            error_mask = error > (256 * 0.05)
-
-        for k in range(keypoint.shape[0]):
-            if w[k] > 0:
-                kp = keypoint[k, :2]
-                c = (1, 0, 0, 0.75) if w[k] == 1 else (0, 0, 1, 0.6)
-                if error_mask is not None and error_mask[k]:
-                    c = (1, 1, 0, 0.75)
-                    patch = plt.Circle(kp,
-                                       radius,
-                                       color=c,
-                                       path_effects=[mpe.withStroke(linewidth=8, foreground='black'),
-                                                     mpe.withStroke(linewidth=4, foreground='white'),
-                                                     mpe.withStroke(linewidth=2, foreground='black'),
-                                                     ],
-                                       zorder=260)
-                    axes.add_patch(patch)
-                    axes.text(kp[0], kp[1], k, fontsize=10, color='black', ha="center", va="center", zorder=320, )
-                else:
-                    patch = plt.Circle(kp,
-                                       radius,
-                                       color=c,
-                                       path_effects=[mpe.withStroke(linewidth=2, foreground='black')],
-                                       zorder=200)
-                    axes.add_patch(patch)
-                    axes.text(kp[0], kp[1], k, fontsize=(radius + 4), color='white', ha="center", va="center",
-                              zorder=300,
-                              path_effects=[
-                                  mpe.withStroke(linewidth=max(1, int((radius + 4) / 5)), foreground='black')])
-                # axes.text(kp[0], kp[1], k)
-                plt.draw()
-
-        if adj is not None:
-            # Make max value 6
-            draw_skeleton = adj ** 1
-            max_skel_val = np.max(draw_skeleton)
-            draw_skeleton = draw_skeleton / max_skel_val * 6
-            for i in range(1, keypoint.shape[0]):
-                for j in range(0, i):
-                    if w[i] > 0 and w[j] > 0 and original_skeleton[i][j] > 0:
-                        if color is None:
-                            num_colors = int((skeleton > 0.05).sum() / 2)
-                            color = iter(plt.cm.rainbow(np.linspace(0, 1, num_colors + 1)))
-                            c = next(color)
-                        elif isinstance(color, str):
-                            c = color
-                        elif isinstance(color, collections.Iterable):
-                            c = next(color)
-                        else:
-                            raise ValueError("Color must be a string or an iterable")
-                    if w[i] > 0 and w[j] > 0 and skeleton[i][j] > 0:
-                        width = draw_skeleton[i][j]
-                        stroke_width = width + (width / 3)
-                        patch = plt.Line2D([keypoint[i, 0], keypoint[j, 0]],
-                                           [keypoint[i, 1], keypoint[j, 1]],
-                                           linewidth=width, color=c, alpha=0.6,
-                                           path_effects=[mpe.withStroke(linewidth=stroke_width, foreground='black')],
-                                           zorder=1)
-                        axes.add_artist(patch)
+    img = query_img
+    w = query_w
+    keypoint = prediction
+    adj = skeleton
+    color = None
+    f, axes = plt.subplots()
+    plt.imshow(img, alpha=img_alpha)
+    for k in range(keypoint.shape[0]):
+        if w[k] > 0:
+            kp = keypoint[k, :2]
+            c = (1, 0, 0, 0.75) if w[k] == 1 else (0, 0, 1, 0.6)
+            patch = plt.Circle(kp,
+                               radius,
+                               color=c,
+                               path_effects=[mpe.withStroke(linewidth=2, foreground='black')],
+                               zorder=200)
+            axes.add_patch(patch)
+            axes.text(kp[0], kp[1], k, fontsize=(radius + 4), color='white', ha="center", va="center",
+                      zorder=300,
+                      path_effects=[
+                          mpe.withStroke(linewidth=max(1, int((radius + 4) / 5)), foreground='black')])
+            plt.draw()
+
+    if adj is not None:
+        max_skel_val = np.max(adj)
+        draw_skeleton = adj / max_skel_val * 6
+        for i in range(1, keypoint.shape[0]):
+            for j in range(0, i):
+                if w[i] > 0 and w[j] > 0 and original_skeleton[i][j] > 0:
+                    if color is None:
+                        num_colors = int((adj > 0.05).sum() / 2)
+                        color = iter(plt.cm.rainbow(np.linspace(0, 1, num_colors + 1)))
+                        c = next(color)
+                    elif isinstance(color, str):
+                        c = color
+                    elif isinstance(color, collections.Iterable):
+                        c = next(color)
+                    else:
+                        raise ValueError("Color must be a string or an iterable")
+                if w[i] > 0 and w[j] > 0 and adj[i][j] > 0:
+                    width = draw_skeleton[i][j]
+                    stroke_width = width + (width / 3)
+                    patch = plt.Line2D([keypoint[i, 0], keypoint[j, 0]],
+                                       [keypoint[i, 1], keypoint[j, 1]],
+                                       linewidth=width, color=c, alpha=0.6,
+                                       path_effects=[mpe.withStroke(linewidth=stroke_width, foreground='black')],
+                                       zorder=1)
+                    axes.add_artist(patch)
 
         plt.axis('off')  # command for hiding the axis.
-        plt.subplots_adjust(0, 0, 1, 1, 0, 0)
         return plt
 
 @spaces.GPU(duration=30)
 def process(query_img, state,
             cfg_path='configs/test/1shot_split1.py',
             checkpoint_path='ckpt/1shot_split1.pth'):
+    print(state)
+    device = print_memory_usage()
     cfg = Config.fromfile(cfg_path)
-    width, height, _ = state['images']['image_orig'].shape
+    width, height, _ = np.array(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()
+    kp_src_np[:, 0] = kp_src_np[:, 0] / width * 256
+    kp_src_np[:, 1] = kp_src_np[:, 1] / height * 256
+    kp_src_np = kp_src_np.copy()
     kp_src_tensor = torch.tensor(kp_src_np).float()
     preprocess = transforms.Compose([
         transforms.ToTensor(),
         transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
-        Resize_Pad(cfg.model.encoder_config.img_size,
-                   cfg.model.encoder_config.img_size)])
+        Resize_Pad(256, 256)
+    ])
 
     if len(state['skeleton']) == 0:
         state['skeleton'] = [(0, 0)]
@@ -154,8 +134,7 @@ def process(query_img, state,
     # Create heatmap from keypoints
     genHeatMap = TopDownGenerateTargetFewShot()
     data_cfg = cfg.data_cfg
-    data_cfg['image_size'] = np.array([cfg.model.encoder_config.img_size,
-                                       cfg.model.encoder_config.img_size])
+    data_cfg['image_size'] = np.array([256, 256])
     data_cfg['joint_weights'] = None
     data_cfg['use_different_joint_weights'] = False
     kp_src_3d = torch.cat(
@@ -172,22 +151,23 @@ def process(query_img, state,
         torch.tensor(target_weight_s).float()[None])
 
     data = {
-        'img_s': [support_img],
-        'img_q': q_img,
-        'target_s': [target_s],
-        'target_weight_s': [target_weight_s],
+        'img_s': [support_img.to(device)],
+        'img_q': q_img.to(device),
+        'target_s': [target_s.to(device)],
+        'target_weight_s': [target_weight_s.to(device)],
         'target_q': None,
         'target_weight_q': None,
         'return_loss': False,
         'img_metas': [{'sample_skeleton': [state['skeleton']],
                        'query_skeleton': state['skeleton'],
-                       'sample_joints_3d': [kp_src_3d],
-                       'query_joints_3d': kp_src_3d,
+                       'sample_joints_3d': [kp_src_3d.to(device)],
+                       'query_joints_3d': kp_src_3d.to(device),
                        'sample_center': [kp_src_tensor.mean(dim=0)],
                        'query_center': kp_src_tensor.mean(dim=0),
                        'sample_scale': [
                            kp_src_tensor.max(dim=0)[0] -
-                           kp_src_tensor.min(dim=0)[0]],
+                           kp_src_tensor.min(dim=0)[0]
+                       ],
                        'query_scale': kp_src_tensor.max(dim=0)[0] -
                                       kp_src_tensor.min(dim=0)[0],
                        'sample_rotation': [0],
@@ -204,7 +184,7 @@ def process(query_img, state,
     if fp16_cfg is not None:
         wrap_fp16_model(model)
     load_checkpoint(model, checkpoint_path, map_location='cpu')
-    model.eval()
+    model.eval().to(device)
     with torch.no_grad():
         outputs = model(**data)
     # visualize results
@@ -218,39 +198,47 @@ def process(query_img, state,
                        vis_s_weight,
                        None,
                        vis_s_weight,
-                       outputs['skeleton'],
+                       outputs['skeleton'][1],
                        torch.tensor(outputs['points']).squeeze(),
                        original_skeleton=state['skeleton'],
                        img_alpha=1.0,
                        )
-    return out, state
-
-
-def update_examples(support_img, posed_support, query_img, state, r=0.015, width=0.02):
-    state['color_idx'] = 0
-    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)
-    draw_limb = ImageDraw.Draw(posed_support)
-    r = int(r * w)
-    width = int(width * w)
-    for pixel in state['kp_src']:
-        leftUpPoint = (pixel[1] - r, pixel[0] - r)
-        rightDownPoint = (pixel[1] + r, pixel[0] + r)
-        twoPointList = [leftUpPoint, rightDownPoint]
-        draw_pose.ellipse(twoPointList, fill=(255, 0, 0, 255))
-        draw_limb.ellipse(twoPointList, fill=(255, 0, 0, 255))
-    for limb in state['skeleton']:
-        point_a = state['kp_src'][limb[0]][::-1]
-        point_b = state['kp_src'][limb[1]][::-1]
-        if state['color_idx'] < len(COLORS):
-            c = COLORS[state['color_idx']]
-            state['color_idx'] += 1
-        else:
-            c = random.choices(range(256), k=3)
-        draw_limb.line([point_a, point_b], fill=tuple(c), width=width)
-    return support_img, posed_support, query_img, state
+    return out
+
+
+def update_examples(support_img, query_image, global_state_str):
+    example_state = json.loads(global_state_str)
+    example_state["load_example"] = True
+    example_state["curr_type_point"] = "start"
+    example_state["prev_point"] = None
+    example_state['images'] = {}
+    example_state['images']['image_orig'] = support_img
+    example_state['images']['image_kp'] = support_img
+    example_state['images']['image_skeleton'] = support_img
+    image_draw = example_state['images']['image_orig'].copy()
+    for xy in example_state['points']:
+        image_draw = update_image_draw(
+            image_draw,
+            xy,
+            example_state
+        )
+    kp_image = image_draw.copy()
+    example_state['images']['image_kp'] = kp_image
+    pts_list = example_state['points']
+    for limb in example_state['skeleton']:
+        prev_point = pts_list[limb[0]]
+        curr_point = pts_list[limb[1]]
+        points = [prev_point, curr_point]
+        image_draw = draw_limbs_on_image(image_draw,
+                                         points
+                                         )
+    skel_image = image_draw.copy()
+    example_state['images']['image_skel'] = skel_image
+    return (support_img,
+            kp_image,
+            skel_image,
+            query_image,
+            example_state)
 
 
 def get_select_coords(global_state,
@@ -260,17 +248,6 @@ def get_select_coords(global_state,
     """
     xy = evt.index
     global_state["points"].append(xy)
-    # point_idx = get_latest_points_pair(points)
-    # if point_idx is None:
-    #     points[0] = {'start': xy, 'target': None}
-    #     print(f'Click Image - Start - {xy}')
-    # elif points[point_idx].get('target', None) is None:
-    #     points[point_idx]['target'] = xy
-    #     print(f'Click Image - Target - {xy}')
-    # else:
-    #     points[point_idx + 1] = {'start': xy, 'target': None}
-    #     print(f'Click Image - Start - {xy}')
-
     image_raw = global_state['images']['image_kp']
     image_draw = update_image_draw(
         image_raw,
@@ -362,7 +339,9 @@ def print_memory_usage():
                            torch.cuda.max_memory_allocated()
         print(f"Available GPU memory: {available_memory / 1e9} GB")
     else:
+        device = "cpu"
         print("No GPU available")
+    return device
 
 def draw_limbs_on_image(image,
                         points,):