基于ONNX模型推理和图像处理技术实现证件照背景一键转换

结合 ONNX模型推理和图像处理技术，实现了从图像中去除背景并添加新背景的功能。通过预处理和后处理步骤，确保图像适应模型输入和输出要求，同时提供了与Gradio 界面交互的支持，使用户可以方便地选择背景颜色和图像大小，并提供可执行的exe文件，实现一键转换。

效果展示

一键背景变换——点击即可使用

图像预处理

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, :, :, :]

preprocess_image函数用于预处理输入图像，将图像调整为模型所需的尺寸和格式，并进行归一化处理。

背景添加

def add_background_to_image(input_image_path, output_image_path, background_color, out_size=None):
    image = Image.open(input_image_path)

    if image.format != 'PNG':
        image = image.convert('RGBA')
    
    if out_size is None:
        out_size = image.size

    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)

add_background_to_image函数用于给透明背景的人像图像添加任意颜色的背景。通过Pillow库将图像转换为RGBA格式，并创建一个新的背景图像。

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))
    print('input_image_path:',input_image_path)
    name = os.path.basename(input_image_path)
    path = os.path.dirname(input_image_path)
    filename, ext = os.path.splitext(name)
    ext = '.png'
    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")

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

    result = net(image)
    result_image = postprocess_image(result[0][0], orig_im_size)

    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, "PNG")

    print(background_color, out_size)
    add_background_to_image(out_path, output_image_path, background_color, out_size)
    return out_path, output_image_path

根据输入图像路径加载图像并处理图像尺寸，调用ONNX模型去除背景并后处理结果图像，最后添加自定义背景颜色并保存最终图像。

下载地址：点击