假设问题实际上要求对高斯函数进行卷积(即高斯模糊,这是标题和已接受答案给我的印象),而不是进行乘法(即晕影效果,这是问题的演示代码产生的效果),那么下面是一个纯PyTorch版本的代码,它不需要安装torchvision(否则可以使用torchvision.transforms.GaussianBlur(),这是
Mushfirat Mohaimin's answer提出的方案):
from math import ceil
import torch
from torch.nn.functional import conv2d
from torch.distributions import Normal
def gaussian_kernel_1d(sigma: float, num_sigmas: float = 3.) -> torch.Tensor:
radius = ceil(num_sigmas * sigma)
support = torch.arange(-radius, radius + 1, dtype=torch.float)
kernel = Normal(loc=0, scale=sigma).log_prob(support).exp_()
return kernel.mul_(1 / kernel.sum())
def gaussian_filter_2d(img: torch.Tensor, sigma: float) -> torch.Tensor:
kernel_1d = gaussian_kernel_1d(sigma)
padding = len(kernel_1d) // 2
img = img.unsqueeze(0).unsqueeze_(0)
img = conv2d(img, weight=kernel_1d.view(1, 1, -1, 1), padding=(padding, 0))
img = conv2d(img, weight=kernel_1d.view(1, 1, 1, -1), padding=(0, padding))
return img.squeeze_(0).squeeze_(0)
if __name__ == "__main__":
import matplotlib.pyplot as plt
img = torch.rand(size=(100, 100))
img_filtered = gaussian_filter_2d(img, sigma=1.5)
plt.subplot(121)
plt.imshow(img)
plt.subplot(122)
plt.imshow(img_filtered)
plt.show()
该代码使用了Andrei Bârsan在
this answer的评论中提到的可分离滤波器的基本思想。这意味着,使用2D高斯核进行卷积可以通过分别沿图像的列和行进行一次卷积来替换。总体上更有效率,因为对于边长为N的核,每个像素需要使用2N而不是N²次乘法。
因此,在提供的代码中,我们首先使用
gaussian_kernel_1d()
创建一个1D高斯核,然后在
gaussian_filter_2d()
中应用两次。
关于代码,还有一些注意事项:
参数
num_sigmas
控制我们实际采样高斯函数凸起的标准差数量,从而产生卷积核。由于高斯函数理论上具有无限的支持(意味着它永远不为零),这就需要在精度和卷积核大小之间进行权衡(这会影响速度和内存使用)。长度为
3 * sigma
应该足够了,因为它将覆盖对应高斯函数下99.7%的面积,通常可以满足支持的两个半部分。
与其使用
Normal().log_prob().exp_()
来生成卷积核,我们可以在此明确编写正态分布的函数,这可能会更有效率。事实上,我们可以编写
kernel = support.square_().mul_(-.5 / (sigma ** 2)).exp_()
,从而(1)原地更改
support
的值(因为我们不再需要它们),并且(2)甚至省略正态分布的归一化常数(因为我们必须在返回卷积核之前对其进行归一化)。
虽然我们使用
conv2d()
而不是
conv1d()
,但实际上我们仍然有两个1D卷积,因为我们在
conv2d()
中应用了一个N×1和1×N的卷积核。我们本可以使用
conv1d()
,但使用
conv2d()
时代码要简单得多。
在较新的PyTorch版本中,我们可以使用
conv2d(…, padding="same")
,而不是自己计算填充量。无论哪种情况,使用
conv2d()
的
padding
参数都意味着使用零填充。如果我们想要更多的填充选项,我们可以在卷积之前手动使用
torch.nn.functional.pad()
填充图像。
gaussian_weights
应该是由高斯函数生成的核函数,就像使用scipy.signal.gaussian
一样。 - jkr