​Win10下RMBG‑2.0纯净官方源码包的部署及调用实践

RMBG‑2.0 是 BRIA AI 在 2024 年发布的最新一代开源背景去除模型，相比 v1.4 在准确率上从 73.26% 提升至 90.14%，采用 BiRefNet 双边参考架构，在高分辨率图像分割任务上达到业界领先水平​。本文将从官方纯净源码包出发，分别演示在 Windows10 + VSCode 环境下的 CPU 与 GPU 两种推理配置，及完成模型启动后的测试调用。

目录

  1. 环境准备

  2. 下载官方源码

  3. CPU 模式部署

  4. GPU 模式部署

  5. 总结与建议

  6. 代码及工具包

环境准备

  • 操作系统：Windows 10

  • 编辑器：VSCode

  • Python 环境：如未安装请从​​​ https://www.python.org/ftp/python/3.13.3/python-3.13.3-amd64.exe下载安装，建议 v3.10+则可。

  • 依赖库：torch、torchvision、pillow、kornia、

  • 建议以不同的虚拟环境分别安装及配置CPU版和CUDA版，本文CPU版使用虚拟环境rmbg_env，CUDA版使用虚拟环境rmbg_env_cuda，激活虚拟环境后，安装transformers和timm库，下面命令中指定版本号是因为在我的Python环境中，这两个版本能正常兼容工作，如不需要指定版本号安装都能正常工作，则不必为这两个库指定版本号安装：

    pip install kornia transformers==4.48.1 timm==0.9.16
  • 即为 RMBG‑2.0 在源码包要求的最小依赖列表​。

  • CPU 版本 Transformers+PyTorch 安装：

    pip install "transformers[torch]"

    此命令会在仅需 CPU 支持时一并安装 PyTorch。

  • GPU 版本：

    1. 安装对应 CUDA 驱动（打开PowerShell，执行：

      nvidia-smi

       查看返回结果，在表格中能看到 CUDA Version 字段，它就是你的显卡驱动支持的 ​​最高 CUDA 版本（例如我的返回信息是​​：

      ），而此处的最高支持版本并不一定是我们实际要安装的版本，实际安装的版本具体还是由 PyTorch 决定，毕竟只有脚最清楚穿什么鞋子最舒服嘛）。此外，Driver Version 字段是你的显卡当前的驱动版本，具体的NVIDIA 驱动与 CUDA 版本对照表可以参考Nvidia文档：CUDA 12.8 Update 1 发行说明 — 发行说明 12.8 文档

    2. 安装支持 CUDA 的 Python版深度机器学习、视觉、图像处理等库，如：torch torchvision pillow等库，这里需要注意的是，必须通过PyTorch网站指定CUDA版本号进行索引安装，这样才能正常适配。经反复测试，发现CUDA  v1.24对应的上述库最适合本机的虚拟环境（建议多尝试不同版本，找出最适合自己显卡的版本，结合CUDA版代码测试，若无明显异常则视为最佳选择），安装如下：

    pip install torch torchvision pillow --index-url https://download.pytorch.org/whl/cu124

  • 若有条件，可参考 Hugging Face 官方文档中关于 GPU 加速的说明​。

下载官方源码

  1. 打开命令行，切换到项目准备目录，例如 D:\projects\test_rmbg。

  2. 克隆源码仓库（官方纯净包），若无法克隆成功，或者发现克隆后ONNX模型文件很小（仅文件指针，网络或源码官方的其它原因……），则请移步本文结尾处提供的连接下载。

    # 确保已安装 Git LFS
git lfs install
git clone https://github.com/Ttawannn/RMBG-2.0.git
cd RMBG-2.0

    该仓库包含模型代码与权重下载脚本等，保持与 BRIA 官方保持一致。

CPU 模式部署

  1. 在 VSCode 中打开 D:\projects\test_rmbg 目录，从VSCode中选择菜单“终端”->“新建终端”，在终端中默认路径是 D:\projects\test_rmbg（如果不是则自行切换），执行执行命令创建并激活虚拟环境：

    python -m venv rmbg_env
.\rmbg_env\Scripts\activate

    接着安装依赖（见上）。

  2. 新建Python代码文件 run_cpu.py，并加入以下代码：

     1 from transformers import AutoModelForImageSegmentation
 2 import torch
 3 from PIL import Image
 4 from torchvision import transforms
 5 import numpy as np
 6 from skimage.morphology import binary_dilation
 7 from scipy.ndimage import gaussian_filter
 8 
 9 # 加载模型
10 model = AutoModelForImageSegmentation.from_pretrained(
11     'briaai/RMBG-2.0', trust_remote_code=True
12 )
13 model.to('cpu')
14 model.eval()
15 
16 # 图像预处理
17 transform = transforms.Compose([
18     transforms.Resize((1024, 1024)),
19     transforms.ToTensor(),
20     transforms.Normalize([0.485, 0.456, 0.406],
21                          [0.229, 0.224, 0.225])
22 ])
23 
24 # 进行推理
25 img = Image.open('input1.jpg')
26 inp = transform(img).unsqueeze(0).to('cpu')
27 with torch.no_grad():
28     mask = model(inp)[-1].sigmoid().cpu().squeeze()
29 
30 # 对掩码进行后处理，去除小的连通区域
31 mask_np = mask.numpy()
32 mask_np[mask_np > 0.5] = 1
33 mask_np[mask_np <= 0.5] = 0
34 
35 # 形态学膨胀操作，连接断开的部分进行多次膨胀，目的是减少“断连”
36 selem = np.ones((2, 2))
37 num_iterations = 2
38 dilated_mask = mask_np
39 for _ in range(num_iterations):
40     dilated_mask = binary_dilation(dilated_mask, selem)
41 
42 # 对膨胀后的掩码进行高斯模糊，达到平滑边缘的目的，同时将模糊后的掩码进行二值化
43 blurred_mask = gaussian_filter(dilated_mask.astype(float), sigma=1.5)
44 blurred_mask[blurred_mask > 0.5] = 1
45 blurred_mask[blurred_mask <= 0.5] = 0
46 
47 # 恢复掩码的原始大小
48 mask_img = Image.fromarray((blurred_mask * 255).astype(np.uint8))
49 mask_img = mask_img.resize(img.size, Image.Resampling.NEAREST)
50 
51 # 将掩码应用到原始图像上，背景0完全透明化，从而达到抽离前景的目的
52 img = img.convert("RGBA")
53 datas = img.getdata()
54 newData = []
55 for item in zip(datas, mask_img.getdata()):
56     pixel = item[0]
57     alpha = item[1]
58     if alpha == 0:
59         # 背景颜色设置为#aa22ff，不透明
60         newData.append((0xAA, 0x22, 0xFF, 255))
61     else:
62         # 保留前景颜色，不透明
63         newData.append((pixel[0], pixel[1], pixel[2], 255))    
64 img.putdata(newData)
65 
66 # 存盘
67 img.save('input1_cpu_done.png')

    在终端命令行（确认已激活虚拟环境前面带rmbg_env）执行命令：

    python run_cpu.py

    分别对三张图片处理结果做对比如下（CPU版本）：

     

GPU 模式部署

注：最低 GPU 配置建议

  • 显存（VRAM）：​至少 6GB。

    虽然某些任务可能在 4GB 显存的显卡上运行，但为了确保模型加载和推理的稳定性，建议使用具有 6GB 或更高显存的显卡。​

  • CUDA 支持：​显卡需支持 CUDA，并与所安装的 CUDA Toolkit 版本兼容。​

  • 推荐显卡型号：

    • NVIDIA GTX 1660（6GB）

    • NVIDIA RTX 2060（6GB）

    • NVIDIA RTX 3060（12GB）

    • NVIDIA RTX 4060（8GB）

    这些显卡在性能和价格之间提供了良好的平衡，适合中等规模的图像处理任务。

  1. 首先确保已安装好CUDA Toolkit（文末处有提供），版本号确认请参照前面：环境准备，安装完成后打开命令行通过命令：

    nvcc --version

    进行查看，正常情况应显示与我的类似：​若显示nvcc是无法识别的命令，则说明没有安装CUDA Toolkits，或者安装后环境变量未生效，需要重启电脑。如果嫌重启电脑麻烦，可以安装 Chocolatey 提供的脚本，使用 refreshenv 命令立即刷新环境变量，免去重启。打开记事本，粘贴以下代码：

    @"%SystemRoot%\System32\WindowsPowerShell\v1.0\powershell.exe" -NoProfile -InputFormat None -ExecutionPolicy Bypass -Command "[System.Net.ServicePointManager]::SecurityProtocol = 3072; iex ((New-Object System.Net.WebClient).DownloadString('https://community.chocolatey.org/install.ps1'))" && SET "PATH=%PATH%;%ALLUSERSPROFILE%\chocolatey\bin"

