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)
    # mask = np.where(alpha_channel == 0, 0, 255).astype(np.uint8)

    # return img["background"], mask
    # Save the mask as a PNG file
    # output_path = "mask_output.png"
    # imageio.imwrite(output_path, mask)  # Save the mask as a PNG

    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),
    outputs=[
        # gr.Image(label="Background Image"),
        gr.Image(label="Mask Image", format="png"),
    ],
    description="Upload an image and then draw mask on the image selecting Pen icon!",
    article = "For Auto text guided Mask generation, try @ovi054's [Text Guided Inpainting Mask tool](https://huggingface.co/spaces/ovi054/text-guided-mask-for-inpainting)"
)

demo.launch(debug=True)