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, k_re, k_im):
k_re, k_im = make_contiguous(k_re, k_im)
x_r, x_i = ifft2(k_re, k_im)
if self.norm == 'ortho':
N = np.sqrt(x_r.size(-1) * x_r.size(-2))
x_r *= N
x_i *= N
return x_r, x_i
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 = ifft2(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')