rembg自定义模型：训练专属背景移除模型的完整流程
项目地址: https://gitcode.com/GitHub_Trending/re/rembg 引言：为什么需要自定义模型？
在图像处理领域，背景移除（Background Removal）是一项基础但至关重要的技术。虽然rembg提供了多种预训练模型，但在特定场景下，通用模型可能无法达到最佳效果。比如：
- 特定行业应用：医疗影像、工业检测、艺术品处理等
- 特殊对象类型：特定品牌产品、特殊材质物体、定制化需求
- 精度要求极高：商业级应用需要99%以上的准确率
- 数据隐私保护：敏感数据不能使用云端API
本文将带你从零开始，完整掌握rembg自定义模型的训练、部署和应用全流程。
技术架构深度解析
rembg模型架构概览
核心Session类结构
环境准备与依赖安装
基础环境要求
# 创建Python虚拟环境
python -m venv rembg-training
source rembg-training/bin/activate
# 安装核心依赖
pip install torch==2.0.1 torchvision==0.15.2
pip install onnxruntime-gpu==1.15.1 # 或onnxruntime-cpu
pip install opencv-python==4.8.0.74
pip install Pillow==10.0.0
pip install numpy==1.24.3
pip install scikit-image==0.21.0
训练框架选择对比
| 框架 | 优点 | 缺点 | 适用场景 |
|---|---|---|---|
| PyTorch | 生态丰富，调试方便 | 内存占用较大 | 研究、实验 |
| TensorFlow | 生产环境稳定 | API变化频繁 | 大规模部署 |
| MXNet | 内存效率高 | 社区较小 | 资源受限环境 |
推荐使用PyTorch进行模型训练，然后导出为ONNX格式供rembg使用。
数据准备与预处理
数据集构建标准
import os
from pathlib import Path
from PIL import Image
import numpy as np
class SegmentationDataset:
def __init__(self, image_dir, mask_dir, transform=None):
self.image_dir = Path(image_dir)
self.mask_dir = Path(mask_dir)
self.transform = transform
self.image_paths = sorted(list(self.image_dir.glob("*.jpg")) +
list(self.image_dir.glob("*.png")))
self.mask_paths = sorted(list(self.mask_dir.glob("*.png")))
assert len(self.image_paths) == len(self.mask_paths), \
"图像和掩码数量不匹配"
def __len__(self):
return len(self.image_paths)
def __getitem__(self, idx):
image = Image.open(self.image_paths[idx]).convert("RGB")
mask = Image.open(self.mask_paths[idx]).convert("L")
if self.transform:
image = self.transform(image)
mask = self.transform(mask)
return image, mask
数据增强策略
from torchvision import transforms
# 训练数据增强
train_transform = transforms.Compose([
transforms.Resize((320, 320)),
transforms.RandomHorizontalFlip(p=0.5),
transforms.RandomRotation(degrees=10),
transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
# 验证/测试数据转换
val_transform = transforms.Compose([
transforms.Resize((320, 320)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
模型选择与训练策略
可供选择的骨干网络
| 模型架构 | 参数量 | 精度 | 速度 | 适用场景 |
|---|---|---|---|---|
| U²-Net | 44M | ⭐⭐⭐⭐ | ⭐⭐⭐ | 通用场景 |
| U²-Netp | 4.5M | ⭐⭐⭐ | ⭐⭐⭐⭐ | 移动端 |
| BiRefNet | 58M | ⭐⭐⭐⭐⭐ | ⭐⭐ | 高精度需求 |
| MODNet | 6.5M | ⭐⭐⭐⭐ | ⭐⭐⭐⭐ | 实时应用 |
训练代码示例
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from models.u2net import U2NET
def train_model(config):
# 初始化模型
model = U2NET(3, 1)
model = model.to(config['device'])
# 损失函数和优化器
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(model.parameters(),
lr=config['lr'],
weight_decay=config['weight_decay'])
# 学习率调度器
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
optimizer, mode='min', factor=0.5, patience=5
)
# 数据加载
train_dataset = SegmentationDataset(
config['train_image_dir'],
config['train_mask_dir'],
transform=train_transform
)
train_loader = DataLoader(
train_dataset,
batch_size=config['batch_size'],
shuffle=True,
num_workers=config['num_workers']
)
# 训练循环
for epoch in range(config['epochs']):
model.train()
epoch_loss = 0
for images, masks in train_loader:
images = images.to(config['device'])
masks = masks.to(config['device'])
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, masks)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
avg_loss = epoch_loss / len(train_loader)
print(f'Epoch {epoch+1}/{config["epochs"]}, Loss: {avg_loss:.4f}')
# 学习率调整
scheduler.step(avg_loss)
# 保存检查点
if (epoch + 1) % config['save_interval'] == 0:
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': avg_loss,
}, f'checkpoint_epoch_{epoch+1}.pth')
return model
超参数配置表
training_config:
batch_size: 8
learning_rate: 0.001
weight_decay: 0.0001
epochs: 100
save_interval: 10
device: "cuda" # or "cpu"
data_config:
image_size: [320, 320]
train_split: 0.8
val_split: 0.1
test_split: 0.1
model_config:
architecture: "u2net"
in_channels: 3
out_channels: 1
pretrained: true
模型导出与转换
PyTorch到ONNX转换
import torch
import torch.onnx
from models.u2net import U2NET
def convert_to_onnx(pytorch_model_path, onnx_model_path):
# 加载训练好的模型
model = U2NET(3, 1)
checkpoint = torch.load(pytorch_model_path, map_location='cpu')
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()
# 创建示例输入
dummy_input = torch.randn(1, 3, 320, 320)
# 导出ONNX模型
torch.onnx.export(
model,
dummy_input,
onnx_model_path,
export_params=True,
opset_version=12,
do_constant_folding=True,
input_names=['input'],
output_names=['output'],
dynamic_axes={
'input': {0: 'batch_size'},
'output': {0: 'batch_size'}
}
)
print(f"模型已成功导出到: {onnx_model_path}")
# 使用示例
convert_to_onnx("best_model.pth", "custom_u2net.onnx")
ONNX模型优化
# 安装ONNX优化工具
pip install onnxoptimizer onnxruntime
# 优化ONNX模型
python -m onnxoptimizer custom_u2net.onnx optimized_custom_u2net.onnx
# 验证模型有效性
import onnx
model = onnx.load("optimized_custom_u2net.onnx")
onnx.checker.check_model(model)
print("模型验证通过")
自定义Session集成
创建自定义Session类
# custom_session.py
import os
from typing import List
import numpy as np
import onnxruntime as ort
from PIL import Image
from PIL.Image import Image as PILImage
from rembg.sessions.base import BaseSession
class CustomSession(BaseSession):
"""自定义背景移除模型Session类"""
def __init__(self, model_name: str, sess_opts: ort.SessionOptions, *args, **kwargs):
"""
初始化自定义Session
Args:
model_name: 模型名称
sess_opts: ONNX Runtime会话选项
model_path: 自定义模型路径（必须提供）
"""
model_path = kwargs.get("model_path")
if model_path is None:
raise ValueError("model_path参数必须提供")
# 确保模型文件存在
if not os.path.exists(model_path):
raise FileNotFoundError(f"模型文件不存在: {model_path}")
super().__init__(model_name, sess_opts, *args, **kwargs)
def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
"""
使用自定义模型进行预测
Args:
img: 输入图像（PIL Image）
Returns:
分割掩码列表
"""
# 预处理输入图像
input_data = self.normalize(
img,
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
size=(320, 320)
)
# 模型推理
ort_outs = self.inner_session.run(None, input_data)
# 后处理
pred = ort_outs[0][:, 0, :, :]
ma = np.max(pred)
mi = np.min(pred)
pred = (pred - mi) / (ma - mi)
pred = np.squeeze(pred)
# 生成掩码图像
mask = Image.fromarray((pred * 255).astype("uint8"), mode="L")
mask = mask.resize(img.size, Image.Resampling.LANCZOS)
return [mask]
@classmethod
def download_models(cls, *args, **kwargs):
"""
返回模型文件路径（对于自定义模型，直接返回提供的路径）
"""
model_path = kwargs.get("model_path")
if model_path is None:
raise ValueError("model_path参数必须提供")
return os.path.abspath(os.path.expanduser(model_path))
@classmethod
def name(cls, *args, **kwargs):
"""返回模型名称"""
return "custom_model"
注册自定义Session
# __init__.py
from .sessions.custom_session import CustomSession
# 在session_factory.py中添加注册逻辑
def register_custom_sessions():
"""注册所有自定义Session"""
from .sessions.custom_session import CustomSession
# 添加到全局session映射
SESSION_MAP['custom_model'] = CustomSession
部署与使用
命令行使用
# 使用自定义模型进行背景移除
rembg i -m custom_model -x '{"model_path": "/path/to/custom_model.onnx"}' input.jpg output.png
# 批量处理目录中的图像
rembg p -m custom_model -x '{"model_path": "/path/to/custom_model.onnx"}' input_dir/ output_dir/
Python API集成
from rembg import remove, new_session
from PIL import Image
# 创建自定义模型session
custom_session = new_session(
"custom_model",
model_path="/path/to/custom_model.onnx"
)
# 处理单张图像
input_image = Image.open("input.jpg")
output = remove(input_image, session=custom_session)
output.save("output.png")
# 批量处理
import os
from pathlib import Path
input_dir = Path("input_images")
output_dir = Path("output_images")
output_dir.mkdir(exist_ok=True)
for img_file in input_dir.glob("*.jpg"):
output_file = output_dir / f"{img_file.stem}_output.png"
with Image.open(img_file) as img:
result = remove(img, session=custom_session)
result.save(output_file)
Web服务部署
# app.py
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import FileResponse
import tempfile
from rembg import remove, new_session
from PIL import Image
import io
app = FastAPI(title="Custom Background Removal API")
# 初始化自定义模型session
custom_session = new_session(
"custom_model",
model_path="/path/to/custom_model.onnx"
)
@app.post("/remove-background")
async def remove_background(file: UploadFile = File(...)):
"""移除图像背景API端点"""
# 读取上传的图像
image_data = await file.read()
image = Image.open(io.BytesIO(image_data))
# 移除背景
result = remove(image, session=custom_session)
# 保存结果到临时文件
with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp_file:
result.save(tmp_file.name)
return FileResponse(
tmp_file.name,
media_type="image/png",
filename=f"removed_bg_{file.filename}"
)
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000)
性能优化与监控
模型推理优化
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
