Update app.py

app.py

@@ -1,27 +1,22 @@
 import gradio as gr
+import numpy as np
 import imageio
 
 def dummy(img):
-    imageio.imwrite("output_image.png", img["mask"])
-    return img["image"], img["mask"]
+    # Assuming img is a dictionary with keys 'composite', 'background', and 'layers'
+    composite_image = img["composite"]
+    imageio.imwrite("output_image.png", composite_image)
+
+    alpha_channel = img["layers"][0][:, :, 3]
+    mask = np.where(alpha_channel == 0, 0, 255).astype(np.uint8)
+
+    return mask
 
 with gr.Blocks() as demo:
-    label_text = """This page is only used to generate MASK mask images and cannot be used to submit tasks.
-    本页面仅用于生成MASK遮罩图片，无法用于提交任务
-    Click the RUN button to generate a MASK IMAGE and send it back to the BRAUNDRESS bot
-    点击RUN按钮生成MASK IMAGE后将其发送回BRAUNDRESS机器人"""
-    gr.Label(label_text, font_size=16, font_family="Arial", padding=10)
-    backup_url_text = "backup mask maker: https://huggingface.co/spaces/stevhliu/inpaint-mask-maker  https://huggingface.co/spaces/r3gm/inpaint-mask-maker https://huggingface.co/spaces/daethyra/inpaint-mask-maker"
-    gr.Label(backup_url_text, font_size=10, font_family="Arial", padding=10)
-    with gr.Row():
-        img = gr.Image(
-            tool="sketch", label="base image", show_label=True,
-            style={"height": "500px", "max-width": "100%", "@media (max-width: 768px)": {"height": "300px"}}
-        )
     with gr.Row():
-        img1 = gr.Image()
+        img = gr.ImageMask(sources=["upload"], layers=True, transforms=[], format="png", label="base image", show_label=True)
         img2 = gr.Image(label="mask image", show_label=True)
-    btn = gr.Button()
-    btn.click(dummy, img, [img1, img2])
+    btn = gr.Button("Generate Mask")
+    btn.click(dummy, inputs=img, outputs=img2)
 
-demo.launch(debug=True)
+demo.launch(debug=True)