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

# 定义颜色列表，每个颜色对应一个 mask
colors = [
    '#000000',  # 背景色
    '#2692F3',  # 蓝色
    '#F89E12',  # 橙色
    '#16C232',  # 绿色
    '#F92F6C',  # 粉色
    '#AC6AEB',  # 紫色
]

# 将颜色转换为 RGB 值
palette = np.array([
    tuple(int(s[i + 1:i + 3], 16) for i in (0, 2, 4))
    for s in colors[1:]  # 跳过背景色
])  # (N, 3)

# 计算颜色距离
def color_distance(c1, c2):
    return np.sqrt(np.sum((c1 - c2) ** 2, axis=-1))

# 处理用户绘制的图像，生成 mask
def process_image(img):
    # 保存用户上传的图像（包含绘制的掩码）
    imageio.imwrite("output_image.png", img["composite"])

    # 提取用户绘制的图层
    user_input = np.asarray(img["layers"][0])  # (H, W, 4)
    height, width, _ = user_input.shape

    # 提取 alpha 通道作为前景 mask
    alpha_channel = user_input[..., 3]  # (H, W)
    foreground_mask = alpha_channel > 0  # 用户绘制的区域

    # 提取 RGB 通道
    user_input_rgb = user_input[..., :3]  # (H, W, 3)

    # 使用颜色距离生成 mask
    masks = []
    for i, color in enumerate(palette):
        distance = color_distance(user_input_rgb, color)
        mask = (distance < 100)  # 调整颜色偏差阈值
        mask = mask * foreground_mask  # 只保留用户绘制的区域
        masks.append(mask)

    # 生成合并后的 mask（使用对应颜色的 RGB 值）
    merged_mask = np.zeros((height, width, 4), dtype=np.uint8)  # RGBA 图像
    merged_mask[..., 3] = 255  # 设置背景为完全不透明（黑色背景）
    for i, mask in enumerate(masks):
        merged_mask[mask, :3] = palette[i]  # 使用对应颜色填充 RGB 通道
        merged_mask[mask, 3] = 255  # 设置前景为完全不透明

    # 将 mask 转换为 PIL 图像（使用对应颜色的 RGB 值）
    mask_images = []
    for i, mask in enumerate(masks):
        mask_array = np.zeros((height, width, 4), dtype=np.uint8)  # RGBA 图像
        mask_array[mask, :3] = palette[i]  # 使用对应颜色填充 RGB 通道
        mask_array[mask, 3] = 255  # 设置前景为完全不透明
        mask_image = Image.fromarray(mask_array, mode='RGBA')
        mask_images.append(mask_image)

    merged_mask_image = Image.fromarray(merged_mask, mode='RGBA')

    # 保存 mask 到本地
    output_dir = "output_masks"
    os.makedirs(output_dir, exist_ok=True)
    for i, mask_image in enumerate(mask_images):
        mask_image.save(f"{output_dir}/mask_{i + 1}.png")
    merged_mask_image.save(f"{output_dir}/merged_mask.png")

    print(f"Masks saved to {output_dir}")
    return img["background"], *mask_images, merged_mask_image

# Gradio 界面
with gr.Blocks() as demo:
    # 添加 HTML 标题和介绍
    gr.Markdown("""
    # 🎨 **Multi-Color Mask Generator**
    Welcome to the **Multi-Color Mask Generator**! This tool allows you to upload an image, draw masks using different colors, and generate corresponding mask images. Each mask is saved as an RGBA image, and a merged mask is also created for your convenience.
    """)

    with gr.Row():
        # 左侧：用户上传图像并绘制掩码
        with gr.Column():
            gr.Markdown("### 🖌️ **Draw Your Mask**")
            img = gr.ImageEditor(
                image_mode='RGBA',
                sources=['upload'],  # 允许用户上传图像
                transforms=[],       # 禁用变换功能
                brush=gr.Brush(
                    colors=colors[1:],  # 设置画笔颜色
                    color_mode="fixed",
                ),
                type='pil',
                label="Draw Mask",
                height=800,  # 设置较高的纵向高度
                width=600,   # 设置宽度
            )
        # 右侧：显示处理后的图像
        with gr.Column():
            gr.Markdown("### 🖼️ **Generated Masks**")
            img1 = gr.Image(label="Background Image", show_label=True)  # 显示背景图像
            img2 = gr.Image(label="Mask 1 (Blue)", show_label=True)     # 显示蓝色掩码
            img3 = gr.Image(label="Mask 2 (Orange)", show_label=True)   # 显示橙色掩码
            img4 = gr.Image(label="Mask 3 (Green)", show_label=True)    # 显示绿色掩码
            img5 = gr.Image(label="Mask 4 (Pink)", show_label=True)     # 显示粉色掩码
            img6 = gr.Image(label="Mask 5 (Purple)", show_label=True)   # 显示紫色掩码
            img7 = gr.Image(label="Merged Mask", show_label=True)       # 显示合并后的掩码

    # 按钮：触发图像处理
    btn = gr.Button("Generate Masks")
    btn.click(process_image, img, [img1, img2, img3, img4, img5, img6, img7])

    # 添加 HTML 结尾
    gr.Markdown("""
    ## 📝 **How to Use**
    1. Upload an image using the **Draw Mask** section.
    2. Use the brush tool to draw masks with different colors.
    3. Click the **Generate Masks** button to create and save the masks.
    4. View the generated masks on the right side.
    ## 📂 **Output**
    - Each mask is saved as an RGBA image in the `output_masks` directory.
    - A merged mask is also generated for your convenience.
    ## 🚀 **Enjoy!**
    """)

# 启动应用
demo.launch(debug=True, share=True)