前几天研究了传统的美颜算法，了解到双边滤波（bilateral filtering）。在看懂原理后，为加深理解，抽时间用 pytorch 重新造了个轮子。虽然效率肯定比不上 opencv ，但当个小练习也不错。为了方便复习以及帮助初学者，在此记录。
高斯滤波
高斯核函数
图像领域的高斯滤波器是个二维的矩阵。矩阵中每个元素的值与它与矩阵中心的距离有关，计算公式就是二维高斯函数的公式：
为了让卷积前后的图像亮度保持不变，需要对 (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
高斯滤波器
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」？
双边滤波
高斯滤波器的权重完全由距离决定。在大块颜色差不多、偶有噪点的区域，它可以把颜色平均化，从而过滤掉噪点。但是在颜色变化剧烈的边缘区域，它还是一视同仁地把所有像素做加权平均，这让本应该清晰锐利的边缘也变得模糊不清了，这就造成了如下图所示的效果，在做人像美颜时是不希望看到的。
双边滤波的权重公式也基于高斯函数。但和高斯滤波的区别是，决定卷积核权重的，不单纯是像素之间的空间距离，还包括像素之间的亮度差异。以卷积核中心为坐标原点，该处像素值为。那么，坐标为 (u, v) 处的像素，对应的权重为：
(2) 中 exp 的第一个指数项和高斯核函数相同，与像素的空间距离有关；第二个指数项则是像素值距离的函数。以e为底对这两项做指数运算，再相乘即得到了公式 (2)。根据公式 (2) 计算的卷积核有如下性质:
- 距离中心像素越远的像素，其权重就越小
- 亮度和中心像素亮度差异越大的像素，其权重就越小
第一条比较好理解，第二条的含义是，像素只和与自己相似的像素求平均，不和与自己差别太大的像素求平均。例如上图中，人脸部分的像素和背景部分的像素差异过于显著，这将导致在对边缘区域做卷积时，背景侧像素对人脸侧的像素的权重极小，反之同理。这就达到了保持边缘 (edge-preserving) 的特性。
代码实现
由于 (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
最终结果为下图，雀斑都没了！同时人脸的轮廓和五官的细节依然被很好地保留下来：
需要注意的是，由于 bilateral filter 的权重和像素值相关，因此设置 sigmaColor 时要注意输入图像的像素范围，看清楚到底是 0-1 还是 0-255（上图像素范围为 0-1）。
总结
本文介绍了双边滤波的基本原理，并附带了 pytorch 的实现。虽然不如 opencv 快，但优点是 backward trackable ，适合包装为模块加入网络中。利用 unfold 实现的缺点是很占内存/显存，kernelsize 越大，unfold 出来的冗余数据就越多，如果有大神知道更高效的实现方式，还望不吝赐教。
后记
我发现网上搜到的很多磨皮祛斑的算法，主要的目标是设计一个高通滤波器，从而得到一个基于像素亮度的 mask，亮的地方权重大（对应皮肤区域），暗的地方权重小（对应雀斑、噪点区域）。将原图 I 和 模糊化的图 I_blur（各种模糊化方式都可以，目标是把较暗的斑点模糊掉）利用 mask 融合：I_mask+I_blur(1-mask) 。这种方法既保留了原图的细节，又能模糊掉斑点，不过在不同图片上应用时，仍然免不了调整一些超参数，而真有调参的功夫，直接调一下双边滤波的几个参数，最后得到的效果未必比那些复杂的算法差。