代码收藏家技术教程 2024-08-07

采用Python的rembg库去除图片背景

近期，在工作中提出一个新的需求，需要将用户上传的人脸照片的背景去除，也就是通俗来讲的抠图。调研常用的一些工具和框架，发现了Python的rembg库。该库使用了深度学习模型，尤其是U-Net结构的卷积神经网络（CNN网络），能够高效、准确的提取图片的前景物体，同时去除图片背景，最终生成带有透明通道（Alpha通道）的png图片。

Rembg工具库使用U-Net结构的卷积神经网络（CNN），这是一种在图像分割效果很好的模型。U-Net模型的特点是结合使用上采样和下采样的分割网络，既能捕捉到图片的全局信息，又不忽略图片局部细节，使其具备了强大的图片前景和背景的分割能力。

Rembg库的优势

简单易用，使用Python api接口，数行代码即可实现图片背景去除工作。
处理速度快，使用高效的深度学习模型，可在短时间内处理大量图片，即便是处理高分辨率图片也拥有较快的处理速度。
精度高，即使图片有复杂的光线条件和背景，对毛发、‌透明物体或复杂边缘的处理也能保持较高的精度。
智能识别，采用先进的深度学习算法，‌能够准确识别图像中的主体与背景，‌无需手动调整，‌一键即可完成抠图。
开源免费：‌遵循MIT许可证，‌允许自由使用、‌修改和分享。‌

项目地址

源码：https://github.com/danielgatis/rembg

安装和使用

安装

使用pip安装rembg库

pip install rembg

命令行

首先查看命令行帮助文档

PS D:\pythonCode\pythonProject> rembg –help
Usage: rembg.exe [OPTIONS] COMMAND [ARGS]…

Options:
–version Show the version and exit.
–help Show this message and exit.

Commands:
b for a byte stream as input
d download all models
i for a file as input
p for a folder as input
s for a http server

单图片处理

本人使用E:\img_test\input\目录作为图片输入地址，使用E:\img_test\output作为图片输出地址

rembg i E:\img_test\input\1.jpg E:\img_test\output\1.png

效果如下：

图片目录处理

rembg p E:\img_test\input\ E:\img_test\output\

效果如下：

Web Server运行

rembg s --host 0.0.0.0 --port 7000 --log_level info

    使用rembg s命令用于启动http服务器。启动成功在浏览器打开http://localhost:7000，如下图所示，即可上传图片执行去除背景操作。

python中使用

单图片处理代码示例

from rembg import remove

if __name__=='__main__':

    # 待处理的图片路径
    input_path = 'E:\\img_test\\input\\1.jpg'
    # 处理后存储的图片路径
    output_path = 'E:\\img_test\\output\\1.png'

    with open(input_path, 'rb') as i:
        with open(output_path, 'wb') as o:
            input = i.read()
            output = remove(input)
            o.write(output)

使用PIL库目录批处理文件代码示例

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/16 下午3:19
@Author  : Bill Fang
@File    : RemoveBg.py
@Desc    : 
"""
from PIL import Image
import os
from rembg import remove


input_img = "E:\\img_test\\input"
output_img = "E:\\img_test\\output"

if __name__=='__main__':

    imgs_path = getFileList(input_img)
    for img_path in imgs_path:
        img_name = os.path.basename(img_path)
        if img_name.endswith('.jpg') or img_name.endswith('.jpeg'):
            #解决jpg图片无Alpha通道
            input = Image.open(img_path).convert('RGB')
        else:
            input = Image.open(img_path)
        #rembg去除图片背景
        output = remove(input)
        # 将去除背景后的图片背景设置为白色
        background = Image.new('RGBA', output.size, (255, 255, 255, 255))
        background.paste(output, mask=output)
        out_path = os.path.join(output_img, img_name.replace('.jpg', '.png').replace('.jpeg','.png'))

        background.save(out_path)
        input.close()
        output.close()
        background.close()

#获取目录下全部图片文件
def getFileList(file_path):
    imgs_path = []
    for filename in os.listdir(file_path):
        img_path = os.path.join(file_path, filename).lower()
        imgs_path.append(img_path)

    return imgs_path

遇到问题

问题描述

初次执行代码时需要下载v2net.onnx模型文件，模型文件大小为176M，由于网络条件差，导致模型下载失败，运行报错如下：

解决方法

第一种解决方法

从https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx

地址下载模型文件，然后放入C:\Users\njfan\.u2net\u2net.onnx目录中。

第二种解决方法

使用rembg的自定义模型，代码如下：

import os
from rembg import remove,new_session

if __name__=='__main__': 
    # 加载自定义模型
    model_path = "E:\\img_test\\u2net.onnx"
    
    with open(input_path, 'rb') as i:
        with open(output_path, 'wb') as o:
            input = i.read()
            #注意model_name必须为u2net_custom
            session = new_session(model_name='u2net_custom', model_path=model_path)
            output_image = remove(input, session=session)
            o.write(output_image)

首先通过new_session()函数自定义模型，然后使用remove(input, session=session)函数调用自定义模型去除图片背景。