当前位置：首页 > news >正文

RMBG-2.0代码实例：FastAPI+Uvicorn后端调用逻辑拆解

news 2026/3/26 22:29:07

RMBG-2.0代码实例：FastAPI+Uvicorn后端调用逻辑拆解

1. RMBG-2.0背景移除模型概述

RMBG-2.0是BRIA AI开源的新一代背景移除模型，基于BiRefNet（Bilateral Reference Network）架构。这个模型通过双边参考机制同时建模前景与背景特征，能够实现发丝级精细分割。

模型特点：

支持人像、商品、动物等多场景
单张1024×1024图片处理仅需0.5-1秒（GPU）
采用Transformers框架部署
消费级显卡（24GB）可稳定输出

2. 环境准备与快速部署

2.1 镜像基本信息

镜像名：ins-rmbg-2.0-v1
适用底座：insbase-cuda124-pt250-dual-v7
启动命令：bash /root/start.sh
访问端口：7860

2.2 部署步骤

在平台镜像市场选择本镜像，点击"部署实例"
等待实例状态变为"已启动"（约需1-2分钟初始化）
首次启动需30-40秒加载BiRefNet模型至显存

3. FastAPI后端架构解析

3.1 核心代码结构

from fastapi import FastAPI, UploadFile, File from fastapi.responses import HTMLResponse import torch from transformers import AutoModelForImageSegmentation from PIL import Image import numpy as np import io app = FastAPI() # 模型加载 model = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-2.0") model.eval() torch.set_float32_matmul_precision('high') @app.post("/remove_bg") async def remove_background(file: UploadFile = File(...)): # 图片预处理 image = Image.open(io.BytesIO(await file.read())) image = preprocess_image(image) # 模型推理 with torch.no_grad(): output = model(image) # 后处理 result = postprocess_output(output) return result

3.2 关键组件说明

FastAPI应用初始化：创建FastAPI实例并配置基础路由
模型加载：使用Transformers的AutoModelForImageSegmentation加载预训练模型
请求处理：定义/remove_bg端点处理图片上传
预处理/后处理：包含图片格式转换、归一化等操作

4. Uvicorn服务启动配置

4.1 启动脚本分析

import uvicorn from fastapi import FastAPI app = FastAPI() if __name__ == "__main__": uvicorn.run( app, host="0.0.0.0", port=7860, workers=1, log_level="info" )

4.2 配置参数说明

host：设置为0.0.0.0允许外部访问
port：7860为默认服务端口
workers：设置为1避免多进程显存冲突
log_level：info级别日志记录请求信息

5. 图片处理流程详解

5.1 预处理阶段

def preprocess_image(image): # 转换为RGB格式 if image.mode != "RGB": image = image.convert("RGB") # 调整尺寸 width, height = image.size scale = 1024 / max(width, height) new_size = (int(width * scale), int(height * scale)) image = image.resize(new_size, Image.LANCZOS) # 转换为Tensor image = torch.from_numpy(np.array(image)).permute(2, 0, 1).float() / 255.0 return image.unsqueeze(0).to("cuda")

5.2 后处理阶段

def postprocess_output(output): # 获取mask mask = output.squeeze().cpu().numpy() > 0.5 # 创建透明背景 rgba = np.zeros((*mask.shape, 4), dtype=np.uint8) rgba[..., :3] = 255 # 白色背景 rgba[..., 3] = mask * 255 # alpha通道 # 转换为PIL Image return Image.fromarray(rgba, "RGBA")

6. 前端交互实现

6.1 HTML页面结构

<!DOCTYPE html> <html> <head> <title>RMBG-2.0 背景移除</title> <style> .container { display: flex; gap: 20px; } .preview { border: 1px solid #ddd; padding: 10px; } </style> </head> <body> <h1>RMBG-2.0 背景移除</h1> <input type="file" id="upload" accept="image/*"> <button id="process"> 生成透明背景</button> <div class="container"> <div class="preview"> <h3>原图预览</h3> <img id="original" style="max-width: 500px;"> </div> <div class="preview"> <h3>处理结果</h3> <img id="result" style="max-width: 500px;"> </div> </div> <script src="app.js"></script> </body> </html>

6.2 JavaScript交互逻辑

document.getElementById('upload').addEventListener('change', function(e) { const file = e.target.files[0]; const reader = new FileReader(); reader.onload = function(event) { document.getElementById('original').src = event.target.result; }; reader.readAsDataURL(file); }); document.getElementById('process').addEventListener('click', async function() { const fileInput = document.getElementById('upload'); if (!fileInput.files.length) return; const formData = new FormData(); formData.append('file', fileInput.files[0]); const response = await fetch('/remove_bg', { method: 'POST', body: formData }); const blob = await response.blob(); document.getElementById('result').src = URL.createObjectURL(blob); });