Update app.py

app.py

@@ -162,6 +162,122 @@ def process(query_img, state,
     return out, 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()
+    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
+
+
+def get_select_coords(kp_support,
+                      limb_support,
+                      state,
+                      evt: gr.SelectData,
+                      r=0.015):
+    pixels_in_queue = set()
+    pixels_in_queue.add((evt.index[1], evt.index[0]))
+    while len(pixels_in_queue) > 0:
+        pixel = pixels_in_queue.pop()
+        if pixel[0] is not None and pixel[1] is not None and pixel not in \
+                state['kp_src']:
+            state['kp_src'].append(pixel)
+        else:
+            continue
+        if limb_support is None:
+            canvas_limb = kp_support
+        else:
+            canvas_limb = limb_support
+        canvas_kp = kp_support
+        w, h = canvas_kp.size
+        draw_pose = ImageDraw.Draw(canvas_kp)
+        draw_limb = ImageDraw.Draw(canvas_limb)
+        r = int(r * w)
+        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))
+    return canvas_kp, canvas_limb, state
+
+
+def get_limbs(kp_support,
+              state,
+              evt: gr.SelectData,
+              r=0.02, width=0.02):
+    curr_pixel = (evt.index[1], evt.index[0])
+    pixels_in_queue = set()
+    pixels_in_queue.add((evt.index[1], evt.index[0]))
+    canvas_kp = kp_support
+    w, h = canvas_kp.size
+    r = int(r * w)
+    width = int(width * w)
+    while len(pixels_in_queue) > 0 and curr_pixel != state['prev_clicked']:
+        pixel = pixels_in_queue.pop()
+        state['prev_clicked'] = pixel
+        closest_point = min(state['kp_src'],
+                            key=lambda p: (p[0] - pixel[0]) ** 2 +
+                                          (p[1] - pixel[1]) ** 2)
+        closest_point_index = state['kp_src'].index(closest_point)
+        draw_limb = ImageDraw.Draw(canvas_kp)
+        if state['color_idx'] < len(COLORS):
+            c = COLORS[state['color_idx']]
+        else:
+            c = random.choices(range(256), k=3)
+        leftUpPoint = (closest_point[1] - r, closest_point[0] - r)
+        rightDownPoint = (closest_point[1] + r, closest_point[0] + r)
+        twoPointList = [leftUpPoint, rightDownPoint]
+        draw_limb.ellipse(twoPointList, fill=tuple(c))
+        if state['count'] == 0:
+            state['prev_pt'] = closest_point[1], closest_point[0]
+            state['prev_pt_idx'] = closest_point_index
+            state['count'] = state['count'] + 1
+        else:
+            if state['prev_pt_idx'] != closest_point_index:
+                # Create Line and add Limb
+                draw_limb.line(
+                    [state['prev_pt'], (closest_point[1], closest_point[0])],
+                    fill=tuple(c),
+                    width=width)
+                state['skeleton'].append(
+                    (state['prev_pt_idx'], closest_point_index))
+                state['color_idx'] = state['color_idx'] + 1
+            else:
+                draw_limb.ellipse(twoPointList, fill=(255, 0, 0, 255))
+            state['count'] = 0
+    return canvas_kp, state
+
+
+def set_query(support_img, state, example=False):
+    if not example:
+        state['skeleton'].clear()
+        state['kp_src'].clear()
+    state['original_support_image'] = np.array(support_img)[:, :, ::-1].copy()
+    width, height = support_img.size
+    support_img = support_img.resize((width // 4, width // 4),
+                                     Image.Resampling.LANCZOS)
+    return support_img, support_img, state
+
+
 with gr.Blocks() as demo:
     state = gr.State({
         'kp_src': [],
@@ -182,7 +298,8 @@ with gr.Blocks() as demo:
     this crucial structural information, our method enhances the accuracy of 
     keypoint localization, marking a significant departure from conventional 
     CAPE techniques that treat keypoints as isolated entities.
-    ### [Paper](https://arxiv.org/abs/2311.17891) | [Official Repo](https://github.com/orhir/PoseAnything) 
+    ### [Paper](https://arxiv.org/abs/2311.17891) | [Official Repo](
+    https://github.com/orhir/PoseAnything) 
     ## Instructions
     1. Upload an image of the object you want to pose on the **left** image.
     2. Click on the **left** image to mark keypoints.
@@ -207,103 +324,65 @@ with gr.Blocks() as demo:
         eval_btn = gr.Button(value="Evaluate")
     with gr.Row():
         output_img = gr.Plot(label="Output Image", height=400, width=400)
-
-
-    def get_select_coords(kp_support,
-                          limb_support,
-                          state,
-                          evt: gr.SelectData,
-                          r=0.015):
-        # global original_support_image
-        # if len(kp_src) == 0:
-        #     original_support_image = np.array(kp_support)[:, :,
-        #                              ::-1].copy()
-        pixels_in_queue = set()
-        pixels_in_queue.add((evt.index[1], evt.index[0]))
-        while len(pixels_in_queue) > 0:
-            pixel = pixels_in_queue.pop()
-            if pixel[0] is not None and pixel[
-                1] is not None and pixel not in state['kp_src']:
-                state['kp_src'].append(pixel)
-            else:
-                print("Invalid pixel")
-            if limb_support is None:
-                canvas_limb = kp_support
-            else:
-                canvas_limb = limb_support
-            canvas_kp = kp_support
-            w, h = canvas_kp.size
-            draw_pose = ImageDraw.Draw(canvas_kp)
-            draw_limb = ImageDraw.Draw(canvas_limb)
-            r = int(r * w)
-            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))
-
-        return canvas_kp, canvas_limb, state
-
-
-    def get_limbs(kp_support,
-                  state,
-                  evt: gr.SelectData,
-                  r=0.02, width=0.02):
-        curr_pixel = (evt.index[1], evt.index[0])
-        pixels_in_queue = set()
-        pixels_in_queue.add((evt.index[1], evt.index[0]))
-        canvas_kp = kp_support
-        w, h = canvas_kp.size
-        r = int(r * w)
-        width = int(width * w)
-        while len(pixels_in_queue) > 0 and curr_pixel != state['prev_clicked']:
-            pixel = pixels_in_queue.pop()
-            state['prev_clicked'] = pixel
-            closest_point = min(state['kp_src'],
-                                key=lambda p: (p[0] - pixel[0]) ** 2 +
-                                              (p[1] - pixel[1]) ** 2)
-            closest_point_index = state['kp_src'].index(closest_point)
-            draw_limb = ImageDraw.Draw(canvas_kp)
-            if state['color_idx'] < len(COLORS):
-                c = COLORS[state['color_idx']]
-            else:
-                c = random.choices(range(256), k=3)
-            leftUpPoint = (closest_point[1] - r, closest_point[0] - r)
-            rightDownPoint = (closest_point[1] + r, closest_point[0] + r)
-            twoPointList = [leftUpPoint, rightDownPoint]
-            draw_limb.ellipse(twoPointList, fill=tuple(c))
-            if state['count'] == 0:
-                state['prev_pt'] = closest_point[1], closest_point[0]
-                state['prev_pt_idx'] = closest_point_index
-                state['count'] = state['count'] + 1
-            else:
-                if state['prev_pt_idx'] != closest_point_index:
-                    # Create Line and add Limb
-                    draw_limb.line(
-                        [state['prev_pt'], (closest_point[1], closest_point[0])],
-                        fill=tuple(c),
-                        width=width)
-                    state['skeleton'].append((state['prev_pt_idx'], closest_point_index))
-                    state['color_idx'] = state['color_idx'] + 1
-                else:
-                    draw_limb.ellipse(twoPointList, fill=(255, 0, 0, 255))
-                state['count'] = 0
-        return canvas_kp, state
-
-
-    def set_qery(support_img, state):
-        state['skeleton'].clear()
-        state['kp_src'].clear()
-        state['original_support_image'] = np.array(support_img)[:, :, ::-1].copy()
-        width, height = support_img.size
-        support_img = support_img.resize((width // 4, width // 4), Image.Resampling.LANCZOS)
-        return support_img, support_img, state
-
+    with gr.Row():
+        gr.Markdown("## Examples")
+    with gr.Row():
+        gr.Examples(
+            examples=[
+                ['examples/dog2.png',
+                 'examples/dog2.png',
+                 'examples/dog1.png',
+                 {'kp_src': [(50, 58), (51, 78), (66, 57), (118, 79),
+                             (154, 79), (217, 74), (218, 103), (156, 104),
+                             (152, 151), (215, 162), (213, 191),
+                             (152, 174), (108, 171)],
+                  'skeleton': [(0, 1), (1, 2), (0, 2), (3, 4), (4, 5),
+                               (3, 7), (7, 6), (3, 12), (12, 8), (8, 9),
+                               (12, 11), (11, 10)], 'count': 0,
+                  'color_idx': 0, 'prev_pt': (174, 152),
+                  'prev_pt_idx': 11, 'prev_clicked': (207, 186),
+                  'original_support_image': None,
+                  }
+                 ],
+                ['examples/sofa1.jpg',
+                 'examples/sofa1.jpg',
+                 'examples/sofa2.jpg',
+                 {
+                     'kp_src': [(82, 28), (65, 30), (52, 26), (65, 50),
+                                (84, 52), (53, 54), (43, 52), (45, 71),
+                                (81, 69), (77, 39), (57, 43), (58, 64),
+                                (46, 42), (49, 65)],
+                     'skeleton': [(0, 1), (3, 1), (3, 4), (10, 9), (11, 8),
+                                  (1, 10), (10, 11), (11, 3), (1, 2), (7, 6),
+                                  (5, 13), (5, 3), (13, 11), (12, 10), (12, 2),
+                                  (6, 10), (7, 11)], 'count': 0,
+                     'color_idx': 23, 'prev_pt': (71, 45), 'prev_pt_idx': 7,
+                     'prev_clicked': (56, 63),
+                     'original_support_image': None,
+                 }],
+                ['examples/person1.jpeg',
+                 'examples/person1.jpeg',
+                 'examples/person2.jpeg',
+                 {
+                     'kp_src': [(121, 95), (122, 160), (154, 130), (184, 106),
+                                (181, 153)],
+                     'skeleton': [(0, 1), (1, 2), (0, 2), (2, 3), (2, 4),
+                                  (4, 3)], 'count': 0, 'color_idx': 6,
+                     'prev_pt': (153, 181), 'prev_pt_idx': 4,
+                     'prev_clicked': (181, 108),
+                     'original_support_image': None,
+                 }]
+            ],
+            inputs=[support_img, posed_support, query_img, state],
+            outputs=[support_img, posed_support, query_img, state],
+            fn=update_examples,
+            run_on_click=True,
+        )
 
     support_img.select(get_select_coords,
                        [support_img, posed_support, state],
                        [support_img, posed_support, state])
-    support_img.upload(set_qery,
+    support_img.upload(set_query,
                        inputs=[support_img, state],
                        outputs=[support_img, posed_support, state])
     posed_support.select(get_limbs,