pyTEMlib 0.2025.4.2__py3-none-any.whl → 0.2025.9.1__py3-none-any.whl

pyTEMlib/image/image_distortion.py

@@ -0,0 +1,299 @@
+"""
+image_distortion part of pycroscopy
+
+Author: Gerd Duscher
+
+Distortions of scanned images are determined by comparing experimentally 
+obtained unit cells with ideal ones.
+
+"""
+import numpy as np
+import scipy
+import skimage
+
+import SimpleITK 
+
+from tqdm.auto import trange
+
+####################
+# Distortion Matrix
+####################
+def get_distortion_matrix(atoms: np.ndarray, ideal_lattice: np.ndarray) -> np.ndarray:
+    """    Calculates distortion matrix
+
+    Calculates the distortion matrix by comparing ideal and distorted Voronoi tiles
+    Parameters
+    ----------
+    atoms: numpy array (Nx2)
+        atomic positions
+    ideal_lattice: numpy array (Mx2)
+        ideal lattice positions
+
+    Returns
+    -------
+    numpy array (Nx4)
+        distortion matrix
+    """
+    vor = scipy.spatial.Voronoi(atoms)
+    # determine a middle Voronoi tile
+    ideal_vor = scipy.spatial.Voronoi(ideal_lattice)
+    near_center = np.average(ideal_lattice, axis=0)
+    index = np.argmin(np.linalg.norm(ideal_lattice - near_center, axis=0))
+
+    # the ideal vertices fo such an Voronoi tile (are there crystals with more than one voronoi?)
+    ideal_vertices = ideal_vor.vertices[ideal_vor.regions[ideal_vor.point_region[index]]]
+    ideal_vertices = get_significant_vertices(ideal_vertices - np.average(ideal_vertices, axis=0))
+
+    distortion_matrix = []
+    for index in trange(vor.points.shape[0]):
+
+        # determine vertices of Voronoi polygons of an atom with number index
+        poly_point = vor.points[index]
+        vertices = vor.vertices[vor.regions[vor.point_region[index]]]
+        poly_vertices = get_significant_vertices(vertices - poly_point)
+
+        # where ATOM has to be moved (not pixel)
+        ideal_point = ideal_lattice[index]
+
+        # transform voronoi to ideal one and keep transformation matrix A
+        uncorrected, corrected, _ = transform_voronoi(poly_vertices, ideal_vertices)
+
+        # pixel positions
+        corrected = corrected + ideal_point + (np.rint(poly_point) - poly_point)
+        for i in range(len(corrected)):
+            # original image pixels
+            x, y = uncorrected[i] + np.rint(poly_point)
+            # collect the two origin and target coordinates and store
+            distortion_matrix.append([x, y, corrected[i, 0], corrected[i, 1]])
+    print()
+    return np.array(distortion_matrix)
+
+
+def undistort(distortion_matrix, image_data):
+    """ Undistort image according to distortion matrix
+    
+    Uses the griddata interpolation of scipy to apply distortion matrix to image.
+    The distortion matrix contains in origin and target pixel coordinates
+    target is where the pixel has to be moved (floats)
+    
+    Parameters
+    ----------
+    distortion_matrix: numpy array (Nx2)
+        distortion matrix (format N x 2)
+    image_data: numpy array or sidpy.Dataset
+        image 
+        
+    Returns
+    -------
+    interpolated: numpy array
+        undistorted image
+    """
+
+    intensity_values = image_data[(distortion_matrix[:, 0].astype(int),
+                                   distortion_matrix[:, 1].astype(int))]
+
+    corrected = distortion_matrix[:, 2:4]
+
+    size_x, size_y = 2 ** np.round(np.log2(image_data.shape[0:2]))  # nearest power of 2
+    size_x = int(size_x)
+    size_y = int(size_y)
+    grid_x, grid_y = np.mgrid[0:size_x - 1:size_x * 1j, 0:size_y - 1:size_y * 1j]
+    print('interpolate')
+
+    interpolated = scipy.interpolate.griddata(np.array(corrected),
+                                              np.array(intensity_values),
+                                              (grid_x, grid_y), method='linear')
+    return interpolated
+
+
+def transform_voronoi(vertices, ideal_voronoi):
+    """ find transformation matrix A between a distorted polygon and a 
+        perfect reference one
+
+    Returns
+    -------
+    uncorrected: list of points: 
+        all points on a grid within original polygon
+    corrected: list of points: 
+        coordinates of these points where pixel have to move to
+    aa: 2x2 matrix A:  
+        transformation matrix
+    """
+    # Find Transformation Matrix, note polygons have to be ordered first.
+    sort_vert = []
+    for vert in ideal_voronoi:
+        sort_vert.append(np.argmin(np.linalg.norm(vertices - vert, axis=1)))
+    vertices = np.array(vertices)[sort_vert]
+
+    # Solve the least squares problem X * A = Y
+    # to find our transformation matrix aa = A
+    aa, _, _, _ = np.linalg.lstsq(vertices, ideal_voronoi, rcond=None)
+
+    # expand polygon to include more points in distortion matrix
+    vertices2 = vertices + np.sign(vertices)  # +np.sign(vertices)
+
+    ext_v = int(np.abs(vertices2).max() + 1)
+
+    polygon_grid = np.mgrid[0:ext_v * 2 + 1, :ext_v * 2 + 1] - ext_v
+    polygon_grid = np.swapaxes(polygon_grid, 0, 2)
+    polygon_array = polygon_grid.reshape(-1, polygon_grid.shape[-1])
+
+    p = skimage.measure.points_in_poly(polygon_array, vertices2)
+    uncorrected = polygon_array[p]
+    corrected = np.dot(uncorrected, aa)
+    return uncorrected, corrected, aa
+
+
+def get_maximum_view(distortion_matrix: np.ndarray):
+    """
+    Determines the largest rectangular view within a distorted image matrix 
+    ----------
+    distortion_matrix : np.ndarray
+        A 3D numpy array representing the distortion matrix of the image, 
+        where invalid pixels are marked with -1000.
+    Returns
+    -------
+    np.ndarray
+        A 1D numpy array of four integers [row_start, row_end, col_start, col_end] 
+        representing the coordinates of the maximal valid view within the 
+        distortion matrix.
+    """
+    distortion_matrix_extent = np.ones(distortion_matrix.shape[1:], dtype=int)
+    distortion_matrix_extent[distortion_matrix[0] == -1000.] = 0
+
+    area = distortion_matrix_extent
+    view_square = np.array([0, distortion_matrix.shape[1] - 1, 0,
+                            distortion_matrix.shape[2] - 1], dtype=int)
+    while np.array(np.where(area == 0)).shape[1] > 0:
+        view_square = view_square + [1, -1, 1, -1]
+        area = distortion_matrix_extent[view_square[0]:view_square[1],
+                                        view_square[2]:view_square[3]]
+
+    change = [-int(np.sum(np.min(distortion_matrix_extent[:view_square[0], view_square[2]:view_square[3]], axis=1))),
+              int(np.sum(np.min(distortion_matrix_extent[view_square[1]:, view_square[2]:view_square[3]], axis=1))),
+              -int(np.sum(np.min(distortion_matrix_extent[view_square[0]:view_square[1], :view_square[2]], axis=0))),
+              int(np.sum(np.min(distortion_matrix_extent[view_square[0]:view_square[1], view_square[3]:], axis=0)))]
+
+    return np.array(view_square) + change
+
+
+def get_significant_vertices(vertices, distance=3):
+    """Calculate average for  all points that are closer than distance apart, 
+    otherwise leave the points alone
+        
+        Parameters
+        ----------
+        vertices: numpy array (n,2)
+            list of points
+        distance: float
+            (in same scale as points )
+        
+        Returns
+        -------
+        ideal_vertices: list of floats
+            list of points that are all a minimum of 3 apart.
+    """
+    tt = scipy.spatial.KDTree(np.array(vertices))
+    near = tt.query_ball_point(vertices, distance)
+    ideal_vertices = []
+    for indices in near:
+        if len(indices) == 1:
+            ideal_vertices.append(vertices[indices][0])
+        else:
+            ideal_vertices.append(np.average(vertices[indices], axis=0))
+    ideal_vertices = np.unique(np.array(ideal_vertices), axis=0)
+    angles = np.arctan2(ideal_vertices[:, 1], ideal_vertices[:, 0])
+    ang_sort = np.argsort(angles)
+    ideal_vertices = ideal_vertices[ang_sort]
+
+    return ideal_vertices
+
+
+def undistort_sitk(image_data, distortion_matrix):
+    """    use simple ITK to undistort image
+    
+    Parameters
+    ----------
+    image_data: numpy array with size NxM 
+    distortion_matrix: sidpy.Dataset or numpy array with size 2 x P x Q
+    with P, Q >= M, N
+    
+    Returns
+    -------
+    image: numpy array MXN
+      
+    """
+    resampler = SimpleITK.ResampleImageFilter()
+    resampler.SetReferenceImage(SimpleITK.GetImageFromArray(image_data))
+    resampler.SetInterpolator(SimpleITK.sitkBSpline)
+    resampler.SetDefaultPixelValue(0)
+
+    distortion_matrix2 = distortion_matrix[:, :image_data.shape[0], :image_data.shape[1]]
+
+    displ2 = SimpleITK.Compose([SimpleITK.GetImageFromArray(-distortion_matrix2[1]),
+                           SimpleITK.GetImageFromArray(-distortion_matrix2[0])])
+    out_tx = SimpleITK.DisplacementFieldTransform(displ2)
+    resampler.SetTransform(out_tx)
+    out = resampler.Execute(SimpleITK.GetImageFromArray(image_data))
+    return SimpleITK.GetArrayFromImage(out)
+
+
+def undistort_stack_sitk(distortion_matrix, image_stack):
+    """
+    use simple ITK to undistort stack of image
+    input:
+    image: numpy array with size NxM 
+    distortion_matrix: h5 Dataset or numpy array with size 2 x P x Q
+    with P, Q >= M, N
+    output:
+    image M, N
+    """
+    resampler = SimpleITK.ResampleImageFilter()
+    resampler.SetReferenceImage(SimpleITK.GetImageFromArray(image_stack[0]))
+    resampler.SetInterpolator(SimpleITK.sitkBSpline)
+    resampler.SetDefaultPixelValue(0)
+
+    displ2 = SimpleITK.Compose([SimpleITK.GetImageFromArray(-distortion_matrix[1]),
+                           SimpleITK.GetImageFromArray(-distortion_matrix[0])])
+    out_tx = SimpleITK.DisplacementFieldTransform(displ2)
+    resampler.SetTransform(out_tx)
+    interpolated = np.zeros(image_stack.shape)
+    nimages = image_stack.shape[0]
+    for i in trange(nimages):
+        out = resampler.Execute(SimpleITK.GetImageFromArray(image_stack[i]))
+        interpolated[i] = SimpleITK.GetArrayFromImage(out)
+    return interpolated
+
+
+def undistort_stack(distortion_matrix, data):
+    """ Undistort stack with distortion matrix
+    
+    Use the griddata interpolation of scipy to apply distortion matrix to image
+    The distortion matrix contains in each pixel where the pixel has to be 
+    moved (floats)
+
+    Parameters
+    ----------
+    distortion_matrix: numpy array
+        distortion matrix to undistort image (format image.shape[0], 
+                                            image.shape[2], 2)
+    data: numpy array or sidpy.Dataset
+        image
+    """
+    corrected = distortion_matrix[:, 2:4]
+    intensity_values = data[:, distortion_matrix[:, 0].astype(int),
+                            distortion_matrix[:, 1].astype(int)]
+    size_x, size_y = 2 ** np.round(np.log2(data.shape[1:]))  # nearest power of 2
+    size_x = int(size_x)
+    size_y = int(size_y)
+
+    grid_x, grid_y = np.mgrid[0:size_x - 1:size_x * 1j, 0:size_y - 1:size_y * 1j]
+    interpolated = np.zeros([data.shape[0], size_x, size_y])
+    nimages = data.shape[0]
+    for i in trange(nimages):
+        interpolated[i, :, :] = scipy.interpolate.griddata(corrected,
+                                                           intensity_values[i, :],
+                                                           (grid_x, grid_y),
+                                                           method='linear')
+    print(':-) \n You have successfully completed undistortion of image stack')
+    return interpolated

