no mask as default. add toggle to enable mask

app.py

@@ -18,6 +18,7 @@ from PIL import Image
 import random
 from copy import deepcopy
 from huggingface_hub import hf_hub_download
+from gradio_toggle import Toggle
 try:
     import spaces
 except:
@@ -30,7 +31,6 @@ LENGTH = 480
 placeholder = cv2.cvtColor(cv2.imread("placeholder.png"), cv2.COLOR_BGR2RGB)
 NEW_MODEL = True
 MODEL_EPOCH = 6
-NO_MASK = False
 HF = False
 pre_device = "cpu" if HF else "cuda"
 spaces_60_fn = spaces.GPU(duration=60) if HF else (lambda f: f)
@@ -282,30 +282,34 @@ def prepare_anno(ref, ref_is_user):
     else:
         return img, None
 
-def get_ref_anno(img, keypts):
+def get_ref_anno(img, keypts, use_mask, use_pose):
+    no_mask, no_pose = not use_mask, not use_pose
     if img.sum() == 0:  # clear_all
         return None, gr.update(), None, gr.update(), True
     elif keypts is None:  # hand not detected
         no_hands = cv2.resize(np.array(Image.open("no_hands.png"))[..., :3], (LENGTH, LENGTH))
         return None, no_hands, None, no_hands_open, False
     missing_keys, extra_keys = autoencoder.load_state_dict(vae_state_dict, strict=False)
-    if isinstance(keypts, list):
-        if len(keypts[0]) == 0:
-            keypts[0] = np.zeros((21, 2))
-        elif len(keypts[0]) == 21:
-            keypts[0] = np.array(keypts[0], dtype=np.float32)
-        else:
-            gr.Info("Number of right hand keypoints should be either 0 or 21.")
-            return None, None, None, gr.update(), gr.update()
-        if len(keypts[1]) == 0:
-            keypts[1] = np.zeros((21, 2))
-        elif len(keypts[1]) == 21:
-            keypts[1] = np.array(keypts[1], dtype=np.float32)
-        else:
-            gr.Info("Number of left hand keypoints should be either 0 or 21.")
-            return None, None, None, gr.update(), gr.update()
-        keypts = np.concatenate(keypts, axis=0)
-    if NO_MASK:
+    if no_pose:
+        keypts = np.zeros((42, 2))
+    else:
+        if isinstance(keypts, list):
+            if len(keypts[0]) == 0:
+                keypts[0] = np.zeros((21, 2))
+            elif len(keypts[0]) == 21:
+                keypts[0] = np.array(keypts[0], dtype=np.float32)
+            else:
+                gr.Info("Number of right hand keypoints should be either 0 or 21.")
+                return None, None, None, gr.update(), gr.update()
+            if len(keypts[1]) == 0:
+                keypts[1] = np.zeros((21, 2))
+            elif len(keypts[1]) == 21:
+                keypts[1] = np.array(keypts[1], dtype=np.float32)
+            else:
+                gr.Info("Number of left hand keypoints should be either 0 or 21.")
+                return None, None, None, gr.update(), gr.update()
+            keypts = np.concatenate(keypts, axis=0)
+    if no_mask:
         hand_mask = np.zeros_like(img[:,:, 0])
         ref_pose = visualize_hand(keypts, img)
     else:
@@ -391,7 +395,7 @@ def get_ref_anno(img, keypts):
     print(f"autoencoder encoder before operating dtype: {next(autoencoder.encoder.parameters()).dtype}")
     latent = opts.latent_scaling_factor * autoencoder.encode(image).sample()
     print(f"latent.max(): {latent.max()}, latent.min(): {latent.min()}")
-    if NO_MASK:
+    if no_mask:
         mask = torch.zeros_like(mask)
     print(f"heatmaps.max(): {heatmaps.max()}, heatmaps.min(): {heatmaps.min()}")
     print(f"mask.max(): {mask.max()}, mask.min(): {mask.min()}")
@@ -921,6 +925,7 @@ def clear_all():
         42,
         3.0,
         gr.update(interactive=False),
+        False,
     )
 
 def fix_clear_all():
@@ -1365,6 +1370,9 @@ with gr.Blocks(css=custom_css, theme="soft") as demo:
         target_auto_cond = gr.State(value=None)
         target_cond = gr.State(value=None)
 
+        # config
+        use_pose = gr.State(value=True)
+
         # main tabs
         with gr.Row():
             # ref column
@@ -1386,7 +1394,8 @@ with gr.Blocks(css=custom_css, theme="soft") as demo:
                     crop_size="1:1",
                 )
                 gr.Examples(example_ref_imgs, [ref], examples_per_page=20)
-                with gr.Accordion(label="See hand pose", open=False):
+                use_mask = Toggle(label="Use mask", value=False, interactive=True)
+                with gr.Accordion(label="See hand pose & mask", open=False):
                     with gr.Tab("Automatic hand keypoints"):
                         ref_pose = gr.Image(
                             type="numpy",
@@ -1478,9 +1487,9 @@ with gr.Blocks(css=custom_css, theme="soft") as demo:
                             interactive=False,
                             visible=True
                         )
