import gradio as gr
import numpy as np
import imageio

def process_image(img):
    imageio.imwrite("output_image.png", img["composite"])

    alpha_channel = img["layers"][0][:, :, 3]
    mask = np.where(alpha_channel == 0, 0, 255)

    # return img["background"], mask
    return mask

# with gr.Blocks() as demo:
#   with gr.Row():
#     img = gr.ImageMask(sources=["upload"], layers=False, transforms=[], format="png", label="base image", show_label=True)
#     img1 = gr.Image()
#     img2 = gr.Image(label="mask image", show_label=True)
#   btn = gr.Button()
#   btn.click(dummy, img, [img1, img2])

# Define the Gradio interface
demo = gr.Interface(
    fn=process_image,
    inputs=gr.ImageMask(sources=["upload"], layers=False, transforms=[], format="png", label="base image", show_label=True, height = 600),
    outputs=[
        # gr.Image(label="Background Image"),
        gr.Image(label="Mask Image"),
    ],
)

demo.launch(debug=True)