Add mask drawing

app.py

@@ -76,12 +76,20 @@ def inpaint(input_img, mask, option):
         mask = Image.fromarray(result[0]['mask'])
     else:
         mask = mask.resize((width,height))
-    
-    if width is not 512 or height is not 512:
+    mask = mask.convert('L')
+
+    if width != 512 or height != 512:
         input_img = input_img.resize((512, 512))
         mask = mask.resize((512, 512))
+    
+    
+    if option == 'Manual':
+        mask = (np.array(mask) - np.array(input_img.convert('L'))) > 0.
+        mask = mask * 1.
+        kernel = np.ones((5, 5), np.uint8)
+        mask = cv2.dilate(mask, kernel) 
+        mask = mask * 255.
 
-    mask = mask.convert('L')
     mask = np.array(mask) / 255.
     mask_tensor = torch.from_numpy(mask).to(torch.float32)
     mask_tensor = mask_tensor.unsqueeze(0)
@@ -106,30 +114,34 @@ def inpaint(input_img, mask, option):
 gradio_inputs = [gr.inputs.Image(type='pil',
                                  tool=None,
                                  label="Input Image"),
-                gr.inputs.Image(type='pil',source="canvas", label="Mask", invert_colors=True),
+                # gr.inputs.Image(type='pil',source="canvas", label="Mask", invert_colors=True),
+                gr.inputs.Image(type='pil',
+                                 tool="editor",
+                                 label="Mask"),
                 gr.inputs.Radio(choices=["Automatic", "Manual"], type="value", default="Manual", label="Masking Choice")
                 ]
 
 gradio_outputs = [gr.outputs.Image(label='Image with Hole'),
                  gr.outputs.Image(label='Inpainted Image')]
 
-examples = [['test_512/person512.png', 'test_512/mask_auto.png', 'Automatic'],
-            ['test_512/a_org.png', 'test_512/a_mask.png', 'Manual'],
-            ['test_512/c_org.png', 'test_512/b_mask.png', 'Manual'],
-            ['test_512/b_org.png', 'test_512/c_mask.png', 'Manual'],
-            ['test_512/d_org.png', 'test_512/d_mask.png', 'Manual'],
-            ['test_512/e_org.png', 'test_512/e_mask.png', 'Manual'],
-            ['test_512/f_org.png', 'test_512/f_mask.png', 'Manual'],
-            ['test_512/g_org.png', 'test_512/g_mask.png', 'Manual'],
-            ['test_512/h_org.png', 'test_512/h_mask.png', 'Manual'],
-            ['test_512/i_org.png', 'test_512/i_mask.png', 'Manual']]
+examples = [['test_512/person512.png', 'test_512/person512.png', 'Automatic'],
+            ['test_512/a_org.png', 'test_512/a_overlay.png', 'Manual'],
+            ['test_512/b_org.png', 'test_512/b_overlay.png', 'Manual'],
+            ['test_512/c_org.png', 'test_512/c_overlay.png', 'Manual'],
+            ['test_512/d_org.png', 'test_512/d_overlay.png', 'Manual'],
+            ['test_512/e_org.png', 'test_512/e_overlay.png', 'Manual'],
+            ['test_512/f_org.png', 'test_512/f_overlay.png', 'Manual'],
+            ['test_512/g_org.png', 'test_512/g_overlay.png', 'Manual'],
+            ['test_512/h_org.png', 'test_512/h_overlay.png', 'Manual'],
+            ['test_512/i_org.png', 'test_512/i_overlay.png', 'Manual']]
 
 title = "FcF-Inpainting"
 description = "[Note: Queue time may take upto 20 seconds! The image and mask are resized to 512x512 before inpainting.] To use FcF-Inpainting: \n \
                 (1) Upload an Image; \n \
-                (2) Draw (Manual) a Mask on the White Canvas or Generate a mask using U2Net by selecting the Automatic option; \n \
+                (2a) Draw (Manual) a Mask using the brush (click on the edit option in the top right of the Mask View) \n \
+                (2b) You can also generate a mask using U2Net by selecting the Automatic option; \n \
                 (3) Click on Submit and witness the MAGIC! 🪄 ✨ ✨"
-article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2112.10741' target='_blank'> Keys to Better Image Inpainting: Structure and Texture Go Hand in Hand</a> | <a href='https://github.com/SHI-Labs/FcF-Inpainting' target='_blank'>Github Repo</a></p>"
+article = "<p style='text-align: center'><a href='https://github.com/SHI-Labs/FcF-Inpainting' target='_blank'> Keys to Better Image Inpainting: Structure and Texture Go Hand in Hand</a> | <a href='https://github.com/SHI-Labs/FcF-Inpainting' target='_blank'>Github Repo</a></p>"
 
 css = ".image-preview {height: 32rem; width: auto;} .output-image {height: 32rem; width: auto;} .panel-buttons { display: flex; flex-direction: row;}"