37K 下载！开源图像背景去除项目，精确到发丝，数字人福音 - 文章 - 开发者社区

picture.image

猴哥的第 136 期分享，欢迎追看

最近 AI 数字人领域进展迅速，多款开源方案，效果炸裂，直逼付费方案-HeyGen！

我们在各大短视频平台刷到的数字人，背后的制作流程大致包括:

底模数字人视频生成
AI 换脸（FaceFusion 等）
背景去除
背景合成

作为数字人背后的关键技术--背景去除 ，很大程度上决定了最终生成视频的真实度。

今天，我们不聊数字人，分享一款开源免费、轻量高效 的背景去除方案 - RMBG-2.0，并带大家本地部署实操。

项目地址：https://github.com/ai-anchorite/BRIA-RMBG-2.0

1.RMBG-2.0 亮点

老规矩，简单介绍下项目亮点：

高精度背景移除 ：能够精确识别并移除复杂图像中的背景，提供边缘清晰、自然度高的专业级图像质量。
丰富多样的训练数据 ：该模型在超过 15,000 张涵盖多种领域的高质量图像上进行训练，确保了其准确性和广泛的适用性。
高效的处理速度 ：单张 1024x1024 图像，GPU推理耗时约 0.15s。

先看下官方测试效果：

picture.image

感兴趣的朋友，可前往官方体验地址：https://huggingface.co/spaces/briaai/BRIA-RMBG-2.0

测试效果如下：

picture.image

本地部署

项目开源了模型权重，为此，可下载到本地自行部署，使用也非常简单，几行代码搞定！

2.1 依赖安装

项目依赖以下仓库：

  
torch  
torchvision  
pillow  
kornia  
transformers

可选择新建 requirents.txt，填入上述依赖库，一键安装。

  
pip install -r requirents.txt

2.2 权重下载

模型权重托管在 huggingface 上：

模型下载：https://huggingface.co/briaai/RMBG-2.0

考虑到国内访问 huggingface 比较麻烦，推荐大家从 ModelScope 下载：

  
git lfs install  
git clone https://www.modelscope.cn/AI-ModelScope/RMBG-2.0.git

2.2 本地推理

推理示例代码：

  
from PIL import Image  
import torch  
from torchvision import transforms  
from transformers import AutoModelForImageSegmentation  
  
model = AutoModelForImageSegmentation.from_pretrained('RMBG-2.0', trust_remote_code=True)  
torch.set_float32_matmul_precision(['high', 'highest'][0])  
model.to('cuda')  
model.eval()  
# Data settings  
transform_image = transforms.Compose([  
    transforms.Resize((1024, 1024)),  
    transforms.ToTensor(),  
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  
])  
image = Image.open('elon-musk.jpg')  
input_images = transform_image(image).unsqueeze(0).to('cuda')  
  
# Prediction  
with torch.no_grad():  
    for i in range(10):  
        preds = model(input_images)[-1].sigmoid().cpu()  
pred = preds[0].squeeze()  
pred_pil = transforms.ToPILImage()(pred)  
mask = pred_pil.resize(image.size)  
image.putalpha(mask)  
  
image.save("elon-musk\_no\_bg\_image.png")

模型预设 1024x1024 输入，在单张 4080 显卡上耗时统计如下，基本稳定在 0.15s/img ：

  
Inference time: 0.147s  
Inference time: 0.147s  
Inference time: 0.150s  
Inference time: 0.147s  
Inference time: 0.147s

显存占用如何？

  
1  NVIDIA GeForce RTX 4080        On  |   00000000:3B:00.0 Off |                  N/A |  
| 38%   28C    P2             86W /  320W |    4667MiB /  16376MiB |      0%      Default