当前位置: 代码网 > it编程>前端脚本>Python > Python去除图片背景的两种方式介绍

Python去除图片背景的两种方式介绍

2025年06月18日 Python 我要评论
1.使用rembg去除图像背景rembg 是一个开源的 python 库,专门用于去除图像背景,它利用深度神经网络能够准确地识别并去除图像背景,使用户无需进行复杂的手动编辑,只需几行代码即可获得专业效

1.使用rembg去除图像背景

rembg 是一个开源的 python 库,专门用于去除图像背景,它利用深度神经网络能够准确地识别并去除图像背景,使用户无需进行复杂的手动编辑,只需几行代码即可获得专业效果。rembg 基于 u2-net 架构,有多种架构修改和经过测试的方法以提供最佳结果,还提供对 gpu 安装的访问以实现更快的处理。

安装依赖:

pip install rembg

示例代码:

from rembg import remove
from pil import image
import numpy as np
import cv2
 
def remove_background(source_image_path):
    # 打开要处理的图片
    input_image = image.open(source_image_path)
 
    # 使用 rembg 去除背景
    output_image = remove(input_image, 
        alpha_matting=false
    )
 
    # 将 pil.image 转换为 opencv 格式
    output_image_array = np.array(output_image)
 
    # 如果是 rgba 图像,转换为 bgra(opencv 使用 bgra 格式)
    if output_image_array.shape[2] == 4:
        output_image_bgra = cv2.cvtcolor(output_image_array, cv2.color_rgba2bgra)
    else:
        output_image_bgra = cv2.cvtcolor(output_image_array, cv2.color_rgb2bgr)
    return output_image_bgra

2.使用u2-net去除图像背景

u2-net 是一种用于显著目标检测的深度学习模型,在 cvpr2020 开源后备受关注。其具有两层嵌套的 u 型结构,底层是带有新颖的 residual u-block(rsu)模块,可在不降低特征图分辨率的情况下提取多尺度特征;顶层则类似于 u-net 结构,每个阶段都由 rsu 填充,这种结构使网络能够在不显著增加内存和计算成本的情况下,加深网络并获得高分辨率特征图,从而有效捕捉显著目标的多尺度信息。

示例代码:

import torch
import os
import cv2
import numpy as np
from torchvision import transforms
from pil import image
import torch.nn as nn
from torch.nn import batchnorm2d
from u2net import u2net
from base_util import get_images_in_dir
 
# 定义u-2-net模型结构
class rebnconv(nn.module):
    def __init__(self, in_ch=3, out_ch=3, dirate=1):
        super(rebnconv, self).__init__()
        self.conv_s1 = nn.conv2d(in_ch, out_ch, 3, padding=1 * dirate, dilation=1 * dirate)
        self.bn_s1 = batchnorm2d(out_ch)
        self.relu_s1 = nn.relu(inplace=true)
 
    def forward(self, x):
        hx = x
        xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
        return xout
 
# 定义u-2-net模型的其他组件...
def normpred(d):
    ma = torch.max(d)
    mi = torch.min(d)
    dn = (d - mi) / (ma - mi)
    return dn
 
def remove_background(input_path, output_path, model_path='models/u2net.pth'):
    # 加载模型
    net = u2net(3,1)
    net.load_state_dict(torch.load(model_path))
    if torch.cuda.is_available():
        net.cuda()
    net.eval()
    
    # 读取图像
    img = image.open(input_path).convert('rgb')
 
    # 图像预处理
    transforms_list = []
    transforms_list.append(transforms.resize((320, 320)))
    transforms_list.append(transforms.totensor())
    transforms_list.append(transforms.normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]))
    transform = transforms.compose(transforms_list)
    
    image = transform(img)
    image = image.unsqueeze(0)
    
    # 模型推理
    if torch.cuda.is_available():
        image = image.cuda()
    
    with torch.no_grad():
        d1, d2, d3, d4, d5, d6, d7 = net(image)
    
    # 获取预测结果
    pred = d1[:, 0, :, :]
    pred = normpred(pred)
    
    # 将预测结果转换为掩码
    predict = pred.squeeze()
    predict_np = predict.cpu().data.numpy()
    mask = image.fromarray((predict_np * 255).astype(np.uint8))
    mask = mask.resize(img.size, resample=image.bilinear)
    
    # 应用阈值处理掩码
    mask_threshold = 0.5  # 阈值可以根据实际情况调整
    mask_array = np.array(mask)
    mask_array = np.where(mask_array > mask_threshold * 255, 255, 0).astype(np.uint8)
 
    kernel = np.ones((3, 3), np.uint8)  # 定义膨胀核
    mask_array = cv2.dilate(mask_array, kernel, iterations=1)
 
    # 应用掩码到原始图像
    img_array = np.array(img)
    
    # 创建透明背景的图像
    transparent_img = np.zeros((img_array.shape[0], img_array.shape[1], 4), dtype=np.uint8)
    transparent_img[:, :, :3] = img_array
    transparent_img[:, :, 3] = mask_array
 
    # 保存结果
    image.fromarray(transparent_img).save(output_path)
    print(f"已保存处理后的图像到: {output_path}")

到此这篇关于python去除图片背景的两种方式介绍的文章就介绍到这了,更多相关python去除图片背景内容请搜索代码网以前的文章或继续浏览下面的相关文章希望大家以后多多支持代码网!

(0)

相关文章:

版权声明:本文内容由互联网用户贡献,该文观点仅代表作者本人。本站仅提供信息存储服务,不拥有所有权,不承担相关法律责任。 如发现本站有涉嫌抄袭侵权/违法违规的内容, 请发送邮件至 2386932994@qq.com 举报,一经查实将立刻删除。

发表评论

验证码:
Copyright © 2017-2025  代码网 保留所有权利. 粤ICP备2024248653号
站长QQ:2386932994 | 联系邮箱:2386932994@qq.com