import gradio as gr
import numpy as np
import imageio

def dummy(img):
    # 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:
    with gr.Row():
        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("Generate Mask")
    btn.click(dummy, inputs=img, outputs=img2)

demo.launch(debug=True)