Create app.py

app.py

@@ -0,0 +1,132 @@
+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, 3), dtype=np.uint8)  # RGB 图像
+    for i, mask in enumerate(masks):
+        merged_mask[mask] = palette[i]  # 使用对应颜色填充
+
+    # 将 mask 转换为 PIL 图像（使用对应颜色的 RGB 值）
+    mask_images = []
+    for i, mask in enumerate(masks):
+        mask_array = np.zeros((height, width, 3), dtype=np.uint8)  # RGB 图像
+        mask_array[mask] = palette[i]  # 使用对应颜色填充
+        mask_image = Image.fromarray(mask_array, mode='RGB')
+        mask_images.append(mask_image)
+
+    merged_mask_image = Image.fromarray(merged_mask, mode='RGB')
+
+    # 保存 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 RGB 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 RGB image in the `output_masks` directory.
+    - A merged mask is also generated for your convenience.
+
+    ## 🚀 **Enjoy!**
+    """)
+
+# 启动应用
+demo.launch(debug=True, share=True)