pyTEMlib/image/image_fft.py

@@ -0,0 +1,277 @@
+"""
+Image_fft
+part of pycroscopy
+
+author: Gerd Duscher, UTK
+
+"""
+
+import itertools
+import collections
+
+import numpy as np
+import scipy
+import sklearn
+
+import sidpy
+
+def fourier_transform(dset: sidpy.Dataset) -> sidpy.Dataset:
+    """
+        Reads information into dictionary 'tags', performs 'FFT', and provides 
+        a smoothed FT and reciprocal and intensity limits for visualization.
+
+        Parameters
+        ----------
+        dset: sidpy.Dataset
+            image
+
+        Returns
+        -------
+        fft_dset: sidpy.Dataset
+            Fourier transform with correct dimensions
+
+        Example
+        -------
+        >>> fft_dataset = fourier_transform(sidpy_dataset)
+        >>> fft_dataset.plot()
+    """
+
+    assert isinstance(dset, sidpy.Dataset), 'Expected a sidpy Dataset'
+
+    selection = []
+    image_dims = dset.get_image_dims(return_axis=True)
+    if dset.data_type.name == 'IMAGE_STACK':
+        stack_dim = dset.get_dimensions_by_type('TEMPORAL')
+
+        if len(image_dims) != 2:
+            raise ValueError('need at least two SPATIAL dimension for an image stack')
+
+        for i in range(dset.ndim):
+            if i in image_dims:
+                selection.append(slice(None))
+            if len(stack_dim) == 0:
+                selection.append(slice(None))
+            elif i in stack_dim:
+                selection.append(slice(None))
+            else:
+                selection.append(slice(0, 1))
+
+        image_stack = np.squeeze(np.array(dset)[selection])
+        new_image = np.sum(np.array(image_stack), axis=stack_dim)
+    elif dset.data_type.name == 'IMAGE':
+        new_image = np.array(dset)
+    else:
+        return None
+
+    new_image = new_image - new_image.min()
+
+    fft_transform = (np.fft.fftshift(np.fft.fft2(np.array(new_image))))
+
+    image_dims = dset.get_image_dims(return_axis=True)
+
+    units_x = '1/' + image_dims[0].units
+    units_y = '1/' + image_dims[1].units
+
+    fft_dset = sidpy.Dataset.from_array(fft_transform)
+    fft_dset.quantity = dset.quantity
+    fft_dset.units = 'a.u.'
+    fft_dset.data_type = 'IMAGE'
+    fft_dset.source = dset.title
+    fft_dset.modality = 'fft'
+
+    fft_dset.set_dimension(0, sidpy.Dimension(np.fft.fftshift(np.fft.fftfreq(new_image.shape[0],
+                                                                            d=dset.x[1]-dset.x[0])),
+                                              name='u', units=units_x, dimension_type='RECIPROCAL',
+                                              quantity='reciprocal_length'))
+    fft_dset.set_dimension(1, sidpy.Dimension(np.fft.fftshift(np.fft.fftfreq(new_image.shape[1],
+                                                                            d=dset.y[1]- dset.y[0])),
+                                              name='v', units=units_y, dimension_type='RECIPROCAL',
+                                              quantity='reciprocal_length'))
+    return fft_dset
+
+
+def power_spectrum(dset: sidpy.Dataset, smoothing: int=3) -> sidpy.Dataset:
+    """
+    Calculate power spectrum
+
+    Parameters
+    ----------
+    dset: sidpy.Dataset
+        image
+    smoothing: int
+        Gaussian smoothing
+
+    Returns
+    -------
+    power_spec: sidpy.Dataset
+        power spectrum with correct dimensions
+
+    """
+
+    fft_transform = fourier_transform(dset)  # dset.fft()
+    fft_mag = np.abs(fft_transform)
+    fft_mag2 = scipy.ndimage.gaussian_filter(fft_mag, sigma=(smoothing, smoothing), order=0)
+
+    power_spec = fft_transform.like_data(np.log(1.+fft_mag2))
+
+    # prepare mask
+    x, y = np.meshgrid(power_spec.v.values, power_spec.u.values)
+    mask = np.zeros(power_spec.shape)
+
+    mask_spot = x ** 2 + y ** 2 > 1 ** 2
+    mask = mask + mask_spot
+    mask_spot = x ** 2 + y ** 2 < 11 ** 2
+    mask = mask + mask_spot
+
+    mask[np.where(mask == 1)] = 0  # just in case of overlapping disks
+
+    minimum_intensity = np.array(power_spec)[np.where(mask == 2)].min() * 0.95
+    maximum_intensity = np.array(power_spec)[np.where(mask == 2)].max() * 1.05
+    power_spec.metadata = {'fft': {'smoothing': smoothing,
+                                   'minimum_intensity': minimum_intensity, 
+                                   'maximum_intensity': maximum_intensity}}
+    power_spec.title = 'power spectrum ' + power_spec.source
+    return power_spec
+
+
+def diffractogram_spots(dset: sidpy.Dataset,
+                        spot_threshold: float,
+                        return_center: bool = True,
+                        eps: float=0.1) -> tuple[np.ndarray, np.ndarray | list[float]]:
+    """Find spots in diffractogram and sort them by distance from center
+
+    Uses blob_log from scipy.spatial
+
+    Parameters
+    ----------
+    dset: sidpy.Dataset
+        diffractogram
+    spot_threshold: float
+        threshold for blob finder
+    return_center: bool, optional
+        return center of image if true
+    eps: float, optional
+        threshold for blob finder
+
+    Returns
+    -------
+    spots: numpy array
+        sorted position (x,y) and radius (r) of all spots
+    """
+
+    # spot detection (for future reference there is no symmetry assumed here)
+    data = np.array(np.log(1+np.abs(dset)))
+    data = data - data.min()
+    data = data/data.max()
+    # some images are strange and blob_log does not work on the power spectrum
+    try:
+        spots_random = scipy.features.blob_log(data, max_sigma=5, threshold=spot_threshold)
+    except ValueError:
+        spots_random = scipy.features.peak_local_max(np.array(data.T),
+                                                     min_distance=3,
+                                                     threshold_rel=spot_threshold)
+        spots_random = np.hstack(spots_random, np.zeros((spots_random.shape[0], 1)))
+
+    # print(f'Found {spots_random.shape[0]} reflections')
+
+    # Needed for conversion from pixel to Reciprocal space
+    image_dims = dset.get_image_dims(return_axis=True)
+    rec_scale = np.array([image_dims[0].slope, image_dims[1].slope])
+
+    spots_random[:, :2] = spots_random[:, :2]*rec_scale+[dset.u.values[0], dset.v.values[0]]
+    # sort reflections
+    spots_random[:, 2] = np.linalg.norm(spots_random[:, 0:2], axis=1)
+    spots_index = np.argsort(spots_random[:, 2])
+    spots = spots_random[spots_index]
+    # third row is angles
+    spots[:, 2] = np.arctan2(spots[:, 0], spots[:, 1])
+
+    center = [0, 0]
+
+    if return_center:
+        points = spots[:, 0:2]
+
+        # Calculate the midpoints between all points
+        reshaped_points = points[:, np.newaxis, :]
+        midpoints = (reshaped_points + reshaped_points.transpose(1, 0, 2)) / 2.0
+        midpoints = midpoints.reshape(-1, 2)
+
+        # Find the most dense cluster of midpoints
+        dbscan = sklearn.cluster.DBSCAN(eps=eps, min_samples=2)
+        labels = dbscan.fit_predict(midpoints)
+        cluster_counter = collections.Counter(labels)
+        largest_cluster_label = max(cluster_counter, key=cluster_counter.get)
+        largest_cluster_points = midpoints[labels == largest_cluster_label]
+
+        # Average of these midpoints must be the center
+        center = np.mean(largest_cluster_points, axis=0)
+
+    return spots, center
+
+
+def adaptive_fourier_filter(dset: sidpy.Dataset,
+                            spots: np.ndarray,
+                            low_pass: float = 3,
+                            reflection_radius: float = 0.3) -> sidpy.Dataset:
+    """
+    Use spots in diffractogram for a Fourier Filter
+
+    Parameters:
+    -----------
+    dset: sidpy.Dataset
+        image to be filtered
+    spots: np.ndarray(N,2)
+        sorted spots in diffractogram in 1/nm
+    low_pass:  float
+        low pass filter in center of diffractogram in 1/nm
+    reflection_radius:  float
+        radius of masked reflections in 1/nm
+
+    Output:
+    -------
+            Fourier filtered image
+    """
+
+    fft_transform = fourier_transform(dset)
+
+    # prepare mask
+    x, y = np.meshgrid(fft_transform.v.values, fft_transform.u.values)
+    mask = np.zeros(dset.shape)
+
+    # mask reflections
+    for spot in spots:
+        mask_spot = (x - spot[1]) ** 2 + (y - spot[0]) ** 2 < reflection_radius ** 2  # make a spot
+        mask = mask + mask_spot  # add spot to mask
+
+    # mask zero region larger (low-pass filter = intensity variations)
+    mask_spot = x ** 2 + y ** 2 < low_pass ** 2
+    mask = mask + mask_spot
+    mask[np.where(mask > 1)] = 1
+    fft_filtered = np.array(fft_transform * mask)
+
+    filtered_image = dset.like_data(np.fft.ifft2(np.fft.fftshift(fft_filtered)).real)
+    filtered_image.title = 'Fourier filtered ' + dset.title
+    filtered_image.source = dset.title
+    filtered_image.metadata = {'analysis': 'adaptive fourier filtered', 'spots': spots,
+                               'low_pass': low_pass, 'reflection_radius': reflection_radius}
+    return filtered_image
+
+
+def rotational_symmetry_diffractogram(spots: np.ndarray) -> list[int]:
+    """ Test rotational symmetry of diffraction spots"""
+
+    rotation_symmetry = []
+    for n in [2, 3, 4, 6]:
+        cc = np.array(
+            [[np.cos(2 * np.pi / n), np.sin(2 * np.pi / n), 0],
+             [-np.sin(2 * np.pi / n), np.cos(2 * np.pi / n), 0],
+             [0, 0, 1]])
+        sym_spots = np.dot(spots, cc)
+        dif = []
+        for p0, p1 in itertools.product(sym_spots[:, 0:2], spots[:, 0:2]):
+            dif.append(np.linalg.norm(p0 - p1))
+        dif = np.array(sorted(dif))
+
+        if dif[int(spots.shape[0] * .7)] < 0.2:
+            rotation_symmetry.append(n)
+    return rotation_symmetry