import gradio as gr
import imageio

def create_mask(image):
    imageio.imwrite("output_image.png", image["mask"])
    return image["mask"]

with gr.Blocks() as demo:
   gr.Markdown(
       """
       # Image Mask Creator
       Draw on the image to create a mask.
       """
   )

   with gr.Row():
       base_image = gr.Image(tool="sketch", label="Base Image", width=512, height=512, show_label=True)
       mask_image = gr.Image(label="Mask Image", show_label=True)

   btn = gr.Button("Create Mask")
   btn.click(create_mask, inputs=base_image, outputs=mask_image)

demo.launch(debug=True)