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

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

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

1.使用rembg去除圖像背景

rembg 是一個開源的 Python 庫，專門用于去除圖像背景，它利用深度神經網絡能夠準確地識別并去除圖像背景，使用戶無需進行復雜的手動編輯，只需幾行代碼即可獲得專業(yè)效果。Rembg 基于 U2-Net 架構，有多種架構修改和經過測試的方法以提供最佳結果，還提供對 GPU 安裝的訪問以實現(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 轉換為 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)
    
    # 創(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
 
    # 保存結果
    Image.fromarray(transparent_img).save(output_path)
    print(f"已保存處理后的圖像到: {output_path}")

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

您可能感興趣的文章: