update

app.py

@@ -78,10 +78,11 @@ def inference_single_image(ref_image,
                            strength, 
                            ddim_steps, 
                            scale, 
-                           seed, 
+                           seed,
+                           enable_shape_control 
                            ):
     raw_background = tar_image.copy()
-    item = process_pairs(ref_image, ref_mask, tar_image, tar_mask)
+    item = process_pairs(ref_image, ref_mask, tar_image, tar_mask, enable_shape_control = enable_shape_control)
 
     ref = item['ref']
     hint = item['hint']
@@ -133,7 +134,7 @@ def inference_single_image(ref_image,
     return raw_background
 
 
-def process_pairs(ref_image, ref_mask, tar_image, tar_mask, max_ratio = 0.8):
+def process_pairs(ref_image, ref_mask, tar_image, tar_mask, max_ratio = 0.8, enable_shape_control = False):
     # ========= Reference ===========
     # ref expand 
     ref_box_yyxx = get_bbox_from_mask(ref_mask)
@@ -189,21 +190,23 @@ def process_pairs(ref_image, ref_mask, tar_image, tar_mask, max_ratio = 0.8):
 
     collage_mask = cropped_target_image.copy() * 0.0
     collage_mask[y1:y2,x1:x2,:] = 1.0
-    collage_mask = np.stack([cropped_tar_mask,cropped_tar_mask,cropped_tar_mask],-1)
+    if enable_shape_control:
+        collage_mask = np.stack([cropped_tar_mask,cropped_tar_mask,cropped_tar_mask],-1)
 
     # the size before pad
     H1, W1 = collage.shape[0], collage.shape[1]
 
     cropped_target_image = pad_to_square(cropped_target_image, pad_value = 0, random = False).astype(np.uint8)
     collage = pad_to_square(collage, pad_value = 0, random = False).astype(np.uint8)
-    collage_mask = pad_to_square(collage_mask, pad_value = 0, random = False).astype(np.uint8)
+    collage_mask = pad_to_square(collage_mask, pad_value = 2, random = False).astype(np.uint8)
 
     # the size after pad
     H2, W2 = collage.shape[0], collage.shape[1]
 
     cropped_target_image = cv2.resize(cropped_target_image.astype(np.uint8), (512,512)).astype(np.float32)
     collage = cv2.resize(collage.astype(np.uint8), (512,512)).astype(np.float32)
-    collage_mask  = (cv2.resize(collage_mask.astype(np.uint8), (512,512)).astype(np.float32) > 0.5).astype(np.float32)
+    collage_mask  = cv2.resize(collage_mask.astype(np.uint8), (512,512),  interpolation = cv2.INTER_NEAREST).astype(np.float32)
+    collage_mask[collage_mask == 2] = -1
 
     masked_ref_image = masked_ref_image  / 255 
     cropped_target_image = cropped_target_image / 127.5 - 1.0
@@ -225,13 +228,6 @@ ref_list.sort()
 image_list=[os.path.join(image_dir,file) for file in os.listdir(image_dir) if '.jpg' in file or '.png' in file or '.jpeg' in file]
 image_list.sort()
 
-def process_image_mask(image_np, mask_np):
-    img = torch.from_numpy(image_np.transpose((2, 0, 1)))
-    img_ten = img.float().div(255).unsqueeze(0)
-    mask_ten = torch.from_numpy(mask_np).float().unsqueeze(0).unsqueeze(0)
-    return img_ten, mask_ten
-
-
 def mask_image(image, mask):
     blanc = np.ones_like(image) * 255
     mask = np.stack([mask,mask,mask],-1) / 255
@@ -247,49 +243,38 @@ def run_local(base,
     ref_mask = ref["mask"].convert("L")
     image = np.asarray(image)
     mask = np.asarray(mask)
-    mask = np.where(mask > 128, 255, 0).astype(np.uint8)
+    mask = np.where(mask > 128, 1, 0).astype(np.uint8)
     ref_image = np.asarray(ref_image)
     ref_mask = np.asarray(ref_mask)
     ref_mask = np.where(ref_mask > 128, 1, 0).astype(np.uint8)
 
-    # refine the user annotated coarse mask
-    if use_interactive_seg:
-        img_ten, mask_ten = process_image_mask(ref_image, ref_mask)
-        ref_mask = iseg_model(img_ten, mask_ten)['instances'][0,0].detach().numpy() > 0.5
-
-    processed_item = process_pairs(ref_image.copy(), ref_mask.copy(), image.copy(), mask.copy(), max_ratio = 0.8)
-    masked_ref = (processed_item['ref']*255)
-
-    mased_image = mask_image(image, mask)
-    #synthesis = image
     synthesis = inference_single_image(ref_image.copy(), ref_mask.copy(), image.copy(), mask.copy(), *args)
     synthesis = torch.from_numpy(synthesis).permute(2, 0, 1)
     synthesis = synthesis.permute(1, 2, 0).numpy()
-
-    masked_ref = cv2.resize(masked_ref.astype(np.uint8), (512,512))
     return [synthesis]
 
+
+
 with gr.Blocks() as demo:
     with gr.Column():
         gr.Markdown("#  Play with AnyDoor to Teleport your Target Objects! ")
         with gr.Row():
             baseline_gallery = gr.Gallery(label='Output', show_label=True, elem_id="gallery", columns=1, height=768)
             with gr.Accordion("Advanced Option", open=True):
-                #num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1)
                 num_samples = 1
                 strength = gr.Slider(label="Control Strength", minimum=0.0, maximum=2.0, value=1.0, step=0.01)
                 ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=30, step=1)
                 scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=30.0, value=4.5, step=0.1)
                 seed = gr.Slider(label="Seed", minimum=-1, maximum=999999999, step=1, value=-1)
+                enable_shape_control = gr.Checkbox(label='Enable Shape Control', value=False)
                 gr.Markdown(" Higher guidance-scale makes higher fidelity, while lower guidance-scale leads to more harmonized blending.")
     
-
         gr.Markdown("# Upload / Select Images for the Background (left) and Reference Object (right)")
         gr.Markdown("### Your could draw coarse masks on the background to indicate the desired location and shape.")
         gr.Markdown("### <u>Do not forget</u> to annotate the target object on the reference image.")
         with gr.Row():
-            base = gr.Image(label="Background", tool="sketch", type="pil", height=512, brush_color='#FFFFFF', mask_opacity=0.5)
-            ref = gr.Image(label="Reference", tool="sketch", type="pil", height=512, brush_color='#FFFFFF', mask_opacity=0.5)
+            base = gr.Image(label="Background", source="upload", tool="sketch", type="pil", height=512, brush_color='#FFFFFF', mask_opacity=0.5)
+            ref = gr.Image(label="Reference", source="upload", tool="sketch", type="pil", height=512, brush_color='#FFFFFF', mask_opacity=0.5)
         run_local_button = gr.Button(label="Generate", value="Run")
 
         with gr.Row():
@@ -304,7 +289,8 @@ with gr.Blocks() as demo:
                                    strength, 
                                    ddim_steps, 
                                    scale, 
-                                   seed, 
+                                   seed,
+                                   enable_shape_control, 
                                    ], 
                            outputs=[baseline_gallery]
                         )