1 开发背景
AI抠图技术已经非常成熟，并且有效果非常好的开源模型。 日常中可以用于替换证件照背景
但是网上许多的证件照替换背景 竟然需要收费
鉴于此，便将目前最好的（SOTA）开源抠图模型 BRIA Background Removal v1.4 Model 进行部署
也可以直接体验部署好的应用
点击链接体验应用
2 项目内容
-
展示如何配置环境，这种配置环境的方法是所有项目通用的，配置一次就可以，做到一劳永逸
-
展示如何调用模型，进行背景的抠去和背景颜色的替换
-
最后展示了如何将该项目 部署成一个应用，这个应用比较简单，也是具有教学意义
3 配置环境
In [ ]
# 配置环境（第一次运行次项目的时候配置即可，之后无需安装）
# 安装包
!mkdir /home/aistudio/external-libraries
!pip install onnxruntime scikit-image -i https://mirrors.aliyun.com/pypi/simple/ \
-t /home/aistudio/external-librariesIn [1]
# 设置环境变量 （每次重新进入此项目或者重启内核后，都要执行此代码）
# 在py文件里面运行环境
import sys
sys.path.append('/home/aistudio/external-libraries')
# 在ipynb里面运行环境
import os
os.environ['PYTHONPATH'] = os.environ.get('PYTHONPATH', '') + ':/home/aistudio/external-libraries'4 定义remove background的函数
In [2]
# 在py文件里面运行环境
import sys
sys.path.append('/home/aistudio/external-libraries')
# 在ipynb里面运行环境
import os
os.environ['PYTHONPATH'] = os.environ.get('PYTHONPATH', '') + ':/home/aistudio/external-libraries'
import onnxruntime as ort
import cv2
import numpy as np
from PIL import Image, ImageDraw
import re
from skimage import io
class BriaRMBG_ONNX:
def __init__(self, model_path):
self.session = ort.InferenceSession(model_path)
def __call__(self, input_tensor):
input_name = self.session.get_inputs()[0].name
outputs = self.session.run(None, {input_name: input_tensor})
return outputs
def preprocess_image(im: np.ndarray, model_input_size: list) -> np.ndarray:
if len(im.shape) < 3:
im = cv2.cvtColor(im, cv2.COLOR_GRAY2BGR) # 将灰度图像转换为BGR格式
im = cv2.resize(im, model_input_size, interpolation=cv2.INTER_LINEAR) # 调整图像大小
im = im.astype(np.float32) # 将图像数据类型转换为float32
im /= 255.0 # 归一化到[0, 1]范围
mean = [0.5, 0.5, 0.5]
std = [1.0, 1.0, 1.0]
im -= mean
im /= std
return im[np.newaxis, :, :, :]
def postprocess_image(result: np.ndarray, im_size: list) -> np.ndarray:
result = result[0] # 移除batch维度
ma = np.max(result)
mi = np.min(result)
result = (result - mi) / (ma - mi) # 归一化到[0, 1]范围
result = (result * 255).astype(np.uint8) # 将数据类型转换回uint8
result = cv2.resize(result, [im_size[1], im_size[0]], interpolation=cv2.INTER_LINEAR) # 调整图像大小
return result
def add_background_to_image(input_image_path, output_image_path, background_color, out_size=None):
"""
给透明背景的PNG人像图像添加任意颜色的背景。
:param input_image_path: 输入图像的路径
:param output_image_path: 输出图像的路径
:param background_color: 背景颜色 (R, G, B)
:param size: 输出图像的大小 (width, height) 默认与输入图像相同
"""
# 打开输入图像
image = Image.open(input_image_path)
# 如果图像不是PNG格式，先转换为PNG
if image.format != 'PNG':
image = image.convert('RGBA')
if out_size is None:
out_size = image.size
image
out_image = Image.new('RGB', image.size, background_color)
out_image.paste(image, (0,0), image)
out_image.resize(out_size)
# 保存新的图像
out_image.save(output_image_path)
def rmbg(input_image_path, background_color, out_size_w, out_size_h, size_opt):
if size_opt == "保持原图大小":
shape = cv2.imread(input_image_path).shape
out_size = (int(shape[0]),int(shape[1]))
else:
out_size = (int(out_size_w), int(out_size_h))
match = re.search(r'^(.+/)([^.]+)(\..+)$', input_image_path) # 使用正则表达式找到图片名称和扩展名
path, filename, ext = match.groups() # 获取组：路径、文件名、扩展名
new_filename = filename + "_rmgb" + ext # 修改文件名
out_path = path + new_filename # 抠图
new_filename = filename + "_bg" + ext
output_image_path = path + new_filename # 证件照
net = BriaRMBG_ONNX(f"/home/aistudio/rmbg/onnx/model.onnx" )
# prepare input
model_input_size = [1024,1024]
orig_im = io.imread(input_image_path)
orig_im_size = orig_im.shape[0:2]
image = preprocess_image(orig_im, model_input_size)
image = np.transpose(image, (0, 3, 1, 2)) # ONNX通常需要CHW格式
# 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(input_image_path)
no_bg_image.paste(orig_image, mask=pil_im)
no_bg_image.save(out_path)
print(background_color, out_size)
add_background_to_image(out_path, output_image_path, background_color, out_size)
return out_path, output_image_path5 调用函数进行抠图
In [3]
import cv2
import matplotlib.pyplot as plt
# 输入图片的路径
input_img = '/home/aistudio/rmbg/photo/wj.png'
# 证件照的背景颜色
# color = "#FFFFFF" # 白色（用于护照、签证、身份证等）
color = "#438EDB" # 蓝色（用于毕业证、工作证等）
# color = "#FF0000" # 红色（用于一些特殊的证件照）
# 证件照的大小
width = 295
height = 413 # 一寸（295像素 x 413像素）
# 是否保持原图大小
# size_opt = "不保持原图大小"
size_opt = "保持原图大小" # 如果选了这个会保持输入图片的大小，忽略上面的 证件照的大小 参数
# color, width, height 这三个参数不影响抠图，只会影响证件照的结果
out_path, output_image_path = rmbg(input_img, color, width, height, size_opt)
print('抠图后的图片： ', out_path)
print('证件照： ', output_image_path)
#438EDB (1024, 1596) 抠图后的图片： /home/aistudio/rmbg/photo/wj_rmgb.png 证件照： /home/aistudio/rmbg/photo/wj_bg.png
6 结果展示
In [4]
# 展示图片
from PIL import Image
import matplotlib.pyplot as plt
# print('原图')
image_input = Image.open(input_img)
# print('抠图后的图片')
image_rmbg = Image.open(out_path)
# print('证件照')
image_bg = Image.open(output_image_path)
# 设定图片显示的大小
fig, axs = plt.subplots(3, 1, figsize=(5, 15))
# 在每个子图上显示一张图片
axs[0].imshow(image_input)
axs[0].axis('off') # 不显示坐标轴
axs[0].set_title(' original')
axs[1].imshow(image_rmbg)
axs[1].axis('off') # 不显示坐标轴
axs[1].set_title(' remove background')
axs[2].imshow(image_bg)
axs[2].axis('off') # 不显示坐标轴
axs[2].set_title('with background')
# 调整子图之间的间距
plt.tight_layout()
# 显示画布
plt.show()
<Figure size 500x1500 with 3 Axes>
7 部署自己的应用
应用的代码我已经写好了 /home/aistudio/untitled.gradio.py
直接点击右上角的“部署” 即可，步骤如下
分割线
