一键AI抠图，证件照换背景，可部署成自己的应用

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-libraries

In [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 ioclass 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].nameoutputs = self.session.run(None, {input_name: input_tensor})return outputsdef 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)  # 将图像数据类型转换为float32im /= 255.0  # 归一化到[0, 1]范围mean = [0.5, 0.5, 0.5]std = [1.0, 1.0, 1.0]im -= meanim /= stdreturn 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)  # 将数据类型转换回uint8result = cv2.resize(result, [im_size[1], im_size[0]], interpolation=cv2.INTER_LINEAR)  # 调整图像大小return resultdef 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格式，先转换为PNGif image.format != 'PNG':image = image.convert('RGBA')if out_size is None:out_size = image.sizeimageout_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).shapeout_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" + extoutput_image_path = path + new_filename # 证件照net = BriaRMBG_ONNX(f"/home/aistudio/rmbg/onnx/model.onnx" )# prepare inputmodel_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格式# inferenceresult = net(image)# post processresult_image = postprocess_image(result[0][0], orig_im_size)# save resultpil_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_path

5 调用函数进行抠图

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()