Rembg WebUI插件开发：自定义功能扩展

1. 智能万能抠图 - Rembg

在图像处理与内容创作领域，自动去背景已成为AI赋能生产力的典型代表。传统手动抠图耗时耗力，而基于深度学习的语义分割技术则实现了“一键透明化”的高效流程。其中，Rembg作为开源社区中广受欢迎的去背景工具，凭借其高精度、通用性强和轻量化部署优势，广泛应用于电商修图、设计辅助、AI绘画预处理等场景。

Rembg 的核心模型基于U²-Net（U-square Net）架构，这是一种专为显著性目标检测设计的双编码器-解码器结构网络。它能够在无需任何标注输入的情况下，自动识别图像中的主体对象，并生成高质量的 Alpha 透明通道。相比传统人像专用模型（如 MODNet、PortraitNet），U²-Net 更加“通用”，对非人类主体（如宠物、商品、机械部件）也具备出色的边缘保留能力。

随着 AIGC 工具链的普及，越来越多用户希望将 Rembg 集成到本地工作流中，尤其是与Stable Diffusion WebUI等主流创作平台结合使用。因此，开发一个可扩展、易定制的Rembg WebUI 插件，不仅提升了操作便捷性，也为后续功能增强（如批量处理、API 调用、后处理优化）提供了工程基础。

💡本系列文章聚焦于如何基于 Rembg 开发 Stable Diffusion WebUI 插件，实现无缝集成与功能扩展，帮助开发者构建属于自己的智能图像预处理模块。

2. 基于Rembg(U2NET)模型的功能架构解析

2.1 核心模型机制：U²-Net 工作原理简析

U²-Net 是一种层级式嵌套编码-解码结构（ReSidual U-blocks, RSUs），其最大特点是引入了多尺度特征融合和深层监督机制。整个网络由两个U型结构组成：

第一层U：主干U-Net结构，负责整体轮廓提取；
第二层U：每个RSU内部又是一个小型U-Net，增强局部细节感知能力。

这种“U within U”的设计使得模型既能捕捉全局上下文信息，又能精细还原边缘纹理（如毛发、半透明区域）。输出端通过 SOD（Salient Object Detection）任务训练，直接预测像素级透明度值（0~1），最终生成带Alpha通道的PNG图像。

# 示例：U²-Net 推理核心逻辑（简化版） import numpy as np from rembg import remove def remove_background(input_path, output_path): with open(input_path, "rb") as i: with open(output_path, "wb") as o: input_data = i.read() output_data = remove(input_data) # 调用rembg库 o.write(output_data)

该代码展示了rembg库最基础的调用方式，底层实际加载的是 ONNX 格式的 U²-Net 模型，在 CPU 上即可高效运行，适合无GPU环境部署。

2.2 WebUI 集成架构设计

为了将 Rembg 功能嵌入 Stable Diffusion WebUI，需遵循其插件系统规范（位于extensions/目录下）。基本目录结构如下：

extensions/ └── sd-webui-rembg/ ├── extension.py # 插件入口 ├── scripts/ │ └── rembg_ui.py # UI与逻辑控制 ├── lib/ │ └── rembg_processor.py # 核心处理逻辑 └── models/ # 存放ONNX模型文件 └── u2net.onnx

关键组件说明：

组件	作用
`extension.py`	声明插件名称、版本、钩子函数
`rembg_ui.py`	使用 Gradio 构建界面，绑定事件回调
`rembg_processor.py`	封装图像读取、模型推理、后处理逻辑
`u2net.onnx`	预训练模型，支持跨平台推理

得益于 ONNX Runtime 的跨平台兼容性，该插件可在 Windows、Linux、macOS 上稳定运行，且支持 CPU 加速，极大降低了使用门槛。

3. 自定义功能扩展实践

3.1 扩展需求分析：从基础抠图到智能预处理

虽然默认的 Rembg 插件已支持上传图片并去除背景，但在实际应用中仍存在以下痛点：

