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去除图片背景内容请搜索代码网以前的文章或继续浏览下面的相关文章希望大家以后多多支持代码网!
发表评论