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

def brightness_to_white_alpha(img, bg_img=None, invert_opacity=False):
    if img is None:
        return None

    # Step 1: Convert input to RGB and compute brightness
    img_rgb = img.convert("RGB")
    img_np = np.array(img_rgb).astype(np.float32)
    brightness = np.dot(img_np[..., :3], [0.299, 0.587, 0.114])

    if invert_opacity:
        brightness = 255 - brightness

    alpha = brightness.clip(0, 255).astype(np.uint8)

    # Step 2: Create solid white RGBA image
    height, width = alpha.shape
    white_rgb = np.full((height, width, 3), 255, dtype=np.uint8)
    white_rgba = np.dstack((white_rgb, alpha))  # shape (H, W, 4)
    white_img = Image.fromarray(white_rgba, mode="RGBA")

    # Step 3: Overlay white image on background (if provided)
    if bg_img:
        bg = bg_img.convert("RGBA").resize(white_img.size)
        composite = Image.alpha_composite(bg, white_img)
        return composite

    return white_img

# Gradio UI
iface = gr.Interface(
    fn=brightness_to_white_alpha,
    inputs=[
        gr.Image(type="pil", label="Input Image"),
        gr.Image(type="pil", label="Optional Background Image"),
        gr.Checkbox(label="Invert Brightness → Alpha", value=False)
    ],
    outputs=gr.Image(type="pil", label="Final Output"),
    title="💡 White Brightness Layer + Optional Background",
    description="Generates a pure white image with brightness as transparency. Optionally overlays it on a background."
)

if __name__ == "__main__":
    iface.launch()