Upload mask_app.py

mask_app.py

@@ -0,0 +1,125 @@
+import os
+import gradio as gr
+from gradio_imageslider import ImageSlider
+from loadimg import load_img
+#import spaces
+from transformers import AutoModelForImageSegmentation
+import torch
+from torchvision import transforms
+import numpy as np
+from PIL import Image
+
+# 检查 CUDA 是否可用
+if torch.cuda.is_available():
+    device = "cuda"
+else:
+    device = "cpu"
+
+torch.set_float32_matmul_precision(["high", "highest"][0])
+
+birefnet = AutoModelForImageSegmentation.from_pretrained(
+    "briaai/RMBG-2.0", trust_remote_code=True
+)
+birefnet.to(device)
+transform_image = transforms.Compose(
+    [
+        transforms.Resize((1024, 1024)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+    ]
+)
+
+output_folder = 'output_images'
+if not os.path.exists(output_folder):
+    os.makedirs(output_folder)
+
+# 定义颜色列表，每个颜色对应一个 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 fn(image, mask_color):
+    im = load_img(image, output_type="pil")
+    im = im.convert("RGB")
+    origin = im.copy()
+    image, mask = process(im, mask_color)    
+    image_path = os.path.join(output_folder, "no_bg_image.png")
+    mask_path = os.path.join(output_folder, "mask_image.png")
+    image.save(image_path)
+    mask.save(mask_path)
+    return (image, origin), image_path, mask
+
+#@spaces.GPU
+def process(image, mask_color):
+    image_size = image.size
+    input_images = transform_image(image).unsqueeze(0).to(device)
+    # Prediction
+    with torch.no_grad():
+        preds = birefnet(input_images)[-1].sigmoid().cpu()
+    pred = preds[0].squeeze()
+    pred_pil = transforms.ToPILImage()(pred)
+    mask = pred_pil.resize(image_size)
+    
+    # 创建一个新的透明背景图像
+    transparent_image = Image.new("RGBA", image_size, (0, 0, 0, 0))
+    transparent_image.paste(image, (0, 0), mask)
+    
+    # 创建一个带有颜色的 mask 图像
+    mask_color_rgb = tuple(int(mask_color[i + 1:i + 3], 16) for i in (0, 2, 4))
+    colored_mask = Image.new("RGBA", image_size, mask_color_rgb + (255,))
+    colored_mask.putalpha(mask)
+    
+    return transparent_image, colored_mask
+
+# 示例数据
+example_image = "giraffe.jpg"  # 确保该文件存在于当前目录
+example_url = "http://farm9.staticflickr.com/8488/8228323072_76eeddfea3_z.jpg"
+
+# 定义 Gradio 组件
+with gr.Blocks() as demo:
+    gr.Markdown("# 🖼️ RMBG-2.0 for Background Removal")
+    with gr.Row():
+        # 左侧列：输入
+        with gr.Column():
+            gr.Markdown("## Input")
+            image_input = gr.Image(label="Upload an image")
+            text_input = gr.Textbox(label="Paste an image URL")
+            color_input = gr.Dropdown(label="Mask Color", choices=colors[1:], value=colors[1])
+            run_button = gr.Button("Run")
+
+        # 右侧列：输出
+        with gr.Column():
+            gr.Markdown("## Output")
+            slider_output = ImageSlider(label="RMBG-2.0", type="pil")
+            file_output = gr.File(label="Output PNG File")
+            mask_output = gr.Image(label="Mask Image")
+
+    # 示例数据
+    gr.Examples(
+        examples=[[example_image, colors[1]], [example_url, colors[1]]],
+        inputs=[image_input, color_input],
+        outputs=[slider_output, file_output, mask_output],  # 添加 outputs 参数
+        fn=fn,
+        cache_examples=True
+    )
+
+    # 绑定事件
+    run_button.click(
+        fn=fn,
+        inputs=[image_input, color_input],
+        outputs=[slider_output, file_output, mask_output]
+    )
+
+if __name__ == "__main__":
+    demo.launch(share=True, show_error=True)