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

最新推荐文章于 2024-09-14 06:56:53 发布

军哥说AI

最新推荐文章于 2024-09-14 06:56:53 发布

阅读量708

点赞数 8

分类专栏： AI写算法文章标签：人工智能 python 迁移学习

本文链接：https://blog.csdn.net/class4715/article/details/139351154

版权

AI写算法专栏收录该内容

56 篇文章 31 订阅

订阅专栏

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