Update app.py

app.py

@@ -49,19 +49,21 @@ def process_image(img):
         masks.append(mask)
 
     # 生成合并后的 mask（使用对应颜色的 RGB 值）
-    merged_mask = np.zeros((height, width, 3), dtype=np.uint8)  # RGB 图像
+    merged_mask = np.zeros((height, width, 4), dtype=np.uint8)  # RGBA 图像
     for i, mask in enumerate(masks):
-        merged_mask[mask] = palette[i]  # 使用对应颜色填充
+        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.ones((height, width, 3), dtype=np.uint8) * 255  # 白色背景
-        mask_array[mask] = palette[i]  # 使用对应颜色填充
-        mask_image = Image.fromarray(mask_array, mode='RGB')
+        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='RGB')
+    merged_mask_image = Image.fromarray(merged_mask, mode='RGBA')
 
     # 保存 mask 到本地
     output_dir = "output_masks"
@@ -78,7 +80,7 @@ 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.
+    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():
@@ -121,7 +123,7 @@ with gr.Blocks() as demo:
     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.
+    - Each mask is saved as an RGBA image in the `output_masks` directory.
     - A merged mask is also generated for your convenience.
     ## 🚀 **Enjoy!**
     """)