社区所有版块导航
Python
python开源   Django   Python   DjangoApp   pycharm  
DATA
docker   Elasticsearch  
aigc
aigc   chatgpt  
WEB开发
linux   MongoDB   Redis   DATABASE   NGINX   其他Web框架   web工具   zookeeper   tornado   NoSql   Bootstrap   js   peewee   Git   bottle   IE   MQ   Jquery  
机器学习
机器学习算法  
Python88.com
反馈   公告   社区推广  
产品
短视频  
印度
印度  
Py学习  »  aigc

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

机器学习算法工程师 • 1 年前 • 264 次点击  

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

设为星标,干货直达!

近期,一家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
 
264 次点击