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

# 定义全局变量保存 masks
masks = []

# 颜色列表
colors = [(255,0,0), (0,255,0), (0,0,255), (255,255,0), (255,0,255), (0,255,255)]

def add_mask(img_dict, color_idx):
    color_idx = int(color_idx)
    mask = img_dict['mask']
    orig_img = img_dict['image']

    if orig_img.ndim == 2:
        orig_img = np.stack([orig_img]*3, axis=-1)
    if orig_img.dtype != np.uint8:
        orig_img = orig_img.astype(np.uint8)

    img_pil = Image.fromarray(orig_img)
    mask_pil = Image.fromarray(mask).convert("L")

    mask_np = np.array(mask_pil)
    color = colors[color_idx % len(colors)]

    # 创建 RGBA 彩色 mask
    colored_mask_np = np.zeros((mask_np.shape[0], mask_np.shape[1], 4), dtype=np.uint8)
    mask_bool = mask_np > 0
    for c in range(3):
        colored_mask_np[..., c][mask_bool] = color[c]
    colored_mask_np[..., 3][mask_bool] = 100

    colored_mask = Image.fromarray(colored_mask_np, mode="RGBA")

    base = Image.alpha_composite(img_pil.convert("RGBA"), colored_mask)

    # 👇 关键修改：返回一张“空白”的 mask 而不是 None
    empty_mask = np.zeros_like(mask)

    return base, orig_img


def reset_masks():
    global masks
    masks = []
    return None

# 写一个单纯的赋值函数
def assign_value(src):
    return src

with gr.Blocks() as demo:
    with gr.Row():
        # image_input = gr.Image(type="pil", label="上传图片")
        mask_canvas = gr.Image(type="numpy", tool="sketch", interactive=True, label="在上面画 mask")
    
    color_index = gr.Number(value=0, label="当前颜色索引 (自动循环)")
    
    add_btn = gr.Button("添加当前 mask")
    reset_btn = gr.Button("清空所有 mask")
    
    output = gr.Image(type="pil", label="叠加后的 mask 结果")
    # orig_img = 0
    add_btn.click(
        add_mask, 
        inputs=[mask_canvas, color_index], 
        outputs=[output, mask_canvas]
    ).then(
        lambda idx: (idx + 1) % len(colors), 
        inputs=color_index, 
        outputs=color_index
    ).then(
        lambda _: None,  # 清空 mask_canvas
        inputs=None,
        outputs=mask_canvas
    )

    reset_btn.click(reset_masks, outputs=output)
    
demo.launch()