简单使用rembg并捣鼓ONNX Runtime
最近捣鼓了一下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 Runtime | CUDA | cuDNN | Notes |
---|---|---|---|
1.20.x | 12.x | 9.x | Avaiable in PyPI. Compatible with PyTorch >= 2.4.0 for CUDA 12.x. |
1.19.x | 12.x | 9.x | Avaiable in PyPI. Compatible with PyTorch >= 2.4.0 for CUDA 12.x. |
1.18.1 | 12.x | 9.x | cuDNN 9 is required. No Java package. |
1.18.0 | 12.x | 8.x | Java package is added. |
1.17.x | 12.x | 8.x | Only C++/C# Nuget and Python packages are released. No Java package. |
可以看到CUDA 12.2是对应的ONNX Runtime 1.19.2,从pip里也能查到版本:
记得加上-gpu
并指定版本,否则安装出来是cpu版本的,没有图形加速。
pip install onnxruntime-gpu
ONNX Runtime | Getting-started
缺少运行环境
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
- 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
本作品采用 知识共享署名-相同方式共享 4.0 国际许可协议 进行许可。