BRIA.AI开源最强AI一键抠图模型RMBG，超简上手体验

本文首发: AIWalker
欢迎关注AIWalker，近距离接触底层视觉与基础AI技术

近日，BRIA.AI团队于HuggingFace开源了一个基于ISNet背景移除模型RMBG-1.4，它可以有效对前景与背景进行分离。RMBG-1.4在精心构建的数据集上训练而来，该数据包含常规图像、电商、游戏以及广告内容，该方案达到了商业级性能，但仅限于非商业用途。关于所用到的训练数据：12000+高质量&高分辨率像素级精度手工标注。更详细的数据分布介绍请移步[RMBG-1.4]

著名的HuggingFace上已有该背景移除模型的体验Demo，见：https://huggingface.co/spaces/briaai/BRIA-RMBG-1.4，用户只需要上传图片即可体验。

当然，也有效果不那么好的，比如下面这张：

快速实战

代码下载

git clone https://huggingface.co/briaai/RMBG-1.4
cd RMBG-1.4/
pip install -r requirements.txt

代码调用示例

from skimage import io
import torch, os
from PIL import Image
from briarmbg import BriaRMBG
from utilities import preprocess_image, postprocess_image

im_path = f"{os.path.dirname(os.path.abspath(__file__))}/example_input.jpg"

net = BriaRMBG()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = BriaRMBG.from_pretrained("briaai/RMBG-1.4")
net.to(device)

# prepare input
model_input_size = [1024,1024]
orig_im = io.imread(im_path)
orig_im_size = orig_im.shape[0:2]
image = preprocess_image(orig_im, model_input_size).to(device)

# inference 
result=net(image)

# post process
result_image = postprocess_image(result[0][0], orig_im_size)

# save result
pil_im = Image.fromarray(result_image)
no_bg_image = Image.new("RGBA", pil_im.size, (0,0,0,0))
orig_image = Image.open(im_path)
no_bg_image.paste(orig_image, mask=pil_im)
no_bg_image.save("example_image_no_bg.png")