Py学习  »  aigc

又一款AIGC工具火了:一键去除背景

机器学习算法工程师 • 2 月前 • 56 次点击  

点蓝色字关注“机器学习算法工程师

设为星标,干货直达!

近期,一家AIGC公司BRIA开源了一个出圈的模型:RMBG-1.4,它可以实现高质量地一键去除图片中的背景。下面是一些具体的例子,可以看到这个模型可以实现非常精细的“抠图”。


RMBG v1.4模型是基于22年所提出的IS-Net,但BRIA采用了私有数据集和特定的训练策略进行优化,这些改进显着提高了模型在不同图像处理场景中的准确性和有效性。

具体来说,这个模型使用超过 12,000 张高质量、高分辨率、手动标记(像素精度)、完全许可的图像进行训练。为了让模型有足够的泛化性,训练数据集包含各种类别的图片,下面是具体的数据类别分布:

从效果上看,RMBG v1.4模型完全可以媲美一些商业化产品:

目前RMBG v1.4模型已经在huggingface上开源:https://huggingface.co/briaai/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
from huggingface_hub import hf_hub_download

model_path = hf_hub_download("briaai/RMBG-1.4", 'model.pth')
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.load_state_dict(torch.load(model_path, map_location=device))
net.to(device)
net.eval()

# 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")

如果想快速体验,你也可以直接使用huggingface上搭建的demo:https://huggingface.co/spaces/briaai/BRIA-RMBG-1.4。用一张神仙姐姐的照片测试,效果还是刚刚的:

和StabilityAI的在线工具https://clipdrop.co/remove-background的效果不相上下:


用RMBG v1.4模型来抠图,然后搭配stable diffusion inpainting就可以实现一键换背景了!不过目前RMBG v1.4模型虽然是开源免费的,但不能用于商业用途,如果要商用,就需要征得BRIA公司的同意。


值得一提的,这个模型也已经有人快速实现了一个ComfyUI节点:https://github.com/ZHO-ZHO-ZHO/ComfyUI-BRIA_AI-RMBG,感兴趣的也可以去尝试。


推荐阅读

使用PyTorch 2.0加速Transformer:训练推理均拿下!

硬核解读Stable Diffusion(系列三)

硬核解读Stable Diffusion(系列二)

硬核解读Stable Diffusion(系列一)

带你入门扩散模型:DDPM


机器学习算法工程师


                                    一个用心的公众号


Python社区是高质量的Python/Django开发社区
本文地址:http://www.python88.com/topic/167087
 
56 次点击