Pytorch 實現(xiàn)雙邊濾波

前幾天研究了傳統(tǒng)的美顏算法，了解到雙邊濾波（bilateral filtering）。在看懂原理后，為加深理解，抽時間用 pytorch 重新造了個輪子。雖然效率肯定比不上 opencv ，但當(dāng)個小練習(xí)也不錯。為了方便復(fù)習(xí)以及幫助初學(xué)者，在此記錄。

高斯濾波

高斯核函數(shù)

圖像領(lǐng)域的高斯濾波器是個二維的矩陣。矩陣中每個元素的值與它與矩陣中心的距離有關(guān)，計算公式就是二維高斯函數(shù)的公式：

G(u,v)=\frac{1}{2\pi\sigma^2}\exp(-\frac{u^2+v^2}{2\sigma^2})\tag{1}

為了讓卷積前后的圖像亮度保持不變，需要對 (1) 計算的矩陣歸一化（除以矩陣所有元素的和），因此 (1) 中 exp 之前的系數(shù)部分可以省略。生成高斯濾波器的代碼如下：

@torch.no_grad()
def getGaussianKernel(ksize, sigma=0):
    if sigma <= 0:
        # 根據(jù) kernelsize 計算默認(rèn)的 sigma，和 opencv 保持一致
        sigma = 0.3 * ((ksize - 1) * 0.5 - 1) + 0.8 
    center = ksize // 2
    xs = (np.arange(ksize, dtype=np.float32) - center) # 元素與矩陣中心的橫向距離
    kernel1d = np.exp(-(xs ** 2) / (2 * sigma ** 2)) # 計算一維卷積核
    # 根據(jù)指數(shù)函數(shù)性質(zhì)，利用矩陣乘法快速計算二維卷積核
    kernel = kernel1d[..., None] @ kernel1d[None, ...] 
    kernel = torch.from_numpy(kernel)
    kernel = kernel / kernel.sum() # 歸一化
    return kernel

高斯濾波器

pytorch 自帶的 conv2d 可以很方便地對圖像施加高斯濾波，代碼如下：

def GaussianBlur(batch_img, ksize, sigma=None):
    kernel = getGaussianKernel(ksize, sigma) # 生成權(quán)重
    B, C, H, W = batch_img.shape # C：圖像通道數(shù)，group convolution 要用到
    # 生成 group convolution 的卷積核
    kernel = kernel.view(1, 1, ksize, ksize).repeat(C, 1, 1, 1)
    pad = (ksize - 1) // 2 # 保持卷積前后圖像尺寸不變
    # mode=relfect 更適合計算邊緣像素的權(quán)重
    batch_img_pad = F.pad(batch_img, pad=[pad, pad, pad, pad], mode='reflect')
    weighted_pix = F.conv2d(batch_img_pad, weight=kernel, bias=None, 
                           stride=1, padding=0, groups=C)
    return weighted_pix

關(guān)于 group convolution，如果不熟悉可以看我這篇回答：什么是「Grouped Convolution」？

雙邊濾波

高斯濾波器的權(quán)重完全由距離決定。在大塊顏色差不多、偶有噪點的區(qū)域，它可以把顏色平均化，從而過濾掉噪點。但是在顏色變化劇烈的邊緣區(qū)域，它還是一視同仁地把所有像素做加權(quán)平均，這讓本應(yīng)該清晰銳利的邊緣也變得模糊不清了，這就造成了如下圖所示的效果，在做人像美顏時是不希望看到的。

高斯濾波

雙邊濾波的權(quán)重公式也基于高斯函數(shù)。但和高斯濾波的區(qū)別是，決定卷積核權(quán)重的，不單純是像素之間的空間距離，還包括像素之間的亮度差異。以卷積核中心為坐標(biāo)原點，該處像素值為I(0,0)。那么，坐標(biāo)為 (u, v) 處的像素，對應(yīng)的權(quán)重為：
G(u, v)=\exp(-\frac{u^2+v^2}{2\sigma_d^2}-\frac{\|I(u,v)-I(0,0)\|^2}{2\sigma_r^2})\tag{2}

(2) 中 exp 的第一個指數(shù)項和高斯核函數(shù)相同，與像素的空間距離有關(guān)；第二個指數(shù)項則是像素值距離的函數(shù)。以e為底對這兩項做指數(shù)運算，再相乘即得到了公式 (2)。根據(jù)公式 (2) 計算的卷積核有如下性質(zhì):

  • 距離中心像素越遠(yuǎn)的像素，其權(quán)重就越小
  • 亮度和中心像素亮度差異越大的像素，其權(quán)重就越小

第一條比較好理解，第二條的含義是，像素只和與自己相似的像素求平均，不和與自己差別太大的像素求平均。例如上圖中，人臉部分的像素和背景部分的像素差異過于顯著，這將導(dǎo)致在對邊緣區(qū)域做卷積時，背景側(cè)像素對人臉側(cè)的像素的權(quán)重極小，反之同理。這就達到了保持邊緣 (edge-preserving) 的特性。

代碼實現(xiàn)

由于 (2) 中卷積核的權(quán)重不僅僅依賴于空間距離，還依賴于像素的亮度，因此卷積核的權(quán)重是不固定的，不能簡單地利用 pytorch 的 conv2d 來實現(xiàn)。pytorch 的 tensor 自帶了一個 unfold 方法，正好可以用在這里。
unfold 的作用是把圖像拆分成 patch，每個patch 為卷積核覆蓋的像素。下面舉個小例子：

import torch
x = torch.arange(12).view(3, 4)
x
Out[4]: 
tensor([[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11]])

