Code examples

import numpy as npdef fft(arr):
    """Perform a correctly shifted fast Fourier transform"""
    return np.fft.fftshift(np.fft.fftn(np.fft.ifftshift(arr)))


def ifft(arr):
    """Perform a correctly shifted inverse fast Fourier transform"""
    return np.fft.fftshift(np.fft.ifftn(np.fft.ifftshift(arr)))


def comp_log(comp_arr):
    """Take the logarithm of the amplitude of a complex array while keeping the phase intact."""
    amp = np.abs(comp_arr)
    phi = np.angle(comp_arr)
    return np.log(amp+1) * np.exp(1j*phi)


def normalize(arr):
    """Compress the range of an array onto the range (0,1)"""
    return (arr - np.min(arr)) / (np.max(arr) - np.min(arr))


def pad_to_size(arr, N_new):
    """Pad a 2D square array with zeros to a certain output size (bigger than the original)."""
    N = arr.shape[0]
    pad = (N_new - N) // 2
    if 2*pad + N == N_new:
        return np.pad(arr, ((pad, pad), (pad, pad)))
    else 2*pad + N == N_new - 1:
        return np.pad(arr, ((pad, pad+1), (pad, pad+1)))

import numpy as npfrom matplotlib import colors

def comp_to_rgb(comp_img):
    """Converts a complex array into an RGB color image (amplitude -> brightness , phase -> color)"""
    amp = normalize(np.abs(comp_img))  # The normalize function is defined above
    phi = normalize(np.angle(comp_img))
    one = np.ones_like(amp)
    hsv = np.dstack((phi, one, amp))
    return colors.hsv_to_rgb(hsv)

from skimage import io
import tkinter as tkfrom pathlib import Path
import shutilfrom matplotlib import pyplot as plt
from PIL import Image


def get_file():
    """Opens a GUI dialog to select a single file"""
    root = tk.Tk()
    root.withdraw()
    try:
        return fd.askopenfilename()
    except KeyError:
        return None


def get_save_as(ext):
    """Opens a save-as dialog to save a file with the given extension type"""
    root = tk.Tk()
    root.withdraw()
    try:
        return fd.asksaveasfilename(defaultextension=ext, filetypes=[(ext.lower(), ext.upper())])
    except KeyError:
        return


def load_image(fname=None):
    """Loads an image as a 2D numpy array. Color images will be loaded in grayscale."""
    if fname is None:
        fname = get_file()
    return imread(fname, as_gray=True)


def save_array_as_image(img, fname=None, cmap='plasma'):
    """Saves a 2D array as an image using one of matplotlib's builtin colormap."""
    if fname is None:
        fname = get_save_as('png')
    plt.imsave(fname, img, cmap=cmap)


def save_imstack_as_gif(imstack, fname=None, cmap='plasma'):
    """Saves a 3D array (image stack) as an animated gif. This method is a little janky. If someone knows a
       better way, please share it."""
    if fname is None:
        fname = get_save_as('gif')
    dname = Path(__file__).parent
    temp = dname / 'temp'
    temp.mkdir()
    [plt.imsave(f'{temp}/im_{n:03}.png', im, cmap=cmap, vmin=imstack.amin(), vmax=imstack.amax())
     for n, im in enumerate(imstack)]
    imstack = [Image.open(f'{temp}/im_{n:03}.png') for n, _ in enumerate(imstack)]
    imstack[0].save(fname, save_all=True, append_images=imstack[1:], duration=300, loop=0)
    shutil.rmtree(temp)