refreshenv

    另存为 “刷新系统环境变量.bat” 批处理文件，通过右键 “以管理员身份运行”，则可。

  2. 在VSCode中新建并激活虚拟环境rmbg_evn_cuda：

    python -m venv rmbg_env_cuda
.\rmbg_env_cuda\Scripts\activate

  3. 新建python代码文run_cuda.py，并加入以下代码：

     1 from transformers import AutoModelForImageSegmentation
 2 import torch
 3 from PIL import Image
 4 import sys
 5 from torchvision import transforms
 6 import torch.nn.functional as F
 7 import torchvision.transforms.functional as FG
 8 import numpy as np
 9 from scipy.ndimage import binary_erosion
10 
11 if not torch.cuda.is_available():
12     sys.exit("CUDA Toolkit 未正确安装、配置")
13 
14 device = torch.device('cuda')
15 
16 # 加载模型
17 model = AutoModelForImageSegmentation.from_pretrained(
18     'briaai/RMBG-2.0', 
19     trust_remote_code=True
20 ).to(device).eval()
21 
22 # 图像预处理
23 transform = transforms.Compose([
24     transforms.Resize((1024, 1024)),
25     transforms.ToTensor(),
26     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
27 ])
28 
29 with torch.no_grad():
30     img = Image.open('input1.jpg').convert('RGB')
31     inp = transform(img).unsqueeze(0).to(device)  
32     
33     # 模型预测
34     mask = model(inp)[-1].sigmoid().squeeze(1)   
35     mask = torch.where(mask > 0.5, 1.0, 0.0)     
36 
37     # 形态学处理，统一4D格式
38     dilation_kernel = torch.ones(3,3, device=device).view(1,1,3,3)
39     erosion_kernel = torch.ones(2,2, device=device).view(1,1,2,2)
40     
41     # 维度修正：保持4D输入
42     mask_4d = mask.unsqueeze(1)  
43     for _ in range(1):
44         mask_4d = F.conv2d(mask_4d, dilation_kernel, padding=1).clamp(0,1)
45         mask_4d = F.conv2d(mask_4d, erosion_kernel, padding=1).clamp(0,1)
46     mask = mask_4d.squeeze(1)  
47 
48     # 高斯模糊
49     blurred_mask = FG.gaussian_blur(
50         mask.unsqueeze(1),  
51         kernel_size=9, 
52         sigma=1.5
53     ).squeeze(1)  
54 
55     # 插值与锐化
56     mask_img = F.interpolate(
57         blurred_mask.unsqueeze(1),  
58         size=(img.height, img.width), 
59         mode='bicubic'
60     ).squeeze(1).squeeze(0)  
61 
62     # 二值化
63     mask_img = torch.where(mask_img > 0.5, 1.0, 0.0).cpu().numpy()
64 
65 # 边缘补偿
66 alpha_channel = (mask_img * 255).astype(np.uint8)
67 alpha_channel = binary_erosion(alpha_channel, structure=np.ones((2,2))).astype(np.uint8) * 255
68 
69 # 合成背景，PIL格式兼容性修正
70 img.putalpha(Image.fromarray(alpha_channel))
71 target_bg_color = (170, 34, 255)
72 background = Image.new('RGB', img.size, target_bg_color)
73 background.paste(img, (0, 0), img.split()[-1])  
74 background.save('input1_cuda_done.png')

    在终端命令行（确认已激活虚拟环境前面带rmbg_env_cuda）执行命令：

    python run_cuda.py

    再来对三张图片处理结果对比下效果（CUDA版本）： 

总结与建议

  • 模型与依赖：务必保持 torch、transformers 版本与 CUDA 驱动匹配，避免运行时报错。

  • 调优建议：在 GPU 模式下，可通过 torch.set_float32_matmul_precision('high') 等方式微调精度与速度权衡。

huggingface的官方模型源码仓库需特殊方式访问，此外还需注册用户登录，并承诺不用于商业用途，才允许克隆及下载。虽然国内有镜像站，但经尝试发现也是不能正常登录的。若嫌麻烦，也可以从下面连接获取，本文仅从技术学习交流角度探讨纯官方模型部署及使用细节，不涉及商业用途，若需要商用请通过官方渠道获取授权，以避免潜在的版权纠纷，谢谢！祝你本地部署与二次开发顺利！

代码及工具包

RMBG-2.0模型  提取码：GCca

CUDA_Toolkit离线包 提取码：7N6j

​
posted @ 2025-04-29 18:09  智造行者  阅读(836)  评论(0)    收藏  举报