v2.3

app.py

@@ -16,7 +16,8 @@ def process_image(image, mask, progress=gr.Progress(track_tqdm=True)):
         print("Mask Size : ", image["layers"][0].shape)
         print("Unique values: ", np.unique(image["layers"][0]))
         print("Type : ", type(image["layers"][0]))
-        output = process_inpaint(image["background"], image["layers"][0])
+        mask = image["layers"][0]
+        output = process_inpaint(image["background"], mask)
     else: 
         print("\nProcessing Received Mask\n")
         print("Mask Size : ", mask.shape)
@@ -25,10 +26,11 @@ def process_image(image, mask, progress=gr.Progress(track_tqdm=True)):
         print("Mask Size : ", mask["background"].shape)
         print("Unique values: ", np.unique(mask["background"]))
         print("Type : ", type(mask["background"]))
+        mask = mask["background"]
         output = process_inpaint(image["background"], mask["background"])
     img_output = Image.fromarray(output).convert("RGB")
     
-    return img_output
+    return img_output, mask
 
 
 with gr.Blocks() as demo:
@@ -41,6 +43,7 @@ with gr.Blocks() as demo:
             button = gr.Button("Remove")
         with gr.Column():
             output = gr.Image(format="png", label="Output Image")
-        button.click(fn=process_image, inputs=[image, mask], outputs=[output])
+            mask_out = gr.Image(format="png", label="Output Image")
+        button.click(fn=process_image, inputs=[image, mask], outputs=[output, mask_out])
 
 demo.launch(debug=True,show_error=True)