import gradio as gr
import numpy as np
from PIL import Image
import tempfile
import os

def generate_binary_mask(image_data):
    if image_data is None or "layers" not in image_data or not image_data["layers"]:
        raise gr.Error("Please draw a mask before generating!")

    mask = image_data["layers"][0]
    mask_array = np.array(mask)

    if np.all(mask_array < 10):
        raise gr.Error("The mask is empty! Please draw something.")

    # Binary mask logic
    is_black = np.all(mask_array < 10, axis=2)
    binary_mask = Image.fromarray(((~is_black) * 255).astype(np.uint8))

    # Save to temporary file
    temp_dir = tempfile.mkdtemp()
    output_path = os.path.join(temp_dir, "binary_mask.png")
    binary_mask.save(output_path)

    return binary_mask, output_path

with gr.Blocks() as app:
    with gr.Row():
        with gr.Column():
            image_input = gr.ImageMask(
                label="Upload or Paste Image, then draw mask",
                type="pil",
                height=None,
                width=None
            )
            generate_btn = gr.Button("Generate Mask")
        
        with gr.Column():
            mask_preview = gr.Image(label="Mask Preview", type="pil")
            output_file = gr.File(label="Download Mask")

    generate_btn.click(fn=generate_binary_mask, inputs=image_input, outputs=[mask_preview, output_file])

app.launch()