最近捣鼓了一下Python下的开源库rembg,一个非常好用的背景移除工具,支持各种各样的物体用于移除背景。并且支持多种模型,还支持微调训练。具体可以看官网:danielgatis/rembg: Rembg is a tool to remove images background。本文使用的环境是ubuntu 22.04,CUDA 12.2。

环境准备

首先需要知道自己的硬件环境(CUDA)和ONNX运行时的对应关系,rembg是基于ONNX实现的。

官方文档里面有具体的版本说明:NVIDIA - CUDA | onnxruntime,我这里为了方便,引用了官方的表格:

ONNX RuntimeCUDAcuDNNNotes
1.20.x12.x9.xAvaiable in PyPI. Compatible with PyTorch >= 2.4.0 for CUDA 12.x.
1.19.x12.x9.xAvaiable in PyPI. Compatible with PyTorch >= 2.4.0 for CUDA 12.x.
1.18.112.x9.xcuDNN 9 is required. No Java package.
1.18.012.x8.xJava package is added.
1.17.x12.x8.xOnly C++/C# Nuget and Python packages are released. No Java package.

可以看到CUDA 12.2是对应的ONNX Runtime 1.19.2,从pip里也能查到版本:

onnxruntime-gpu · PyPI

记得加上-gpu并指定版本,否则安装出来是cpu版本的,没有图形加速。

pip install onnxruntime-gpu

ONNX Runtime | Getting-started

image-20250101171735231.webp

缺少运行环境

1.安装cuDNN,可以直接从官网下载,并配置环境。

cuDNN Archive | NVIDIA Developer

或者使用pip环境里面的cudnn,envs/环境名称/lib/python3.12/site-packages/nvidia/cudnn/lib

缺少时会提示 error: libcudnn.so.9: cannot open shared object file: No such file or directory

2.安装tensorrt

运行时提示FAIL : Failed to load library libonnxruntime_providers_tensorrt.so with error: libnvinfer.so.10: cannot open shared object file: No such file or directory

并使用pip环境里面的tensorrt,envs/环境名称/lib/python3.12/site-packages/tensorrt_libs

export LD_LIBRARY_PATH=/home/cyqsd/miniconda3/envs/环境名称/lib/python3.12/site-packages/tensorrt_libs:$LD_LIBRARY_PATH

下载BiRefNet-portrait模型

onnx-community/BiRefNet-portrait-epoch_150 at main

# 如果模型自动下载失败,则可以手动拷贝模型文件到以下目录
/home/cyqsd/.u2net

Release中也提供了模型的下载地址:

Release BiRefNet-v1 · ZhengPeng7/BiRefNet

image-20250101171654497.webp

  • u2net : 一个适用于通用场景的预训练模型。
  • u2netp : u2net 模型的轻量级版本。
  • u2net_human_seg : 一个用于人体分割的预训练模型。
  • u2net_cloth_seg : 一个从人像中解析衣物的预训练模型。衣物被分为三类:上身、下身和全身。
  • silueta : 与 u2net 相同,但大小减少到 43Mb。
  • isnet-general-use : 一个新的适用于通用场景的预训练模型。
  • isnet-anime : 一个针对动漫角色的高精度分割模型。
  • sam : 一个适用于任何场景的预训练模型。
  • birefnet-general : 一个适用于通用场景的预训练模型。
  • birefnet-general-lite : 一个适用于通用场景的轻量级预训练模型。
  • birefnet-portrait : 一个用于人像的预训练模型。
  • birefnet-dis : 一个用于二值图像分割(DIS)的预训练模型。
  • birefnet-hrsod : 一个用于高分辨率显著物体检测(HRSOD)的预训练模型。
  • birefnet-cod : 一个用于隐匿物体检测(COD)的预训练模型。
  • birefnet-massive : 一个基于大规模数据集训练的预训练模型

运行

使用官方文档中的示例即可:

输入输出:

from rembg import remove

input_path = 'input.png'
output_path = 'output.png'

with open(input_path, 'rb') as i:
    with open(output_path, 'wb') as o:
        input = i.read()
        output = remove(input)
        o.write(output)

遍历目录中的png文件并输出:

from pathlib import Path
from rembg import remove, new_session

session = new_session()

for file in Path('path/to/folder').glob('*.png'):
    input_path = str(file)
    output_path = str(file.parent / (file.stem + ".out.png"))

    with open(input_path, 'rb') as i:
        with open(output_path, 'wb') as o:
            input = i.read()
            output = remove(input, session=session)
            o.write(output)

new_session中传模型参数即可实现模型的指定。

# 加载自定义模型
model_path = "./40.onnx"
session = new_session(model_name='40', model_path=model_path)

需要额外做微调可以参考这篇:How to convert custom .pth Model to .onnx? · Issue #193 · danielgatis/rembg

文章目录