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

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

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