AI证件照抠图

*免费AI证件照抠图

1. 环境导入

mkdir /home/aistudio/external-libraries
pip install  onnxruntime scikit-image -i https://mirrors.aliyun.com/pypi/simple/ \
     -t /home/aistudio/external-libraries

2. main.py

# AI证件照抠图

# 1 设置环境变量 （每次重新进入此项目或者重启内核后，都要执行此代码）
# 在py文件里面运行环境
import sys

sys.path.append('./external-libraries')
# 在ipynb里面运行环境
import os

os.environ['PYTHONPATH'] = os.environ.get('PYTHONPATH', '') + ':./external-libraries'

# 2定义remove background的函数


import sys

sys.path.append('./external-libraries')

# 在ipynb里面运行环境
import os

os.environ['PYTHONPATH'] = os.environ.get('PYTHONPATH', '') + ':./external-libraries'

import onnxruntime as ort
import cv2
import numpy as np
from PIL import Image, ImageDraw
import re
from skimage import io
#pip install  onnxruntime -i https://mirrors.aliyun.com/pypi/simple/  -t /home/aistudio/external-libraries
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"./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

import cv2
import matplotlib.pyplot as plt



# 输入图片的路径
#input_img = './rmbg/photo/wj.png'
input_img='./jpgtopng/t.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)

3. 实现效果展示

# 展示图片

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

*原图

*抠图

*证件照

可以需要需要选择背景色