def test_expand():
X = torch.randn(2, 2, 4, 4).cuda().double()
zeros = torch.zeros(2, 2, 4, 4).cuda().double()
r1, r2 = cfft.rfft2(X)
c1, c2 = cfft.fft2(X, zeros)
assert np.allclose(cfft.expand(r1).cpu().numpy(), c1.cpu().numpy())
assert np.allclose(
cfft.expand(r2, imag=True).cpu().numpy(),
c2.cpu().numpy())
r1, r2 = cfft.rfft3(X)
c1, c2 = cfft.fft3(X, zeros)
assert np.allclose(cfft.expand(r1).cpu().numpy(), c1.cpu().numpy())
assert np.allclose(
cfft.expand(r2, imag=True).cpu().numpy(),
c2.cpu().numpy())
X = torch.randn(2, 2, 5, 5).cuda().double()
zeros = torch.zeros(2, 2, 5, 5).cuda().double()
r1, r2 = cfft.rfft3(X)
c1, c2 = cfft.fft3(X, zeros)
assert np.allclose(
cfft.expand(r1, odd=True).cpu().numpy(),
c1.cpu().numpy())
assert np.allclose(
cfft.expand(r2, imag=True, odd=True).cpu().numpy(),
c2.cpu().numpy())
def forward(ctx, input, filter_height, filter_width):
ctx.save_for_backward(input)
y_spectral = fft.fft2(input)
y_spectral = _crop_spectrum(y_spectral)
y_spectral = _treat_corner_cases(y_spectral)
cropped = fft.ifft2(y_spectral)
return cropped.float()
def forward(self, X_re, X_im):
X_re, X_im = make_contiguous(X_re, X_im)
k_re, k_im = fft2(X_re, X_im)
if self.norm == 'ortho':
N = np.sqrt(k_re.size(-1) * k_re.size(-2))
k_re /= N
k_im /= N
return k_re, k_im
def backward(ctx, grad_output):
# Retrieve original tensor shape for _pad_spectrum
orig = ctx.saved_variables
orig_shape = orig[0].shape[3]
z = fft.fft2(z)
z = _remove_redundancy(z)
z = _pad_spectrum(z, orig_shape)
z = _recover_map(z)
return Variable(z), None, None
def backward(self, grad_output_re, grad_output_im):
grad_output_re, grad_output_im = make_contiguous(
grad_output_re, grad_output_im)
gi, gr = fft2(grad_output_im, grad_output_re)
if self.norm == 'ortho':
N = np.sqrt(gi.size(-1) * gi.size(-2))
gi /= N
gr /= N
return gr, gi
def run_fft(x, z):
if torch.cuda.is_available():
y1, y2 = cfft.fft2(x, z)
x_np = x.cpu().numpy().squeeze()
y_np = nfft.fft2(x_np)
assert np.allclose(y1.cpu().numpy(), y_np.real)
assert np.allclose(y2.cpu().numpy(), y_np.imag)
# assert np.allclose(y1[1,0].cpu().numpy(), nfft.fft2(x_np[1,0]).real)
x0, z0 = cfft.ifft2(y1, y2)
x0_np = nfft.ifft2(y_np)
assert np.allclose(x0.cpu().numpy(), x0_np.real)
assert np.allclose(z0.cpu().numpy(), x0_np.imag)
else:
print("Cuda not available, cannot test.")
import numpy as np
import torch
import pytorch_fft.fft as fft
from PIL import Image
im = np.array(Image.open('rome.jpg').convert('LA'), dtype=np.uint8)
# Image.fromarray(im, mode='LA').save('newrome.png')
print('min: {}, max: {}'.format(np.min(im), np.max(im)))
im = torch.from_numpy(im).float().cuda()
im = im.permute(2, 0, 1).contiguous()
print(im.type())
print(im.size())
B_real, B_imag = fft.fft2(im, torch.zeros(*im.size()).cuda())
im, _ = fft.ifft2(B_real, torch.zeros(*im.size()).cuda())
im = im.permute(1, 2, 0).contiguous()
im = im.byte().cpu().numpy()
Image.fromarray(im, mode='LA').save('newrometorched.png')
print('min: {}, max: {}'.format(np.min(im), np.max(im)))
import torch
import pytorch_fft.fft as fft
import numpy as np
from PIL import Image
imreal = torch.from_numpy(np.array(Image.open('rome.png').convert('LA'))).cuda()
Image.fromarray(imreal.cpu().numpy(), mode='LA').save('preimg.png')
imreal = imreal.unsqueeze(0).float()
imimag = torch.zeros(*imreal.size()).cuda()
fft_version = fft.fft2(imreal, imimag)[0]
newimreal, newimimag = fft.ifft2(fft_version, torch.zeros(*fft_version.size()).cuda())
def pre(im):
# im = im.abs().log()
# mini = torch.min(im)
# maxi = torch.max(im)
# im = 255 * (im - mini) / (maxi - mini)
# print(im.min())
# print(im.max())
return im.byte()
real_img = pre(newimreal).cpu().numpy()
imag_img = pre(newimimag).cpu().numpy()
# np.save('real_img.npy', real_img)
Image.fromarray(real_img.squeeze(), mode='LA').save('real_img.png')
# Image.fromarray(imag_img.squeeze(), mode='LA').save('real_img.png')