mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-09 21:12:42 +02:00
added plotting capabilities and tests (note: no colorbar yet for colorized plt)
This commit is contained in:
Maddie Cain
@@ -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
|
||||
"""
|
||||
|
||||
|
||||
@@ -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
|
||||
+26
-15
@@ -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]
|
||||
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user