I'm trying to implement a FFT convolution that mimics scipy.fftconvolve
using pyfftw for performance and pictures as input :
import numpy as np
import pyfftw
a = np.ones((6000, 4000), dtype='float32')
b = np.kaiser(25, 8)
b = np.outer(b, b).astype('float32')
class fftconvolve:
def __init__(self, A, B, domain, threads=8):
MK = B.shape[0]
NK = B.shape[1]
M = A.shape[0]
N = A.shape[1]
if domain =="same":
Y = M
X = N
elif domain == "valid":
Y = M - MK + 1
X = N - NK + 1
elif domain == "full":
Y = M + MK - 1
X = N + NK - 1
self.fft_A_obj = pyfftw.builders.rfft2(A, s=(M + MK -1, N + NK -1), threads=threads)
self.fft_B_obj = pyfftw.builders.rfft2(B, s=(M + MK -1, N + NK -1), threads=threads)
self.ifft_obj = pyfftw.builders.irfft2(self.fft_A_obj.output_array, s=(Y, X), threads=threads)
def __call__(self, A, B):
return self.ifft_obj(np.fft.ifftshift(
np.fft.fftshift(self.fft_A_obj(A)) * np.fft.fftshift(self.fft_B_obj(B))
))
Calling it :
plan = fftconvolve(a, b, "full", threads=8)
c_1 = plan(a, b)
c_1
Output :
array([[ 3.89971137e-06, 3.51986018e-05, 1.24518745e-04, ...,
1.25271297e-04, 3.56316777e-05, 4.04627326e-06],
[ 4.91737483e-05, 2.60021159e-04, 8.61040782e-04, ...,
8.63055116e-04, 2.61142646e-04, 4.95371969e-05],
[ 1.26523402e-04, 8.49825097e-04, 2.90915114e-03, ...,
2.90881563e-03, 8.49568460e-04, 1.26304061e-04],
...,
[ 1.28503540e-04, 8.52331228e-04, 2.91197700e-03, ...,
2.91016186e-03, 8.51134886e-04, 1.28111642e-04],
[ 2.14206957e-05, 2.32703838e-04, 8.34190170e-04, ...,
8.34319100e-04, 2.32750244e-04, 2.14206957e-05],
[ -8.42595455e-06, 2.29651105e-05, 1.12404508e-04, ...,
1.12760317e-04, 2.31778213e-05, -8.35505125e-06]], dtype=float32)
Calling the scipy equivalent :
c_2 = scipy.signal.fftconvolve(a, b, "full").astype(np.float32)
c_2
Output :
array([[ 5.47012860e-06, 3.68362089e-05, 1.26135841e-04, ...,
1.26135841e-04, 3.68362089e-05, 5.47012769e-06],
[ 3.68362089e-05, 2.48057506e-04, 8.49407224e-04, ...,
8.49407224e-04, 2.48057506e-04, 3.68362089e-05],
[ 1.26135841e-04, 8.49407224e-04, 2.90856976e-03, ...,
2.90856976e-03, 8.49407224e-04, 1.26135841e-04],
...,
[ 1.26135841e-04, 8.49407224e-04, 2.90856976e-03, ...,
2.90856976e-03, 8.49407224e-04, 1.26135841e-04],
[ 3.68362089e-05, 2.48057506e-04, 8.49407224e-04, ...,
8.49407224e-04, 2.48057506e-04, 3.68362089e-05],
[ 5.47012814e-06, 3.68362089e-05, 1.26135841e-04, ...,
1.26135841e-04, 3.68362089e-05, 5.47012814e-06]], dtype=float32)
Checking the outputs :
c_1 == c_2
Gives :
array([[False, False, False, ..., False, False, False],
[False, False, False, ..., False, False, False],
[False, False, False, ..., False, False, False],
...,
[False, False, False, ..., False, False, False],
[False, False, False, ..., False, False, False],
[False, False, False, ..., False, False, False]], dtype=bool)
And :
np.allclose(c_1, c_2)
gives :
False
So the output is not correct. Removing the fftshift
does not change a thing.
In my project, the scipy version gives a correct image, the pyfftw with my implementation gives blurry output.
Edit
I tested also in double
type (np.float64
) and while the raw result of the convolution is close enough (actually, scipy does the convolution in double), the picture is still bad :
Deconvolved with custom convolution here : it's not only blurry but you get fringes on the edges :
a = np.ones((6000, 4000), dtype='float64')
b = np.kaiser(25, 8)
b = np.outer(b, b).astype('float64')
Now :
np.allclose(c_1, c_2)
Returns :
True
What could give this result ?
This works as expected :
class fftconvolve:
def __init__(self, A, B, domain, threads=8):
MK = B.shape[0]
NK = B.shape[1]
M = A.shape[0]
N = A.shape[1]
if domain =="same":
self.Y = M
self.X = N
elif domain == "valid":
self.Y = M - MK + 1
self.X = N - NK + 1
elif domain == "full":
self.Y = M + MK - 1
self.X = N + NK - 1
self.M = M + MK - 1
self.N = N + NK - 1
a = np.pad(A, ((0, MK - 1), (0, NK - 1)), mode='constant')
b = np.pad(B, ((0, M - 1), (0, N - 1)), mode='constant')
self.fft_A_obj = pyfftw.builders.rfft2(a, s=(self.M, self.N), threads=threads)
self.fft_B_obj = pyfftw.builders.rfft2(b, s=(self.M, self.N), threads=threads)
self.ifft_obj = pyfftw.builders.irfft2(self.fft_A_obj.output_array, s=(self.M, self.N), threads=threads)
self.offset_Y = int(np.floor((self.M - self.Y)/2))
self.offset_X = int(np.floor((self.N - self.X)/2))
def __call__(self, A, B):
MK = B.shape[0]
NK = B.shape[1]
M = A.shape[0]
N = A.shape[1]
a = np.pad(A, ((0, MK - 1), (0, NK - 1)), mode='constant')
b = np.pad(B, ((0, M - 1), (0, N - 1)), mode='constant')
return self.ifft_obj(self.fft_A_obj(a) * self.fft_B_obj(b))[self.offset_Y:self.offset_Y + self.Y, self.offset_X:self.offset_X + self.X]