快捷導(dǎo)航

Python去除圖片背景的兩種方式介紹

更新時(shí)間：2025年06月18日 15:07:33 作者：培根芝士

這篇文章主要為大家詳細(xì)介紹了Python去除圖片背景的兩種方式,一個(gè)是使用rembg,一個(gè)是使用U2-NET,文中的示例代碼講解詳細(xì),感興趣的小伙伴可以了解下

1.使用rembg去除圖像背景

rembg 是一個(gè)開源的 Python 庫，專門用于去除圖像背景，它利用深度神經(jīng)網(wǎng)絡(luò)能夠準(zhǔn)確地識(shí)別并去除圖像背景，使用戶無需進(jìn)行復(fù)雜的手動(dòng)編輯，只需幾行代碼即可獲得專業(yè)效果。Rembg 基于 U2-Net 架構(gòu)，有多種架構(gòu)修改和經(jīng)過測(cè)試的方法以提供最佳結(jié)果，還提供對(duì) GPU 安裝的訪問以實(shí)現(xiàn)更快的處理。

安裝依賴：

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 轉(zhuǎn)換為 OpenCV 格式
    output_image_array = np.array(output_image)
 
    # 如果是 RGBA 圖像，轉(zhuǎn)換為 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 是一種用于顯著目標(biāo)檢測(cè)的深度學(xué)習(xí)模型，在 CVPR2020 開源后備受關(guān)注。其具有兩層嵌套的 U 型結(jié)構(gòu)，底層是帶有新穎的 ReSidual U-Block（RSU）模塊，可在不降低特征圖分辨率的情況下提取多尺度特征；頂層則類似于 U-Net 結(jié)構(gòu)，每個(gè)階段都由 RSU 填充，這種結(jié)構(gòu)使網(wǎng)絡(luò)能夠在不顯著增加內(nèi)存和計(jì)算成本的情況下，加深網(wǎng)絡(luò)并獲得高分辨率特征圖，從而有效捕捉顯著目標(biāo)的多尺度信息。

示例代碼：

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模型結(jié)構(gòu)
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')
 
    # 圖像預(yù)處理
    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)
    
    # 獲取預(yù)測(cè)結(jié)果
    pred = d1[:, 0, :, :]
    pred = normPRED(pred)
    
    # 將預(yù)測(cè)結(jié)果轉(zhuǎn)換為掩碼
    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)
    
    # 應(yīng)用閾值處理掩碼
    mask_threshold = 0.5  # 閾值可以根據(jù)實(shí)際情況調(diào)整
    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)
 
    # 應(yīng)用掩碼到原始圖像
    img_array = np.array(img)
    
    # 創(chuàng)建透明背景的圖像
    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
 
    # 保存結(jié)果
    Image.fromarray(transparent_img).save(output_path)
    print(f"已保存處理后的圖像到: {output_path}")

到此這篇關(guān)于Python去除圖片背景的兩種方式介紹的文章就介紹到這了,更多相關(guān)Python去除圖片背景內(nèi)容請(qǐng)搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家！

您可能感興趣的文章: