brush examples, visualize crop

app.py

@@ -57,11 +57,9 @@ def remove_prefix(text, prefix):
         return text[len(prefix) :]
     return text
 
-
 def unnormalize(x):
     return (((x + 1) / 2) * 255).astype(np.uint8)
 
-
 def visualize_hand(all_joints, img, side=["right", "left"], n_avail_joints=21):
     # Define the connections between joints for drawing lines and their corresponding colors
     connections = [
@@ -127,7 +125,6 @@ def visualize_hand(all_joints, img, side=["right", "left"], n_avail_joints=21):
 
     return numpy_img
 
-
 def mask_image(image, mask, color=[0, 0, 0], alpha=0.6, transparent=True):
     """Overlay mask on image for visualization purpose.
     Args:
@@ -145,7 +142,6 @@ def mask_image(image, mask, color=[0, 0, 0], alpha=0.6, transparent=True):
         out = img
     return out
 
-
 def scale_keypoint(keypoint, original_size, target_size):
     """Scale a keypoint based on the resizing of the image."""
     keypoint_copy = keypoint.copy()
@@ -153,10 +149,8 @@ def scale_keypoint(keypoint, original_size, target_size):
     keypoint_copy[:, 1] *= target_size[1] / original_size[1]
     return keypoint_copy
 
-
 print("Configure...")
 
-
 @dataclass
 class HandDiffOpts:
     run_name: str = "ViT_256_handmask_heatmap_nvs_b25_lr1e-5"
@@ -451,27 +445,39 @@ def get_target_anno(img, keypts):
 
     return img, target_pose, target_cond, keypts, gr.update(), True
 
-# def get_mask_inpaint(ref):
-#     # inpaint_mask = np.zeros_like(img_original[:, :, 0])
-#     # cropped_mask = np.array(ref["layers"][0])[..., -1]
-#     # inpaint_mask[crop_coord[0][1]:crop_coord[1][1], crop_coord[0][0]:crop_coord[1][0]] = cropped_mask
-
-#     return inpaint_mask
-
 def visualize_ref(ref):
     if ref is None:
         return None
     
+    # from user or from example
+    h, w = ref["background"].shape[:2]
+    if ref["layers"][0].sum() == 0:
+        if ref["background"][:, :, -1].sum() == h * w * 255:
+            from_example = False
+        else:
+            from_example = True
+    else:
+        from_example = False
+
     # inpaint mask
-    inpaint_mask = np.array(ref["layers"][0])[..., -1]
-    inpaint_mask = cv2.resize(
-        inpaint_mask, opts.image_size, interpolation=cv2.INTER_AREA
-    )
-    inpaint_mask = (inpaint_mask >= 128).astype(np.uint8)
+    if from_example:
+        inpaint_mask = ref["background"][:, :, -1]
+        inpainted = inpaint_mask.copy()
+        inpaint_mask = cv2.resize(
+           inpaint_mask, opts.image_size, interpolation=cv2.INTER_AREA
+        )
+        inpaint_mask = (inpaint_mask > 128).astype(np.uint8)
+        img = cv2.cvtColor(ref["background"], cv2.COLOR_RGBA2RGB)
+    else:
+        inpaint_mask = np.array(ref["layers"][0])[..., -1]
+        inpaint_mask = cv2.resize(
+            inpaint_mask, opts.image_size, interpolation=cv2.INTER_AREA
+        )
+        inpaint_mask = (inpaint_mask >= 128).astype(np.uint8)
+        inpainted = ref["layers"][0][..., -1]
+        img = ref["background"][..., :3]
 
     # viualization
-    inpainted = ref["layers"][0][..., -1]
-    img = ref["background"][..., :3]
     # img = cv2.resize(img, inpainted.shape[::-1], interpolation=cv2.INTER_AREA)
     mask = inpainted < 128
     # img = img.astype(np.int32)
@@ -481,6 +487,15 @@ def visualize_ref(ref):
     img = mask_image(img, mask)
     return img, inpaint_mask
 
+# def inpaint_from_example(ex_mask):
+#     inpaint_mask = (ex_mask[:,:, -1] == 0).astype(np.uint8)
+#     inpaint_mask = np.tile(inpaint_mask[..., None], (1, 1, 4))
+#     ret = {
+#         "background": cv2.cvtColor(ex_mask, cv2.COLOR_RGB2RGBA),
+#         "layers": [inpaint_mask],
+#         "composite": cv2.cvtColor(ex_mask, cv2.COLOR_RGB2RGBA)
+#     }
+#     return ret
 
 def get_kps(img, keypoints, side: Literal["right", "left"], evt: gr.SelectData):
     if keypoints is None:
@@ -503,7 +518,6 @@ def get_kps(img, keypoints, side: Literal["right", "left"], evt: gr.SelectData):
     vis_hand = visualize_hand(kps, img, side, len_kps)
     return vis_hand, keypoints
 
-
 def undo_kps(img, keypoints, side: Literal["right", "left"]):
     if keypoints is None:
         return img, None
@@ -523,7 +537,6 @@ def undo_kps(img, keypoints, side: Literal["right", "left"]):
     vis_hand = visualize_hand(kps, img, side, len_kps)
     return vis_hand, keypoints
 
-
 def reset_kps(img, keypoints, side: Literal["right", "left"]):
     if keypoints is None:
         return img, None
@@ -535,29 +548,54 @@ def reset_kps(img, keypoints, side: Literal["right", "left"]):
 
 def stay_crop(img, crop_coord):
     if img is not None:
-        crop_coord = [[0, 0], [img.shape[1], img.shape[0]]]
-        cropped = img.copy()
-        return crop_coord, cropped
+        if crop_coord is None:
+            crop_coord = [[0, 0], [img.shape[1], img.shape[0]]]
+            cropped = img.copy()
+            return crop_coord, cropped
+        else:
+            return gr.update(), gr.update()
     else:
         return None, None
 
+def stash_original(img):
+    if img is None:
+        return None
+    else:
+        return img[:,:,:3]
+    
 def process_crop(img, crop_coord, evt:gr.SelectData):
-    if len(crop_coord) == 2:
+    image = img.copy()
+    if len(crop_coord) == 2:  # will add first click
         crop_coord = [list(evt.index)]
-        cropped = img.copy()
+        cropped = image
+        cropped_vis = image.copy()
+        alpha = np.ones_like(cropped_vis[:,:, -1]) * 255
+        cv2.circle(alpha, tuple(crop_coord[0]), 5, 0, 4)
+        cropped_vis[:,:,-1] = alpha
     elif len(crop_coord) == 1:
         new_coord =list(evt.index)
-        if new_coord[0] <= crop_coord[0][0] or new_coord[1] <= crop_coord[0][1]:
+        if new_coord[0] <= crop_coord[0][0] or new_coord[1] <= crop_coord[0][1]:  # will skip
             gr.Warning("Second click should be more under and more right thand the first click. Try second click again.", duration=3)
-            cropped = img.copy()
-        else:
+            cropped = image
+            cropped_vis = image.copy()
+            cropped_vis[:,:,-1] = 255
+        else:  # will add second click
             crop_coord.append(new_coord)
             x1, y1 = crop_coord[0]
             x2, y2 = crop_coord[1]
-            cropped = img.copy()[y1:y2, x1:x2]
+            cropped = image[y1:y2, x1:x2]
+            cropped_vis = image.copy()
+            alpha = np.ones_like(cropped_vis[:,:, -1]) * 255
+            cv2.rectangle(alpha, tuple([x1, y1]), tuple([x2, y2]), 0, 4)
+            cropped_vis[:,:,-1] = alpha
     else:
         gr.Error("Something is wrong", duration=3)
-    return crop_coord, cropped
+    # cropped_editor = {
+    #     "background": cropped,
+    #     "composite": cropped,
+    #     "layers": [cropped_vis]
+    # }
+    return crop_coord, cropped, cropped_vis
 
 def disable_crop(crop_coord):
     if len(crop_coord) == 2:
@@ -749,7 +787,6 @@ def ready_sample(img_cropped, inpaint_mask, keypts):
         vis_mask256,
     )
 
-
 def switch_mask_size(radio):
     if radio == "256x256":
         out = (gr.update(visible=False), gr.update(visible=True))
@@ -829,7 +866,6 @@ def sample_inpaint(
             results_original.append(placeholder)
     return results, results_pose, results_original
 
-
 def flip_hand(
     img, img_raw, pose_img, pose_manual_img,
     manual_kp_right, manual_kp_left,
@@ -884,14 +920,12 @@ def flip_hand(
             manual_keypts[21:, 0] = opts.image_size[1] - manual_keypts[21:, 0]
     return img, img_raw, pose_img, pose_manual_img, manual_kp_right, manual_kp_left, cond, auto_cond, manual_cond, keypts, auto_keypts, manual_keypts
 
-
 def resize_to_full(img):
     img["background"] = cv2.resize(img["background"], (LENGTH, LENGTH))
     img["composite"] = cv2.resize(img["composite"], (LENGTH, LENGTH))
     img["layers"] = [cv2.resize(layer, (LENGTH, LENGTH)) for layer in img["layers"]]
     return img
 
-
 def clear_all():
     return (
         None,
@@ -921,7 +955,6 @@ def clear_all():
         gr.update(interactive=False),
     )
 
-
 def fix_clear_all():
     return (
         None,
@@ -951,7 +984,6 @@ def fix_clear_all():
         10,
     )
 
-
 def enable_component(image1, image2):
     if image1 is None or image2 is None:
         return gr.update(interactive=False)
@@ -971,7 +1003,6 @@ def enable_component(image1, image2):
             return gr.update(interactive=False)
     return gr.update(interactive=True)
 
-
 def set_visible(checkbox, kpts, img_clean, img_pose_right, img_pose_left, done=None, done_info=None):
     if kpts is None:
         kpts = [[], []]
@@ -1162,18 +1193,28 @@ fix_example_imgs = [
     ["bad_hands/1.jpg"],  # "bad_hands/1_mask.jpg"],
     # ["bad_hands/2.jpg"],  # "bad_hands/2_mask.jpg"],
     ["bad_hands/3.jpg"],  # "bad_hands/3_mask.jpg"],
-    # ["bad_hands/4.jpg"],  # "bad_hands/4_mask.jpg"],
+    ["bad_hands/4.jpg"],  # "bad_hands/4_mask.jpg"],
     ["bad_hands/5.jpg"],  # "bad_hands/5_mask.jpg"],
     ["bad_hands/6.jpg"],  # "bad_hands/6_mask.jpg"],
-    # ["bad_hands/7.jpg"],  # "bad_hands/7_mask.jpg"],
+    ["bad_hands/7.jpg"],  # "bad_hands/7_mask.jpg"],
     # ["bad_hands/8.jpg"],  # "bad_hands/8_mask.jpg"],
     # ["bad_hands/9.jpg"],  # "bad_hands/9_mask.jpg"],
     # ["bad_hands/10.jpg"],  # "bad_hands/10_mask.jpg"],
     # ["bad_hands/11.jpg"],  # "bad_hands/11_mask.jpg"],
     # ["bad_hands/12.jpg"],  # "bad_hands/12_mask.jpg"],
     # ["bad_hands/13.jpg"],  # "bad_hands/13_mask.jpg"],
-    ["bad_hands/14.jpg"],
-    ["bad_hands/15.jpg"],
+    # ["bad_hands/14.jpg"],
+    # ["bad_hands/15.jpg"],
+]
+fix_example_brush = [
+    ["bad_hands/1_composite.png"],# "bad_hands/1.jpg"],
+    ["bad_hands/3_composite.png"],# "bad_hands/3.jpg"],
+    ["bad_hands/4_composite.png"],# "bad_hands/4.jpg"],
+    ["bad_hands/5_composite.png"],# "bad_hands/5.jpg"],
+    ["bad_hands/6_composite.png"],# "bad_hands/6.jpg"],
+    ["bad_hands/7_composite.png"],# "bad_hands/7.jpg"],
+    # ["bad_hands/14_mask.jpg"],
+    # ["bad_hands/15_mask.jpg"],
 ]
 custom_css = """
 .gradio-container .examples img {
@@ -1286,6 +1327,7 @@ with gr.Blocks(css=custom_css, theme="soft") as demo:
                     width=LENGTH,
                     interactive=True,
                     visible=True,
+                    image_mode="RGBA"
                 )
                 gr.Markdown(
                     """<p style="text-align: center;">💡 If you crop, the model can focus on more details of the cropped area. Square crops might work better than rectangle crops.</p>"""
@@ -1304,6 +1346,8 @@ with gr.Blocks(css=custom_css, theme="soft") as demo:
                     fix_example_imgs,
                     inputs=[fix_crop],
                     examples_per_page=20,
+                    # run_on_click=True,
+                    # fn=load_brush,
                 )
             with gr.Column():
                 gr.Markdown(
@@ -1326,14 +1370,22 @@ with gr.Blocks(css=custom_css, theme="soft") as demo:
                     ),  # 204, 50, 50
                     image_mode="RGBA",
                     container=False,
-                    interactive=False,
+                    interactive=True,
+                )
+                fix_ex_brush = gr.Examples(
+                    fix_example_brush,
+                    inputs=[fix_ref],
+                    # outputs=[fix_ref],
+                    examples_per_page=20,
+                    # run_on_click=True,
+                    # fn=inpaint_from_example,
                 )
                 # gr.Markdown(
                 #     """<p style="text-align: center;">&#9314; Hit the \"Finish Cropping & Brushing\" button</p>"""
                 # )
