python - 需要 scipy signal.fftconvolve 的 Tensorflow/Keras 等效项

Question

我想在Tensorflow/Keras中使用scipy.signal.fftconvolve，有什么办法可以做到吗？

现在我正在使用以下代码:

window = np.tile(window, (1, 1, 1, 3))
tf.nn.conv2d(img1, window, strides=[1,1,1,1], padding='VALID')

这些行是否相当于:

signal.fftconvolve(img1, window, mode='valid')

Answer 1

实现

FFT 卷积可以在 tensorflow 中相对容易地实现。以下内容非常严格地遵循scipy.signal.fftconvolve

import tensorflow as tf

def _centered(arr, newshape):
    # Return the center newshape portion of the array.
    currshape = tf.shape(arr)[-2:]
    startind = (currshape - newshape) // 2
    endind = startind + newshape
    return arr[..., startind[0]:endind[0], startind[1]:endind[1]]

def fftconv(in1, in2, mode="full"):
    # Reorder channels to come second (needed for fft)
    in1 = tf.transpose(in1, perm=[0, 3, 1, 2])
    in2 = tf.transpose(in2, perm=[0, 3, 1, 2])

    # Extract shapes
    s1 = tf.convert_to_tensor(tf.shape(in1)[-2:])
    s2 = tf.convert_to_tensor(tf.shape(in2)[-2:])
    shape = s1 + s2 - 1

    # Compute convolution in fourier space
    sp1 = tf.spectral.rfft2d(in1, shape)
    sp2 = tf.spectral.rfft2d(in2, shape)
    ret = tf.spectral.irfft2d(sp1 * sp2, shape)

    # Crop according to mode
    if mode == "full":
        cropped = ret
    elif mode == "same":
        cropped = _centered(ret, s1)
    elif mode == "valid":
        cropped = _centered(ret, s1 - s2 + 1)
    else:
        raise ValueError("Acceptable mode flags are 'valid',"
                         " 'same', or 'full'.")

    # Reorder channels to last
    result = tf.transpose(cropped, perm=[0, 2, 3, 1])
    return result

示例

将宽度为 20 像素的高斯平滑应用于标准“面部”图像的简单示例如下:

if __name__ == '__main__':
    from scipy import misc
    import matplotlib.pyplot as plt
    from tensorflow.python.ops import array_ops, math_ops
    session = tf.InteractiveSession()

    # Create gaussian
    std = 20
    grid_x, grid_y = array_ops.meshgrid(math_ops.range(3 * std),
                                        math_ops.range(3 * std))
    grid_x = tf.cast(grid_x[None, ..., None], 'float32')
    grid_y = tf.cast(grid_y[None, ..., None], 'float32')

    gaussian = tf.exp(-((grid_x - 1.5 * std) ** 2 + (grid_y - 1.5 * std) ** 2) / std ** 2)
    gaussian = gaussian / tf.reduce_sum(gaussian)

    face = misc.face(gray=False)[None, ...].astype('float32')

    # Apply convolution
    result = fftconv(face, gaussian, 'same')
    result_r = session.run(result)

    # Show results
    plt.figure('face')
    plt.imshow(face[0, ...] / 256.0)

    plt.figure('convolved')
    plt.imshow(result_r[0, ...] / 256.0)