rembg+gradio 在线抠图服务部署

利用 gradio 实现 rembg 的抠图服务的接口化。实际的代码如下。

# @author helsonlin
# @date 11/28 24
import gradio as gr
import os
from rembg import remove, new_session
from PIL import Image, ImageDraw
import numpy as np

# 自定义模型的位置
basedir = os.path.abspath(os.path.dirname(__file__))
os.environ['U2NET_HOME'] = basedir + '/models'



def segment_image(image, model_type):
    session = new_session(model_type)
    # 使用 rembg 进行抠图
    output_image = remove(image, session=session)
    return output_image


# 创建 Gradio 接口
iface = gr.Interface(
    fn=segment_image,
    inputs=[
        gr.Image(type="numpy", label="选择图片"),
        gr.Dropdown(["u2net", "u2net_human_seg", "silueta", "isnet-general-use"], label="选择抠图模型", value="u2net")
    ],
    outputs=gr.Image(type="numpy", label="输出图片", format="png"),
    title="使用 Rembg 进行图像抠图",
    description="选择模型并上传图像进行抠图"
)

# 启动接口
iface.launch(share=True)

依赖库如下：

rembg==2.0.60
numpy==2.0.2
gradio==5.6.0
pillow==10.4.0
onnxruntime==1.2.1

Docker部署

这里/tmp/models 为我的模型存放的文件夹

拉取镜像

docekr pull h55205l/rembg

准备模型

从rembg Github下载模型，解压直接后丢在一个公共的文件内，文件夹内（/tmp/models）的文件应该都是.onnx 的模型文件。

运行容器

docker run -d -p 7860:7860 -v /tmp/models:/app/models h55205l/rembg:latest

访问测试

打开 http://localhost:7860 查看服务是否正常。