使用Python中的fft2进行“有效”和“完整”卷积

Question

这是一个不完整的Python代码片段，用于通过快速傅里叶变换（FFT）进行卷积运算。

我想对它进行修改，使它支持以下两种卷积方式：1）有效卷积 2）全卷积

import numpy as np
from numpy.fft import fft2, ifft2

image = np.array([[3,2,5,6,7,8],
                  [5,4,2,10,8,1]])

kernel = np.array([[4,5],
                   [1,2]])

fft_size =  # what size should I put here for,
            # 1) valid convolution
            # 2) full convolution

convolution = ifft2(fft2(image, fft_size) * fft2(kernel, fft_size))

提前谢谢你。

Answer 1

对于一维数组 x 和 y，它们的长度分别是 L 和 M，如果你想要使用“完整”模式的快速傅里叶变换（FFT），你需要把数据填充到长度 L + M - 1。对于二维数组，你也要对每个维度都应用这个规则。

使用 numpy 库，你可以在处理二维数组时计算出所需的大小，方法是：

np.array(x.shape) + np.array(y.shape) - 1

如果你想实现“有效”模式，你需要先计算出“完整”的结果，然后从中切出有效的部分。对于一维数组，假设 L 大于 M，那么有效的数据就是完整数据中间的 L - M + 1 个元素。同样，对于二维数组，也要对每个维度应用这个规则。

举个例子，

import numpy as np
from numpy.fft import fft2, ifft2


def fftconvolve2d(x, y, mode="full"):
    """
    x and y must be real 2-d numpy arrays.

    mode must be "full" or "valid".
    """
    x_shape = np.array(x.shape)
    y_shape = np.array(y.shape)
    z_shape = x_shape + y_shape - 1
    z = ifft2(fft2(x, z_shape) * fft2(y, z_shape)).real

    if mode == "valid":
        # To compute a valid shape, either np.all(x_shape >= y_shape) or
        # np.all(y_shape >= x_shape).
        valid_shape = x_shape - y_shape + 1
        if np.any(valid_shape < 1):
            valid_shape = y_shape - x_shape + 1
            if np.any(valid_shape < 1):
                raise ValueError("empty result for valid shape")
        start = (z_shape - valid_shape) // 2
        end = start + valid_shape
        z = z[start[0]:end[0], start[1]:end[1]]

    return z

这是应用于你示例数据的函数：

In [146]: image
Out[146]: 
array([[ 3,  2,  5,  6,  7,  8],
       [ 5,  4,  2, 10,  8,  1]])

In [147]: kernel
Out[147]: 
array([[4, 5],
       [1, 2]])

In [148]: fftconvolve2d(image, kernel, mode="full")
Out[148]: 
array([[  12.,   23.,   30.,   49.,   58.,   67.,   40.],
       [  23.,   49.,   37.,   66.,  101.,   66.,   21.],
       [   5.,   14.,   10.,   14.,   28.,   17.,    2.]])

In [149]: fftconvolve2d(image, kernel, mode="valid")
Out[149]: array([[  49.,   37.,   66.,  101.,   66.]])

你可以增加更多的错误检查，也可以修改这个函数来处理复杂数组和多维数组。如果你做了这些改进，可能会得到类似于 scipy.signal.fftconvolve 的东西（https://github.com/scipy/scipy/blob/master/scipy/signal/signaltools.py#L210）：

In [152]: from scipy.signal import fftconvolve

In [153]: fftconvolve(image, kernel, mode="full")
Out[153]: 
array([[  12.,   23.,   30.,   49.,   58.,   67.,   40.],
       [  23.,   49.,   37.,   66.,  101.,   66.,   21.],
       [   5.,   14.,   10.,   14.,   28.,   17.,    2.]])

In [154]: fftconvolve(image, kernel, mode="valid")
Out[154]: array([[  49.,   37.,   66.,  101.,   66.]])