diff --git a/CDTools/tools/image_processing.py b/CDTools/tools/image_processing.py index 3476785..228c72c 100644 --- a/CDTools/tools/image_processing.py +++ b/CDTools/tools/image_processing.py @@ -7,10 +7,10 @@ from CDTools.tools import cmath def centroid(im, dims=2): """Returns the centroid of an image or a stack of images - + By default, the last two dimensions are used in the calculation and the remainder of the dimensions are passed through. - + Beware that the meaning of the centroid is not well defined if your image contains values less than 0 @@ -34,14 +34,14 @@ def centroid(im, dims=2): def centroid_sq(im, dims=2, comp=False): """Returns the centroid of the square of an image or stack of images - + By default, the last two dimensions are used in the calculation and the remainder of the dimensions are passed through. If the "comp" flag is set, it will be assumed that the last dimension represents the real and imaginary part of a complex number, and the centroid will be calculated for the magnitude squared of those numbers - + Args: im (t.Tensor) : An image or stack of images to calculate from dims (int) : Default 2, how many trailing dimensions to calculate for @@ -56,13 +56,13 @@ def centroid_sq(im, dims=2, comp=False): return centroid(im_sq, dims=dims) - + def find_subpixel_shift(im1, im2, search_around=(0,0), resolution=10): """Calculates the subpixel shift between two images by maximizing the autocorrelation - + This function only searches in a 2 pixel by 2 pixel box around the specified search_around parameter. The calculation is done using the - approach outlined in "Efficient subpixel image registration algorithms", + approach outlined in "Efficient subpixel image registration algorithms", Optics Express (2008) by Manual Guizar-Sicarios et al. Args: @@ -77,7 +77,7 @@ def find_pixel_shift(im1, im2): This function simply takes the circular correlation with an FFT and returns the position of the maximum of that correlation - + Args: im1 (t.Tensor): The first real or complex-valued torch tensor im2 (t.Tensor): The second real or complex-valued torch tensor @@ -93,10 +93,9 @@ def find_shift(im1, im2, resolution=10): This function starts by calculating the maximum shift to integer pixel resolution, and then searchers the nearby area to calculate a subpixel shift - + Args: im1 (t.Tensor): The first real or complex-valued torch tensor im2 (t.Tensor): The second real or complex-valued torch tensor resolution (int): Default is 10, the resolution to calculate to in units of 1/n """ - diff --git a/CDTools/tools/plotting.py b/CDTools/tools/plotting.py new file mode 100644 index 0000000..be7f335 --- /dev/null +++ b/CDTools/tools/plotting.py @@ -0,0 +1,110 @@ +from __future__ import division, print_function, absolute_import + +from CDTools.tools import cmath +import torch as t +import numpy as np +import matplotlib.pyplot as plt +from matplotlib.colors import hsv_to_rgb + + +def colorize(z): + """ Returns RGB values for a complex color plot given a complex array + This function returns a set of RGB values that can be used directly + in a call to imshow based on an input complex numpy array (not a + torch tensor representing a complex field) + + Args: + z (array_like) : A complex-valued array + Returns: + list : A list of arrays for R,G, and B channels of an image. + + """ + + amp = np.abs(z) + rmin = 0 + rmax = np.max(amp) + amp = np.where(amp < rmin, rmin, amp) + amp = np.where(amp > rmax, rmax, amp) + ph = np.angle(z, deg=1) + 90 + # HSV are values in range [0,1] + h = (ph % 360) / 360 + s = 0.85 * np.ones_like(h) + v = (amp - rmin) / (rmax - rmin) + + return hsv_to_rgb(np.dstack((h,s,v))) + + +def plot_1d(im, **kwargs): + pass + +def plot_amplitude(im, fig = None, basis = np.array([[0,-1], [-1,0], [0,0]]), **kwargs): + """ Plots the amplitude of a complex Tensor or numpy array with dimensions NxMx2. + Args: + im (t.Tensor) : An image with dimensions NxMx2. + fig (matplotlib.figure.Figure) : A matplotlib figure to use to plot. If None, + a new figure is created with an Axes subplot at 111. + basis (array-like) : The probe basis, used to put the axis labels in real space units. + Should have dimensions 3x2 + **kwargs: Can be used to set any keyword arguments for the matplotlib.axes.Axes class + (see https://matplotlib.org/api/axes_api.html#the-axes-class) + """ + if fig is None: + fig = plt.figure() + ax = fig.add_subplot(111, **kwargs) + basis_norm = np.linalg.norm(basis, axis = -1) + if isinstance(im, t.Tensor): + absolute = cmath.cabs(im).detach().cpu().numpy() + else: + absolute = np.absolute(im) + plt.imshow(absolute, cmap = 'viridis', extent = [0, absolute.shape[-1]*basis_norm[1], 0, absolute.shape[-2]*basis_norm[0]]) + plt.colorbar() + return fig + +def plot_phase(im, fig = None, basis = np.array([[0,-1], [-1,0], [0,0]]), **kwargs): + """ Plots the phase of a complex Tensor or numpy array with dimensions NxMx2. + Args: + im (t.Tensor) : An image with dimensions NxMx2. + fig (matplotlib.figure.Figure) : A matplotlib figure to use to plot. If None, + a new figure is created with an Axes subplot at 111. + basis (array-like) : The probe basis, used to put the axis labels in real space units. + Should have dimensions 3x2 + **kwargs: Can be used to set any keyword arguments for the matplotlib.axes.Axes class + (see https://matplotlib.org/api/axes_api.html#the-axes-class) + """ + if fig is None: + fig = plt.figure() + ax = fig.add_subplot(111, **kwargs) + # If the user has matplotlib >=3.0, use the preferred colormap + if isinstance(im, t.Tensor): + phase = cmath.cphase(im).detach().cpu().numpy() + else: + phase = np.angle(im) + basis_norm = np.linalg.norm(basis, axis = -1) + try: plt.imshow(phase, cmap = 'twilight', extent = [0, phase.shape[-1]*basis_norm[1], 0, phase.shape[-2]*basis_norm[0]]) + except: plt.imshow(phase, cmap = 'hsv', extent = [0, phase.shape[-1]*basis_norm[1], 0, phase.shape[-2]*basis_norm[0]]) + plt.colorbar() + return fig + +def plot_colorized(im, fig = None, basis = np.array([[0,-1], [-1,0], [0,0]]), **kwargs): + """ Plots the colorized version of a complex Tensor or numpy array with dimensions NxMx2. + The darkness corresponds to the intensity of the image, and the color corresponds + to the phase. + + Args: + im (t.Tensor) : An image with dimensions NxMx2. + fig (matplotlib.figure.Figure) : A matplotlib figure to use to plot. If None, + a new figure is created with an Axes subplot at 111. + basis (array-like) : The probe basis, used to put the axis labels in real space units. + Should have dimensions 3x2 + **kwargs: Can be used to set any keyword arguments for the matplotlib.axes.Axes class + (see https://matplotlib.org/api/axes_api.html#the-axes-class) + """ + if fig is None: + fig = plt.figure() + ax = fig.add_subplot(111, **kwargs) + if isinstance(im, t.Tensor): + im = cmath.torch_to_complex(im.detach().cpu()) + basis_norm = np.linalg.norm(basis, axis = -1) + colorized = colorize(im) + plt.imshow(colorized, extent = [0, im.shape[-1]*basis_norm[1], 0, im.shape[-2]*basis_norm[0]]) + return fig diff --git a/tests/conftest.py b/tests/conftest.py index a062cf6..6e90f24 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -13,6 +13,18 @@ import datetime # # +def pytest_addoption(parser): + parser.addoption( + "--plot", action="store", default=False, help="plot: True to show test plots" + ) + + +@pytest.fixture +def show_plot(request): + return request.config.getoption("--plot") + + + @pytest.fixture(scope='module') def ptycho_cxi_1(): """Creates an example file for CXI ptychography. This file is defined @@ -20,12 +32,12 @@ def ptycho_cxi_1(): defined. It will return both a dictionary describing what is expected to be loaded and a file with the data stored in it. """ - + expected = {} f = h5py.File('ptycho_cxi_1',driver='core',backing_store=False) # Start by defining the basic structure - f.create_dataset('cxi_version', data=150) + f.create_dataset('cxi_version', data=150) f.create_dataset('number_of_entries',data=1) # Then define a bunch of metadata for entry_1 @@ -74,7 +86,7 @@ def ptycho_cxi_1(): energy = np.float32(1.3618e-16) #Joules, = 850 eV source1f['energy'] = energy expected['wavelength'] = np.float32(1.9864459e-25) / energy - source1f['wavelength'] = expected['wavelength'] + source1f['wavelength'] = expected['wavelength'] d1f = i1f.create_group('detector_1') expected['detector'] = {} @@ -96,7 +108,7 @@ def ptycho_cxi_1(): d1f.create_dataset('mask',data=mask) data1f = e1f.create_group('data_1') - + data = np.random.rand(100,256,256).astype(np.float32) expected['data'] = data d1f.create_dataset('data',data=data) @@ -111,7 +123,7 @@ def ptycho_cxi_1(): data1f['translation'] = h5py.SoftLink('/entry_1/sample_1/geometry_1/translation') d1f['translation'] = h5py.SoftLink('/entry_1/sample_1/geometry_1/translation') expected['translations'] = -translations - + yield f, expected f.close() @@ -131,12 +143,12 @@ def ptycho_cxi_2(): * Is missing many allowed metadata attributes """ - + expected = {} f = h5py.File('ptycho_cxi_2',driver='core',backing_store=False) # Start by defining the basic structure - f.create_dataset('cxi_version', data=150) + f.create_dataset('cxi_version', data=150) f.create_dataset('number_of_entries',data=1) # Then define a bunch of metadata for entry_1 @@ -158,7 +170,7 @@ def ptycho_cxi_2(): energy = np.float32(1.3618e-16) #Joules, = 850 eV expected['wavelength'] = np.float32(1.9864459e-25) / energy - source1f['wavelength'] = expected['wavelength'] + source1f['wavelength'] = expected['wavelength'] d1f = i1f.create_group('detector_1') expected['detector'] = {} @@ -176,7 +188,7 @@ def ptycho_cxi_2(): expected['mask'] = None data1f = e1f.create_group('data_1') - + data = np.random.rand(100,256,256).astype(np.float32) expected['data'] = data d1f.create_dataset('data',data=data) @@ -187,7 +199,7 @@ def ptycho_cxi_2(): translations = np.arange(300).reshape((100,3)).astype(np.float32) g1f.create_dataset('translation',data=translations) expected['translations'] = -translations - + yield f, expected f.close() @@ -206,12 +218,12 @@ def ptycho_cxi_3(): * Defines the sample to detector distance but no corner location * Is missing some of the allowed metadata """ - + expected = {} f = h5py.File('ptycho_cxi_3',driver='core',backing_store=False) # Start by defining the basic structure - f.create_dataset('cxi_version', data=150) + f.create_dataset('cxi_version', data=150) f.create_dataset('number_of_entries',data=1) # Then define a bunch of metadata for entry_1 @@ -250,7 +262,7 @@ def ptycho_cxi_3(): d1f.create_dataset('mask',data=mask) data1f = e1f.create_group('data_1') - + data = np.random.rand(100,256,256).astype(np.float32) expected['data'] = data data1f.create_dataset('data',data=data) @@ -261,7 +273,7 @@ def ptycho_cxi_3(): translations = np.arange(300).reshape((100,3)).astype(np.float32) data1f.create_dataset('translation',data=translations) expected['translations'] = -translations - + yield f, expected f.close() @@ -280,4 +292,3 @@ def test_ptycho_cxis(ptycho_cxi_1, ptycho_cxi_2, ptycho_cxi_3): on the cxi files. """ return [ptycho_cxi_1, ptycho_cxi_2, ptycho_cxi_3] - diff --git a/tests/tools/test_plotting.py b/tests/tools/test_plotting.py new file mode 100644 index 0000000..f9d599f --- /dev/null +++ b/tests/tools/test_plotting.py @@ -0,0 +1,54 @@ +from __future__ import division, print_function, absolute_import + +from CDTools.tools import cmath +from CDTools.tools import plotting +from CDTools.tools import initializers +import numpy as np +import pytest +import torch as t +import scipy.misc +import matplotlib.pyplot as plt + + + +def test_plot_amplitude(show_plot): + # Test with tensor + im = cmath.complex_to_torch(scipy.misc.ascent().astype(np.float64)) + plotting.plot_amplitude(im, basis = np.array([[1,1], [1,1], [0,0]]), title = 'Test Amplitude') + if show_plot: + plt.show() + + # Test with numpy array + im = scipy.misc.ascent().astype(np.complex128) + plotting.plot_amplitude(im, title = 'Test Amplitude') + if show_plot: + plt.show() + + +def test_plot_phase(show_plot): + # Test with tensor + im = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1]) + plotting.plot_phase(im, title = 'Test Phase') + if show_plot: + plt.show() + + # Test with numpy array + im = cmath.torch_to_complex(initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1])) + plotting.plot_phase(im, title = 'Test Phase', basis = np.array([[1,1], [1,1], [0,0]])) + if show_plot: + plt.show() + +def test_plot_colorize(show_plot): + # Test with tensor + gaussian = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1]) + im = cmath.cmult(gaussian, cmath.complex_to_torch(scipy.misc.ascent().astype(np.float64))) + plotting.plot_colorized(im, title = 'Test Colorize', basis = np.array([[1,1], [1,1], [0,0]])) + if show_plot: + plt.show() + + # Test with numpy array + gaussian = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1]) + im = cmath.torch_to_complex(cmath.cmult(gaussian, cmath.complex_to_torch(scipy.misc.ascent().astype(np.float64)))) + plotting.plot_colorized(im, title = 'Test Colorize') + if show_plot: + plt.show()