-                    gr.Markdown(
-                        """<p style="text-align: center;">Optionally flip the hand</p>"""
-                    )
+                    # gr.Markdown(
+                    #     """<p style="text-align: center;">Optionally flip the hand</p>"""
+                    # )
                     ref_flip = gr.Checkbox(
                         value=False, label="Flip Handedness (Reference)", interactive=False
                     )
@@ -1596,9 +1605,9 @@ with gr.Blocks(css=custom_css, theme="soft") as demo:
                             interactive=False,
                             visible=True
                         )
-                    gr.Markdown(
-                        """<p style="text-align: center;">Optionally flip the hand</p>"""
-                    )
+                    # gr.Markdown(
+                    #     """<p style="text-align: center;">Optionally flip the hand</p>"""
+                    # )
                     target_flip = gr.Checkbox(
                         value=False, label="Flip Handedness (Target)", interactive=False
                     )
@@ -1687,7 +1696,8 @@ with gr.Blocks(css=custom_css, theme="soft") as demo:
         ref.change(prepare_anno, [ref, ref_is_user], [ref_im_raw, ref_kp_raw])
         ref_kp_raw.change(lambda x: x, ref_im_raw, ref_manual_kp_right)
         ref_kp_raw.change(lambda x: x, ref_im_raw, ref_manual_kp_left)
-        ref_kp_raw.change(get_ref_anno, [ref_im_raw, ref_kp_raw], [ref_img, ref_pose, ref_auto_cond, ref, ref_is_user])
+        ref_kp_raw.change(get_ref_anno, [ref_im_raw, ref_kp_raw, use_mask, use_pose], [ref_img, ref_pose, ref_auto_cond, ref, ref_is_user])
+        use_mask.input(get_ref_anno, [ref_im_raw, ref_kp_raw, use_mask, use_pose], [ref_img, ref_pose, ref_auto_cond, ref, ref_is_user])
         ref_pose.change(enable_component, [ref_kp_raw, ref_pose], ref_use_auto)
         ref_pose.change(enable_component, [ref_img, ref_pose], ref_flip)
         ref_auto_cond.change(lambda x: x, ref_auto_cond, ref_cond)
@@ -1733,7 +1743,7 @@ with gr.Blocks(css=custom_css, theme="soft") as demo:
         )
         ref_manual_done.click(visible_component, [gr.State(0), ref_manual_pose], ref_manual_pose)
         ref_manual_done.click(visible_component, [gr.State(0), ref_use_manual], ref_use_manual)
-        ref_manual_done.click(get_ref_anno, [ref_im_raw, ref_kp_got], [ref_img, ref_manual_pose, ref_manual_cond])
+        ref_manual_done.click(get_ref_anno, [ref_im_raw, ref_kp_got, use_mask, use_pose], [ref_img, ref_manual_pose, ref_manual_cond])
         ref_manual_pose.change(enable_component, [ref_manual_pose, ref_manual_pose], ref_manual_done)
         ref_manual_pose.change(enable_component, [ref_img, ref_manual_pose], ref_flip)
         ref_manual_cond.change(lambda x: x, ref_manual_cond, ref_cond)
@@ -1850,6 +1860,7 @@ with gr.Blocks(css=custom_css, theme="soft") as demo:
                 seed,
                 cfg,
                 ref_kp_raw,
+                use_mask,
             ],
         )
         clear.click(
@@ -2156,6 +2167,10 @@ with gr.Blocks(css=custom_css, theme="soft") as demo:
                 postprocess=False,
                 elem_id="fix_examples_all",
             )
+        with gr.Row():
+            gr.Markdown(
+                """<p style="text-align: left;">⚠️ If brushed image doesn't load, please click the example again</p>"""
+            )
 
         # tutorial video
         with gr.Accordion("Tutorial Videos of Demo 2", elem_id="accordion_bold_large_center"):

sbatch/sbatch_demo2.sh

@@ -0,0 +1,40 @@
+#!/bin/bash
+
+# job name
+#SBATCH -J demo2
+
+# partition
+#SBATCH -p 3090-gcondo --gres=gpu:1
+
+##SBATCH --partition=ssrinath-gcondo --gres=gpu:1 --gres-flags=enforce-binding
+##SBATCH --account=ssrinath-gcondo
+
+# ensures all allocated cores are on the same node
+#SBATCH -N 1
+
+# cpu cores
+#SBATCH --ntasks-per-node=4
+
+# memory per node
+#SBATCH --mem=32G
+
+# runtime
+#SBATCH -t 240:00:00
+
+# output
+#SBATCH -o out/demo2.out
+
+# error
+#SBATCH -e err/demo2.err
+
+# email notifiaction
+# SBATCH --mail-type=ALL
+
+module load miniconda3/23.11.0s
+source /oscar/runtime/software/external/miniconda3/23.11.0/etc/profile.d/conda.sh
+conda activate handdiff
+
+cd $HOME/hdd/FoundHand_demo
+echo Directory is `pwd`
+
+python -u app.py