❌ 不支持批量处理多张图像
❌ 缺乏后处理选项（如边缘平滑、阴影保留）
❌ 无法与其他WebUI功能联动（如送入 img2img 或 ControlNet）

为此，我们提出三项可落地的自定义功能扩展方案：

批量图像去背
参数化后处理调节
一键发送至 SD 主功能区

3.2 实现一：支持文件夹级批量处理

修改scripts/rembg_ui.py中的 Gradio 界面，增加“输入目录”和“输出目录”字段，并调用批处理函数。

# rembg_ui.py 片段 import gradio as gr from lib.rembg_processor import batch_remove_background def on_run_batch(input_dir, output_dir, model_name): try: processed_count = batch_remove_background(input_dir, output_dir, model=model_name) return f"✅ 成功处理 {processed_count} 张图像" except Exception as e: return f"❌ 处理失败: {str(e)}" with gr.Blocks() as ui: gr.Markdown("## 📦 批量去背景") with gr.Row(): in_dir = gr.Textbox(label="输入文件夹路径") out_dir = gr.Textbox(label="输出文件夹路径") model_choice = gr.Dropdown(["u2net", "u2netp"], label="选择模型", value="u2net") run_btn = gr.Button("开始处理") result_msg = gr.Textbox(label="处理结果") run_btn.click( fn=on_run_batch, inputs=[in_dir, out_dir, model_choice], outputs=result_msg )

对应lib/rembg_processor.py中实现批量遍历逻辑：

# lib/rembg_processor.py import os from PIL import Image from rembg import remove def batch_remove_background(input_dir, output_dir, model="u2net"): if not os.path.exists(output_dir): os.makedirs(output_dir) count = 0 for filename in os.listdir(input_dir): if filename.lower().endswith(("png", "jpg", "jpeg", "webp")): input_path = os.path.join(input_dir, filename) output_path = os.path.join(output_dir, f"{os.path.splitext(filename)[0]}.png") with open(input_path, "rb") as inp: img_data = inp.read() result = remove(img_data, model_name=model) with open(output_path, "wb") as out: out.write(result) count += 1 return count

✅优势：无需逐张上传，适用于电商商品图批量精修。

3.3 实现二：添加后处理参数调节

通过rembg.remove()支持部分参数配置，例如：

alpha_matting：启用Alpha抠图（更精细）
alpha_matting_erode_size：腐蚀大小，控制边缘紧致度
session：复用推理会话，提升性能

# 参数化调用示例 result = remove( data=img_data, alpha_matting=True, alpha_matting_erode_size=6, session=new_session(model_name="u2net") )

在 WebUI 中暴露这些参数为滑块控件：

with gr.Accordion("🔧 高级设置"): use_matting = gr.Checkbox(True, label="启用Alpha抠图") erode_slider = gr.Slider(0, 10, value=6, step=1, label="边缘腐蚀强度")

用户可根据图像复杂度动态调整，避免出现“毛边残留”或“主体缺失”。

3.4 实现三：与 Stable Diffusion 主流程联动

利用 WebUI 提供的send_to_img2img或send_to_inpaint接口，实现“一键发送”功能。

# 在Gradio界面上添加按钮 send_to_img2img = gr.Button("➡️ 发送到 img2img") send_to_img2img.click( fn=lambda x: x, # 此处x为输出图像 _js="sendImageToImg2imgTab", # 调用前端JS函数 inputs=output_image, outputs=None )

前端需注册 JS 函数（存于javascript/目录）：

// javascript/rembg_send.js function sendImageToImg2imgTab(img) { const tab = gradioApp().querySelector('button[value="img2img"]'); if (tab) tab.click(); setTimeout(() => { const target = gradioApp().querySelector('#img2img_image input[type="file"]'); if (target) { // 触发粘贴图像逻辑（模拟拖拽） const dt = new DataTransfer(); const file = new File([img.src], "rembg_output.png", { type: "image/png" }); dt.items.add(file); target.dispatchEvent(new Event("change", { bubbles: true })); } }, 500); }

此功能打通了“预处理 → 生成 → 编辑”的完整闭环，显著提升创作效率。