-                fix_finish_crop = gr.Button(
-                    value="Finish Croping & Brushing", interactive=False
-                )
+                # fix_finish_crop = gr.Button(
+                #     value="Finish Croping & Brushing", interactive=False
+                # )
             
             # keypoint selection
             with gr.Column():
@@ -1526,18 +1578,18 @@ with gr.Blocks(css=custom_css, theme="soft") as demo:
         
         # listeners
         # fix_crop.change(resize_to_full, fix_crop, fix_ref)
-        fix_crop.change(lambda x: x, fix_crop, fix_original)  # fix_original: (real_H, real_W, 3)
+        fix_crop.change(stash_original, fix_crop, fix_original)  # fix_original: (real_H, real_W, 3)
         fix_crop.change(stay_crop, [fix_crop, fix_crop_coord], [fix_crop_coord, fix_ref])
-        fix_crop.select(process_crop, [fix_crop, fix_crop_coord], [fix_crop_coord, fix_ref])
+        fix_crop.select(process_crop, [fix_crop, fix_crop_coord], [fix_crop_coord, fix_ref, fix_crop])
         # fix_ref.change(disable_crop, fix_crop_coord, fix_crop)
-        fix_ref.change(enable_component, [fix_crop, fix_crop], fix_ref)
-        fix_ref.change(enable_component, [fix_crop, fix_crop], fix_finish_crop)
-        fix_finish_crop.click(visualize_ref, [fix_ref], [fix_img, fix_inpaint_mask])
+        # fix_ref.change(enable_component, [fix_crop, fix_crop], fix_ref)
+        # fix_ref.change(enable_component, [fix_crop, fix_crop], fix_finish_crop)
+        fix_ref.change(visualize_ref, [fix_ref], [fix_img, fix_inpaint_mask])
         # fix_finish_crop.click(get_mask_inpaint, [fix_ref], [])  # fix_ref: (real_cropped_H, real_cropped_W, 3)
         fix_img.change(lambda x: x, [fix_img], [fix_kp_right])
         fix_img.change(lambda x: x, [fix_img], [fix_kp_left])
