import numpy as np
import imageio
import matplotlib.pyplot as plt
from phase_retrieval import fienup_phase_retrieval

np.random.seed(1)
image = imageio.imread('EncryptedImage_fft.png', as_gray=True)

# magnitudes = np.sqrt(np.abs(np.fft.fft2(image)))

magnitudes = np.sqrt(image)

result = fienup_phase_retrieval(magnitudes, steps=5,
                                verbose=False)

plt.show()
plt.subplot(121)
plt.imshow(magnitudes, cmap='gray')
plt.title('Image')
plt.subplot(122)
plt.imshow(result);
plt.title('Reconstruction')
plt.show()