# 沿著行，以步長 1 拆分 x，每個 patch 為 2 行，列保持不變，
y = x.unfold(dimension=0, size=2, step=1) 
y.shape
Out[6]: torch.Size([2, 4, 2])
y[0]
Out[7]: 
tensor([[0, 4],
        [1, 5],
        [2, 6],
        [3, 7]])
y[1]
Out[8]: 
tensor([[ 4,  8],
        [ 5,  9],
        [ 6, 10],
        [ 7, 11]])

# 直接對 y 的第二個維度拆分，例如拆分成 3 列，步長仍為 1
z = y.unfold(dimension=1, size=3, step=1)
z.shape
Out[10]: torch.Size([2, 2, 2, 3])

# 觀察 z[0, 0]，發(fā)現(xiàn)正是 x 左上角的六個元素
z[0,0]
Out[11]: 
tensor([[0, 1, 2],
        [4, 5, 6]])

# z[0, 1] 也同樣符合預(yù)期
z[0,1]
Out[12]: 
tensor([[1, 2, 3],
        [5, 6, 7]])

實現(xiàn)的思路是：把原始圖像 unfold 成一個個的 patch，對每個 patch 計算權(quán)重以及加權(quán)平均。代碼如下：

def bilateralFilter(batch_img, ksize, sigmaColor=None, sigmaSpace=None):
    device = batch_img.device
    if sigmaSpace is None:
        sigmaSpace = 0.15 * ksize + 0.35  # 0.3 * ((ksize - 1) * 0.5 - 1) + 0.8
    if sigmaColor is None:
        sigmaColor = sigmaSpace
    
    pad = (ksize - 1) // 2
    batch_img_pad = F.pad(batch_img, pad=[pad, pad, pad, pad], mode='reflect')
    
    # batch_img 的維度為 BxcxHxW, 因此要沿著第 二、三維度 unfold
    # patches.shape:  B x C x H x W x ksize x ksize
    patches = batch_img_pad.unfold(2, ksize, 1).unfold(3, ksize, 1)
    patch_dim = patches.dim() # 6 
    # 求出像素亮度差
    diff_color = patches - batch_img.unsqueeze(-1).unsqueeze(-1)
    # 根據(jù)像素亮度差，計算權(quán)重矩陣
    weights_color = torch.exp(-(diff_color ** 2) / (2 * sigmaColor ** 2))
    # 歸一化權(quán)重矩陣
    weights_color = weights_color / weights_color.sum(dim=(-1, -2), keepdim=True)
    
    # 獲取 gaussian kernel 并將其復(fù)制成和 weight_color 形狀相同的 tensor
    weights_space = getGaussianKernel(ksize, sigmaSpace).to(device)
    weights_space_dim = (patch_dim - 2) * (1,) + (ksize, ksize)
    weights_space = weights_space.view(*weights_space_dim).expand_as(weights_color)
    
    # 兩個權(quán)重矩陣相乘得到總的權(quán)重矩陣
    weights = weights_space * weights_color
    # 總權(quán)重矩陣的歸一化參數(shù)
    weights_sum = weights.sum(dim=(-1, -2))
    # 加權(quán)平均
    weighted_pix = (weights * patches).sum(dim=(-1, -2)) / weights_sum
    return weighted_pix

最終結(jié)果為下圖，雀斑都沒了！同時人臉的輪廓和五官的細(xì)節(jié)依然被很好地保留下來：


輸入圖片尺寸為 256 x 256，ksize=15，sigmaColor=0.15，sigmaSpace=5 。

需要注意的是，由于 bilateral filter 的權(quán)重和像素值相關(guān)，因此設(shè)置 sigmaColor 時要注意輸入圖像的像素范圍，看清楚到底是 0-1 還是 0-255（上圖像素范圍為 0-1）。

總結(jié)

本文介紹了雙邊濾波的基本原理，并附帶了 pytorch 的實現(xiàn)。雖然不如 opencv 快，但優(yōu)點是 backward trackable ，適合包裝為模塊加入網(wǎng)絡(luò)中。利用 unfold 實現(xiàn)的缺點是很占內(nèi)存/顯存，kernelsize 越大，unfold 出來的冗余數(shù)據(jù)就越多，如果有大神知道更高效的實現(xiàn)方式，還望不吝賜教。

后記

我發(fā)現(xiàn)網(wǎng)上搜到的很多磨皮祛斑的算法，主要的目標(biāo)是設(shè)計一個高通濾波器，從而得到一個基于像素亮度的 mask，亮的地方權(quán)重大（對應(yīng)皮膚區(qū)域），暗的地方權(quán)重?。▽?yīng)雀斑、噪點區(qū)域）。將原圖 I 和 模糊化的圖 I_blur（各種模糊化方式都可以，目標(biāo)是把較暗的斑點模糊掉）利用 mask 融合：I_mask+I_blur(1-mask) 。這種方法既保留了原圖的細(xì)節(jié)，又能模糊掉斑點，不過在不同圖片上應(yīng)用時，仍然免不了調(diào)整一些超參數(shù)，而真有調(diào)參的功夫，直接調(diào)一下雙邊濾波的幾個參數(shù)，最后得到的效果未必比那些復(fù)雜的算法差。

最后編輯于 ：
?著作權(quán)歸作者所有,轉(zhuǎn)載或內(nèi)容合作請聯(lián)系作者
【社區(qū)內(nèi)容提示】社區(qū)部分內(nèi)容疑似由AI輔助生成，瀏覽時請結(jié)合常識與多方信息審慎甄別。
平臺聲明：文章內(nèi)容（如有圖片或視頻亦包括在內(nèi)）由作者上傳并發(fā)布，文章內(nèi)容僅代表作者本人觀點，簡書系信息發(fā)布平臺，僅提供信息存儲服務(wù)。

友情鏈接更多精彩內(nèi)容