# app.py
import gradio as gr
import numpy as np
from PIL import Image
import os # 示例：如果需要处理文件

# ... (你的图像处理函数，例如 generate_mask) ...
def generate_mask(input_dict):
    # (你的 mask 生成逻辑) 
    # 假设返回一个 PIL Image 对象
    original_image = input_dict['image']
    sketch_mask = input_dict['mask'] 
    if sketch_mask.mode == 'RGBA':
        alpha_channel = sketch_mask.split()[3]
        binary_mask_np = np.array(alpha_channel)
        binary_mask_np = np.where(binary_mask_np > 0, 255, 0).astype(np.uint8)
        binary_mask_pil = Image.fromarray(binary_mask_np, mode='L')
        return binary_mask_pil
    else:
        return sketch_mask

# 定义 Gradio 界面
iface = gr.Interface(
    fn=generate_mask,
    inputs=gr.Image(type="pil", tool="sketch", label="Draw Mask on Image", source="upload"), 
    outputs=gr.Image(type="pil", label="Generated Mask"),
    title="Image Mask Generator",
    description="Upload an image and draw on it to create a mask."
)

# 启动应用 (在 Hugging Face Spaces 上不需要指定 share=True 或 server_name)
iface.launch()