-        fix_inpaint_mask.change(
-            enable_component, [fix_inpaint_mask, fix_inpaint_mask], fix_checkbox
+        fix_ref.change(
+            enable_component, [fix_ref, fix_ref], fix_checkbox
         )
         fix_inpaint_mask.change(
             enable_component, [fix_inpaint_mask, fix_inpaint_mask], fix_kp_right
@@ -2220,7 +2272,7 @@ with gr.Blocks(css=custom_css, theme="soft") as demo:
     gr.Markdown("<h1>Acknowledgement</h1>")
     gr.Markdown(_ACK_)
     gr.Markdown("<h1>Trouble Shooting</h1>")
-    gr.Markdown("If something doesn't work, <br>1. Try refreshing the page and do it again. <br>2. Leave a message at our HuggingFace Spaces's \"Community\" tab on the top right, at our Github repo's Issues, or email us.<br>3. The problem might be from compatibility issue to HuggingFace or GPU vram limitations. If that's possible, we highly recommend you to clone this repo and try with your own gpu.")
+    gr.Markdown("If something doesn't work, please try the following steps:<br>1. Refresh the page and try again.<br>2. If the issue persists, kindly leave a message on our HuggingFace Spaces Community tab (located at the top right), on our GitHub repository's Issues page, or send us an email. We are happy to help you as soon as possible.<br>3. The issue might be due to compatibility with HuggingFace or GPU VRAM limitations. We recommend cloning this repository and trying it with your own GPU if possible.")
     gr.Markdown("<h1>Citation</h1>")
     gr.Markdown(
         """<p style="text-align: left;">If this was useful, please cite us! ❤️</p>"""

tmp.py

@@ -0,0 +1,18 @@
+import glob
+import os
+from PIL import Image
+import numpy as np
+
+img_dir = "bad_hands"
+masked_paths = sorted(glob.glob(os.path.join(img_dir, "*_mask.jpg")))
+for masked_pth in masked_paths:
+    img_path = masked_pth.replace("_mask.jpg", ".jpg")
+    assert os.path.exists(img_path), f"Image path {img_path} does not exist."
+    masked = np.array(Image.open(masked_pth))
+    mask = (np.all(masked > 245, axis=-1)).astype(np.uint8)*128 + 64
+    img = np.array(Image.open(img_path))
+    composite = np.concatenate((img, mask[..., None]), axis=-1)
+    # img.putalpha(Image.fromarray(mask))
+    composite = Image.fromarray(composite)
+    composite.save(masked_pth.replace("_mask.jpg", "_composite.png"))
+    print(f"Saved composite image {masked_pth.replace('_mask.jpg', '_composite.png')}")