Add mask extension function to inpaint mode.

app.py

@@ -250,6 +250,38 @@ def xywh_to_xyxy(box, sizeW, sizeH):
     box = box.numpy()
     return box
 
+def to_extend_mask(segment_mask, boxes_filt, size, remove_mask_extend, remove_mode):
+    # remove from mask
+    mask_imgs = []
+    masks_shape = segment_mask.shape        
+    boxes_filt_ori_array = boxes_filt.numpy()
+    if inpaint_mode == 'merge':
+        extend_shape_0 = masks_shape[0]
+        extend_shape_1 = masks_shape[1]
+    else:
+        extend_shape_0 = 1
+        extend_shape_1 = 1
+    for i in range(extend_shape_0):
+        for j in range(extend_shape_1):                
+            mask = segment_mask[i][j].cpu().numpy()
+            mask_pil = Image.fromarray(mask)
+        
+            if remove_mode == 'segment':
+                useRectangle = False
+            else:
+                useRectangle = True
+
+            try:
+                remove_mask_extend = int(remove_mask_extend)
+            except:
+                remove_mask_extend = 10
+            mask_pil_exp = mask_extend(copy.deepcopy(mask_pil).convert("RGB"), 
+                            xywh_to_xyxy(torch.tensor(boxes_filt_ori_array[i]), size[0], size[1]),
+                            extend_pixels=remove_mask_extend, useRectangle=useRectangle)
+            mask_imgs.append(mask_pil_exp)
+    mask_pil = mix_masks(mask_imgs)
+    return mask_pil
+
 def run_anything_task(input_image, text_prompt, task_type, inpaint_prompt, box_threshold, text_threshold, 
             iou_threshold, inpaint_mode, mask_source_radio, remove_mode, remove_mask_extend):
 
@@ -372,40 +404,16 @@ def run_anything_task(input_image, text_prompt, task_type, inpaint_prompt, box_t
         if task_type == 'inpainting':
             # inpainting pipeline
             image_source_for_inpaint = image_pil.resize((512, 512))
+            if remove_mask_extend:
+                mask_pil = to_extend_mask(masks_ori, boxes_filt_ori, size, remove_mask_extend, remove_mode)
+                output_images.append(mask_pil.convert("RGB"))
             image_mask_for_inpaint = mask_pil.resize((512, 512))
             image_inpainting = sd_pipe(prompt=inpaint_prompt, image=image_source_for_inpaint, mask_image=image_mask_for_inpaint).images[0]            
         else:
             # remove from mask
             if mask_source_radio == mask_source_segment:
-                mask_imgs = []
-                masks_shape = masks_ori.shape        
-                boxes_filt_ori_array = boxes_filt_ori.numpy()
-                if inpaint_mode == 'merge':
-                    extend_shape_0 = masks_shape[0]
-                    extend_shape_1 = masks_shape[1]
-                else:
-                    extend_shape_0 = 1
-                    extend_shape_1 = 1
-                for i in range(extend_shape_0):
-                    for j in range(extend_shape_1):                
-                        mask = masks_ori[i][j].cpu().numpy()
-                        mask_pil = Image.fromarray(mask)
-                    
-                        if remove_mode == 'segment':
-                            useRectangle = False
-                        else:
-                            useRectangle = True
-
-                        try:
-                            remove_mask_extend = int(remove_mask_extend)
-                        except:
-                            remove_mask_extend = 10
-                        mask_pil_exp = mask_extend(copy.deepcopy(mask_pil).convert("RGB"), 
-                                        # box_convert(torch.tensor(boxes_filt_ori_array[i]), in_fmt="cxcywh", out_fmt="xyxy").numpy(),
-                                        xywh_to_xyxy(torch.tensor(boxes_filt_ori_array[i]), size[0], size[1]),
-                                        extend_pixels=remove_mask_extend, useRectangle=useRectangle)
-                        mask_imgs.append(mask_pil_exp)
-                mask_pil = mix_masks(mask_imgs)
+                if remove_mask_extend:  
+                    mask_pil = to_extend_mask(masks_ori, boxes_filt_ori, size, remove_mask_extend, remove_mode)
                 output_images.append(mask_pil.convert("RGB"))               
             image_inpainting = lama_cleaner_process(np.array(image_pil), np.array(mask_pil.convert("L")))
 
@@ -495,7 +503,7 @@ if __name__ == "__main__":
                         with gr.Column(scale=1):
                             remove_mode = gr.Radio(["segment", "rectangle"],  value="segment", label='remove mode') 
                         with gr.Column(scale=1):
-                            remove_mask_extend = gr.Textbox(label="remove_mask_extend", value='10')
+                            remove_mask_extend = gr.Textbox(label="Enlarge Mask (Empty: no mask extension, default: 10)")
 
             with gr.Column():
                 gallery = gr.Gallery(label="result images", show_label=True, elem_id="gallery", visible=True