【他山之石】pytorch 实现双边滤波

“他山之石，可以攻玉”，站在巨人的肩膀才能看得更高，走得更远。在科研的道路上，更需借助东风才能更快前行。为此，我们特别搜集整理了一些实用的代码链接，数据集，软件，编程技巧等，开辟“他山之石”专栏，助你乘风破浪，一路奋勇向前，敬请关注。

作者：知乎—尹相楠

原文地址：https://zhuanlan.zhihu.com/p/310710051

前几天研究了传统的美颜算法，了解到双边滤波（bilateral filtering）。在看懂原理后，为加深理解，抽时间用 pytorch 重新造了个轮子。虽然效率肯定比不上 opencv ，但当个小练习也不错。为了方便复习以及帮助初学者，在此记录。

01

高斯滤波

1. 高斯核函数

图像领域的高斯滤波器是个二维的矩阵。矩阵中每个元素的值与它与矩阵中心的距离有关，计算公式就是二维高斯函数的公式：

为了让卷积前后的图像亮度保持不变，需要对 (1) 计算的矩阵归一化（除以矩阵所有元素的和），因此 (1) 中 exp 之前的系数部分可以省略。生成高斯滤波器的代码如下：

@torch.no_grad()
def getGaussianKernel(ksize, sigma=0):
    if sigma <= 0:
        # 根据 kernelsize 计算默认的 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)) # 计算一维卷积核
    # 根据指数函数性质，利用矩阵乘法快速计算二维卷积核
    kernel = kernel1d[..., None] @ kernel1d[None, ...] 
    kernel = torch.from_numpy(kernel)
    kernel = kernel / kernel.sum() # 归一化
    return kernel

2. 高斯滤波器

pytorch 自带的 conv2d 可以很方便地对图像施加高斯滤波，代码如下：

def GaussianBlur(batch_img, ksize, sigma=None):
    kernel = getGaussianKernel(ksize, sigma) # 生成权重
    B, C, H, W = batch_img.shape # C：图像通道数，group convolution 要用到
    # 生成 group convolution 的卷积核
    kernel = kernel.view(1, 1, ksize, ksize).repeat(C, 1, 1, 1)
    pad = (ksize - 1) // 2 # 保持卷积前后图像尺寸不变
    # mode=relfect 更适合计算边缘像素的权重
    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

关于 group convolution，如果不熟悉可以看我这篇回答：

什么是「Grouped Convolution」？

https://www.zhihu.com/question/60484190/answer/1507783179

02

双边滤波

高斯滤波器的权重完全由距离决定。在大块颜色差不多、偶有噪点的区域，它可以把颜色平均化，从而过滤掉噪点。但是在颜色变化剧烈的边缘区域，它还是一视同仁地把所有像素做加权平均，这让本应该清晰锐利的边缘也变得模糊不清了，这就造成了如下图所示的效果，在做人像美颜时是不希望看到的。

这里，就引入了双边滤波（bilateral filtering）。

双边滤波的权重公式也基于高斯函数。但和高斯滤波的区别是，决定卷积核权重的，不单纯是像素之间的空间距离，还包括像素之间的亮度差异。以卷积核中心为坐标原点，该处像素值为I(0,0)。那么，坐标为 (u, v) 处的像素，对应的权重为：

(2) 中 exp 的第一个指数项和高斯核函数相同，与像素的空间距离有关；第二个指数项则是像素值距离的函数。以e为底对这两项做指数运算，再相乘即得到了公式 (2)。根据公式 (2) 计算的卷积核有如下性质：

• 距离中心像素越远的像素，其权重就越小
• 亮度和中心像素亮度差异越大的像素，其权重就越小

第一条比较好理解，第二条的含义是，像素只和与自己相似的像素求平均，不和与自己差别太大的像素求平均。例如上图中，人脸部分的像素和背景部分的像素差异过于显著，这将导致在对边缘区域做卷积时，背景侧像素对人脸侧的像素的权重极小，反之同理。这就达到了保持边缘 (edge-preserving) 的特性。

03

代码实现

由于 (2) 中卷积核的权重不仅仅依赖于空间距离，还依赖于像素的亮度，因此卷积核的权重是不固定的，不能简单地利用 pytorch 的 conv2d 来实现。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]，发现正是 x 左上角的六个元素
z[0,0]
Out[11]: 
tensor([[0, 1, 2],
        [4, 5, 6]])

# z[0, 1] 也同样符合预期
z[0,1]
Out[12]: 
tensor([[1, 2, 3],
        [5, 6, 7]])

实现的思路是：把原始图像 unfold 成一个个的 patch，对每个 patch 计算权重以及加权平均。代码如下：

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)
    # 根据像素亮度差，计算权重矩阵
    weights_color = torch.exp(-(diff_color ** 2) / (2 * sigmaColor ** 2))
    # 归一化权重矩阵
    weights_color = weights_color / weights_color.sum(dim=(-1, -2), keepdim=True)
    
    # 获取 gaussian kernel 并将其复制成和 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)
    
    # 两个权重矩阵相乘得到总的权重矩阵
    weights = weights_space * weights_color
    # 总权重矩阵的归一化参数
    weights_sum = weights.sum(dim=(-1, -2))
    # 加权平均
    weighted_pix = (weights * patches).sum(dim=(-1, -2)) / weights_sum
    return weighted_pix

最终结果为下图，雀斑都没了！同时人脸的轮廓和五官的细节依然被很好地保留下来：

输入图片尺寸为 256 x 256，ksize=15，sigmaColor=0.15，sigmaSpace=5 。

需要注意的是，由于 bilateral filter 的权重和像素值相关，因此设置 sigmaColor 时要注意输入图像的像素范围，看清楚到底是 0-1 还是 0-255（上图像素范围为 0-1）。

04

总结

本文介绍了双边滤波的基本原理，并附带了 pytorch 的实现。虽然不如 opencv 快，但优点是 backward trackable ，适合包装为模块加入网络中。利用 unfold 实现的缺点是很占内存/显存，kernelsize 越大，unfold 出来的冗余数据就越多，如果有大神知道更高效的实现方式，还望不吝赐教。

05

后记

我发现网上搜到的很多磨皮祛斑的算法，主要的目标是设计一个高通滤波器，从而得到一个基于像素亮度的 mask，亮的地方权重大（对应皮肤区域），暗的地方权重小（对应雀斑、噪点区域）。将原图 I 和 模糊化的图I_blur（各种模糊化方式都可以，目标是把较暗的斑点模糊掉）利用 mask 融合：I_mask+I_blur(1-mask) 。这种方法既保留了原图的细节，又能模糊掉斑点，不过在不同图片上应用时，仍然免不了调整一些超参数，而真有调参的功夫，直接调一下双边滤波的几个参数，最后得到的效果未必比那些复杂的算法差。

本文目的在于学术交流，并不代表本公众号赞同其观点或对其内容真实性负责，版权归原作者所有，如有侵权请告知删除。

“他山之石”历史文章

更多他山之石专栏文章，

请点击文章底部“阅读原文”查看

分享、点赞、在看，给个三连击呗！