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

build/lib/pyTEMlib/image/image_registration.py

@@ -0,0 +1,316 @@
+"""
+image_registration.py
+by Gerd Duscher, UTK
+part of pycroscopy.image
+MIT license except where stated differently
+"""
+
+import typing
+
+import numpy as np
+import skimage
+import scipy
+
+from tqdm.auto import trange
+
+import sidpy
+_SIMPLEITK_PRESENT = True
+try:
+    import SimpleITK
+except ModuleNotFoundError:
+    _SIMPLEITK_PRESENT = False
+if not _SIMPLEITK_PRESENT:
+    print('SimpleITK not installed; Registration Functions for Image Stacks not available')
+
+#####################################################
+# Registration Functions
+#####################################################
+
+def complete_registration(main_dataset: sidpy.Dataset) -> typing.Tuple[sidpy.Dataset,
+                                                                       sidpy.Dataset]:
+    """Rigid and then non-rigid (demon) registration
+
+    Performs rigid and then non-rigid registration, please see individual functions:
+    - rigid_registration
+    - demon_registration
+
+    Parameters
+    ----------
+    main_dataset: sidpy.Dataset
+        dataset of data_type 'IMAGE_STACK' to be registered
+
+    Returns
+    -------
+    non_rigid_registered: sidpy.Dataset
+    rigid_registered_dataset: sidpy.Dataset
+
+    """
+
+    if main_dataset.data_type.name != 'IMAGE_STACK':
+        raise TypeError('Registration makes only sense for an image stack')
+
+    rigid_registered_dataset = rigid_registration(main_dataset)
+
+    rigid_registered_dataset.data_type = 'IMAGE_STACK'
+
+    non_rigid_registered = demon_registration(rigid_registered_dataset)
+    return non_rigid_registered, rigid_registered_dataset
+
+
+def demon_registration(dataset: sidpy.Dataset, verbose: bool=False) -> sidpy.Dataset:
+    """
+    Diffeomorphic Demon Non-Rigid Registration
+
+    Depends on:
+        simpleITK and numpy
+    Please Cite: http://www.simpleitk.org/SimpleITK/project/parti.html
+    and T. Vercauteren, X. Pennec, A. Perchant and N. Ayache
+    Diffeomorphic Demons Using ITK\'s Finite Difference Solver Hierarchy
+    The Insight Journal, http://hdl.handle.net/1926/510 2007
+
+    Parameters
+    ----------
+    dataset: sidpy.Dataset
+        stack of image after rigid registration and cropping
+    verbose: boolean
+        optional for increased output
+    Returns
+    -------
+        dem_reg: stack of images with non-rigid registration
+
+    Example
+    -------
+    dem_reg = demon_reg(stack_dataset, verbose=False)
+    """
+
+    if dataset.data_type.name != 'IMAGE_STACK':
+        raise TypeError('Registration makes only sense for an image stack')
+
+    dem_reg = np.zeros(dataset.shape)
+    nimages = dataset.shape[0]
+    if verbose:
+        print(nimages)
+    # create fixed image by summing over rigid registration
+
+    fixed_np = np.average(np.array(dataset), axis=0)
+
+    if not _SIMPLEITK_PRESENT:
+        print('This feature is not available:')
+        print('Please install simpleITK with: conda install simpleitk -c simpleitk')
+
+    fixed = SimpleITK.GetImageFromArray(fixed_np)
+    fixed = SimpleITK.DiscreteGaussian(fixed, 2.0)
+
+    demons = SimpleITK.DiffeomorphicDemonsRegistrationFilter()
+
+    demons.SetNumberOfIterations(200)
+    demons.SetStandardDeviations(1.0)
+
+    resampler = SimpleITK.ResampleImageFilter()
+    resampler.SetReferenceImage(fixed)
+    resampler.SetInterpolator(SimpleITK.sitkBSpline)
+    resampler.SetDefaultPixelValue(0)
+
+    for i in trange(nimages):
+        moving = SimpleITK.GetImageFromArray(dataset[i])
+        moving_f = SimpleITK.DiscreteGaussian(moving, 2.0)
+        displacement_field = demons.Execute(fixed, moving_f)
+        out_tx = SimpleITK.DisplacementFieldTransform(displacement_field)
+        resampler.SetTransform(out_tx)
+        out = resampler.Execute(moving)
+        dem_reg[i, :, :] = SimpleITK.GetArrayFromImage(out)
+
+    print(':-)')
+    print('You have successfully completed Diffeomorphic Demons Registration')
+
+    demon_registered = dataset.like_data(dem_reg)
+    demon_registered.title = 'Non-Rigid Registration'
+    demon_registered.source = dataset.title
+
+    demon_registered.metadata =dataset.metadata.copy()
+    if 'analysis' not in demon_registered.metadata:
+        demon_registered.metadata['analysis'] = {}
+    demon_registered.metadata['analysis']['non_rigid_demon_registration'] = {'package': 'simpleITK',
+                                                                             'method': 'DiscreteGaussian',
+                                                                             'variance': 2,
+                                                                             'input_dataset': dataset.source}
+    demon_registered.data_type = 'IMAGE_STACK'
+    return demon_registered
+
+
+# ##############################
+# Rigid Registration New 05/09/2024
+# ##############################
+def rigid_registration(dataset: sidpy.Dataset, normalization: typing.Optional[str] = None) -> sidpy.Dataset:
+    """
+    Rigid registration of image stack with pixel accuracy
+
+    Uses simple cross_correlation
+    (we determine drift from one image to next)
+
+    Parameters
+    ----------
+    dataset: sidpy.Dataset
+        sidpy dataset with image_stack dataset
+    normalization: str or None
+        if 'phase' then phase cross correlation is used, otherwise
+        normalized cross correlation is used
+
+    Returns
+    -------
+    rigid_registered: sidpy.Dataset
+        Registered Stack and drift (with respect to center image)
+    """
+
+    if dataset.data_type.name != 'IMAGE_STACK':
+        raise TypeError('Registration makes only sense for an image stack')
+
+    if isinstance (normalization, str):
+        if normalization.lower() != 'phase':
+            normalization = None
+    else:
+        normalization = None
+    image_dimensions = dataset.get_image_dims(return_axis=True)
+    if dataset.get_dimensions_by_type('TEMPORAL')[0] != 0:
+        x = image_dimensions[0]
+        y = image_dimensions[1]
+        z = dataset.get_dimensions_by_type('TEMPORAL', return_axis=True)[0]
+        metadata = dataset.metadata.copy()
+        original_metadata = dataset.original_metadata.copy()
+        arr = np.rollaxis(np.array(dataset), 2, 0)
+        dataset = sidpy.Dataset.from_array(arr, title=dataset.title, data_type='IMAGE_STACK',
+                                           quantity=dataset.quantity, units=dataset.units)
+        dataset.set_dimension(0, sidpy.Dimension(z.values, name='frame', units='frame', quantity='time',
+                                                  dimension_type='temporal'))
+        dataset.set_dimension(1, x)
+        dataset.set_dimension(2, y)
+        dataset.metadata = metadata
+        dataset.original_metadata = original_metadata
+
+    stack_dim = dataset.get_dimensions_by_type('TEMPORAL', return_axis=True)[0]
+    image_dim = dataset.get_image_dims(return_axis=True)
+    if len(image_dim) != 2:
+        raise ValueError('need at least two SPATIAL dimension for an image stack')
+
+    relative_drift = [[0., 0.]]
+    im1 = np.fft.fft2(np.array(dataset[0]))
+    for i in range(1, len(stack_dim)):
+        im2 = np.fft.fft2(np.array(dataset[i]))
+        shift, error, _ = skimage.registration.phase_cross_correlation(im1, im2,
+                                                                       normalization=normalization,
+                                                                       space='fourier')
+        im1 = im2.copy()
+        relative_drift.append(shift)
+
+    rig_reg, drift = rig_reg_drift(dataset, relative_drift)
+    crop_reg, input_crop = crop_image_stack(rig_reg, drift)
+
+    rigid_registered = sidpy.Dataset.from_array(crop_reg,
+                                                title='Rigid Registration',
+                                                data_type='IMAGE_STACK',
+                                                quantity=dataset.quantity,
+                                                units=dataset.units)
+    rigid_registered.title = 'Rigid_Registration'
+    rigid_registered.source = dataset.title
+    rigid_registered.metadata['analysis'] = {'rigid_registration': {'drift': drift,
+                                 'input_crop': input_crop, 'input_shape': dataset.shape[1:]}}
+
+    if 'experiment' in dataset.metadata:
+        rigid_registered.metadata['experiment'] = dataset.metadata['experiment'].copy()
+    rigid_registered.set_dimension(0, sidpy.Dimension(np.arange(rigid_registered.shape[0]),
+                                          name='frame', units='frame', quantity='time',
+                                          dimension_type='temporal'))
+
+    array_x = image_dim[0].values[input_crop[0]:input_crop[1]]
+    rigid_registered.set_dimension(1, sidpy.Dimension(array_x, name='x',
+                                                      units='nm', quantity='Length',
+                                                      dimension_type='spatial'))
+    array_y =image_dim[1].values[input_crop[2]:input_crop[3]]
+    rigid_registered.set_dimension(2, sidpy.Dimension(array_y, name='y',
+                                                      units='nm', quantity='Length',
+                                                      dimension_type='spatial'))
+    rigid_registered.data_type = 'IMAGE_STACK'
+    return rigid_registered.rechunk({0: 'auto', 1: -1, 2: -1})
+
+
+def rig_reg_drift(dset: sidpy.Dataset,
+                  rel_drift: typing.Union[typing.List[typing.List[float]], np.ndarray]
+                  ) -> typing.Tuple[np.ndarray, np.ndarray]:
+    """ Shifting images on top of each other
+
+    Uses relative drift to shift images on top of each other,
+    with center image as reference.
+    Shifting is done with shift routine of ndimage from scipy.
+    This function is used by rigid_registration routine
+
+    Parameters
+    ----------
+    dset: sidpy.Dataset
+        dataset with image_stack
+    rel_drift:
+        relative_drift from image to image as list of [shiftx, shifty]
+
+    Returns
+    -------
+    stack: numpy array
+    drift: list of drift in pixel
+    """
+
+    frame_dim = []
+    spatial_dim = []
+    selection = []
+
+    for i, axis in dset._axes.items():
+        if axis.dimension_type.name == 'SPATIAL':
+            spatial_dim.append(i)
+            selection.append(slice(None))
+        else:
+            frame_dim.append(i)
+            selection.append(slice(0, 1))
+
+    if len(spatial_dim) != 2:
+        print('need two spatial dimensions')
+    if len(frame_dim) != 1:
+        print('need one frame dimensions')
+
+    rig_reg = np.zeros([dset.shape[frame_dim[0]], dset.shape[spatial_dim[0]], dset.shape[spatial_dim[1]]])
+
+    # absolute drift
+    drift = np.array(rel_drift).copy()
+
+    drift[0] = [0, 0]
+    for i in range(1, drift.shape[0]):
+        drift[i] = drift[i - 1] + rel_drift[i]
+    center_drift = drift[int(drift.shape[0] / 2)]
+    drift = drift - center_drift
+    # Shift images
+    for i in range(rig_reg.shape[0]):
+        selection[frame_dim[0]] = slice(i, i+1)
+        # Now we shift
+        rig_reg[i, :, :] = scipy.ndimage.shift(dset[tuple(selection)].squeeze().compute(),
+                                         [drift[i, 0], drift[i, 1]], order=3)
+    return rig_reg, drift
+
+
+
+def crop_image_stack(rig_reg: np.ndarray, drift: typing.Union[np.ndarray, list]
+                     ) -> typing.Tuple[np.ndarray, list[int]]:
+    """Crop images in stack according to drift
+
+    This function is used by rigid_registration routine
+
+    Parameters
+    ----------
+    rig_reg: numpy array (N,x,y)
+    drift: list (2,B)
+
+    Returns
+    -------
+    numpy array
+    """
+    xpmax = int(rig_reg.shape[1] - -np.floor(np.min(np.array(drift)[:, 0])))
+    xpmin = int(np.ceil(np.max(np.array(drift)[:, 0])))
+    ypmax = int(rig_reg.shape[1] - -np.floor(np.min(np.array(drift)[:, 1])))
+    ypmin = int(np.ceil(np.max(np.array(drift)[:, 1])))
+
+    return rig_reg[:, xpmin:xpmax, ypmin:ypmax:], [xpmin, xpmax, ypmin, ypmax]

build/lib/pyTEMlib/image/image_utilities.py

@@ -0,0 +1,309 @@
+"""
+image_utilities part of image package of pycroscopy
+
+"""
+import numpy as np
+import scipy
+from skimage.restoration import inpaint
+
+import sidpy
+
+def crop_image(dataset: sidpy.Dataset, corners: np.ndarray) -> sidpy.Dataset:
+    """
+    Crops an image according to the corners given in the format of 
+    matplotlib.widget.RectangleSelector.
+
+    Parameters
+    ----------
+    dataset: sidpy.Dataset
+        An instance of sidpy.Dataset representing the image to be cropped.
+    corners: np.ndarray
+        A 1D array of length 4 containing the corners of the rectangular region 
+        to be cropped. The order of the corners should be (x1, y1, x2, y2).
+
+    Returns
+    -------
+    sidpy.Dataset
+        A new instance of sidpy.Dataset representing the cropped image.
+
+    Raises
+    ------
+    ValueError
+        If dataset is not an instance of sidpy.Dataset or if dataset is not an image 
+        dataset. If corners parameter is not of correct shape or size.
+
+    """
+    if not isinstance(dataset, sidpy.Dataset):
+        raise ValueError('Input dataset is not an instance of sidpy.Dataset')
+    if not dataset.data_type.name == 'IMAGE':
+        raise ValueError('Only image datasets are supported at this point')
+
+    if corners.shape != (4,):
+        raise ValueError(f'Input corners parameter should have shape (4,) but got shape {corners.shape}')
+    if corners[2]-corners[0] <= 0 or corners[3]-corners[1] <= 0:
+        raise ValueError('Invalid input corners parameter')
+
+    pixel_size = np.array([dataset.x[1]-dataset.x[0], dataset.y[1]-dataset.y[0]])
+    corners /= pixel_size
+
+    selection = np.stack([np.min(corners[:2])+0.5, np.max(corners[2:])+0.5]).astype(int)
+
+    cropped_dset = dataset.like_data(dataset[selection[0, 0]:selection[1, 0],
+                                             selection[0, 1]:selection[1, 1]])
+    cropped_dset.title = 'cropped_' + dataset.title
+    cropped_dset.source = dataset.title
+    cropped_dset.metadata = {'crop_dimension': selection, 'original_dimensions': dataset.shape}
+
+    return cropped_dset
+
+
+def flatten_image(sid_dset, order=1, flatten_axis = 'row', method = 'line_fit'):
+    """
+    Flattens an image according to the method chosen. Used heavily for AFM/STM images
+
+    Parameters
+    ----------
+    dataset: sidpy.Dataset
+        An instance of sidpy.Dataset representing the image to be flattened.
+    order: integer, 
+        Optional, default = 1. Ordfor the polynomial fit.
+    flatten_axis: string, 
+        Optional, default = 'row'. Axis along which to flatten the image.
+    method: string, 
+        Optional, default = 'line_fit'. Method to use for flattening the image.
+
+    Returns
+    -------
+    sidpy.Dataset
+        A new instance of sidpy.Dataset representing the flattened image.
+    """
+    #TODO: lots of cleanup in this function required...
+    new_sid_dset = sid_dset.copy()
+    assert len(new_sid_dset._axes) == 2, "Dataset must be 2-D for this function"
+    assert new_sid_dset.data_type == sidpy.DataType.IMAGE, "Dataset must IMAGE for this function"
+    #check the spatial dimensions, flatten along each row
+    if flatten_axis == 'row':
+        num_pts = sid_dset.shape[0]  # this is hard coded, it shouldn't be
+    elif flatten_axis == 'col':
+        num_pts = sid_dset.shape[1]  # this is hard coded, but it shouldn't be
+    else:
+        raise ValueError(f"Gave flatten axis of {flatten_axis} but only 'row', 'col' are allowed")
+
+    data_flat = np.zeros(sid_dset.shape) #again this should be the spatial (2 dimensional) part only
+    print(sid_dset.shape, num_pts)
+    if method == 'line_fit':
+        for line in range(num_pts):
+            if flatten_axis=='row':
+                line_data = np.array(sid_dset[:])[line,:]
+            elif flatten_axis=='col':
+                line_data = np.array(sid_dset[:])[:,line]
+            p = np.polyfit(np.arange(len(line_data)), line_data,order)
+            lin_est = np.polyval(p,np.arange(len(line_data)))
+            new_line = line_data - lin_est
+            data_flat[line] = new_line
+    elif method == 'plane_fit':
+        #TODO: implement plane fit
+        pass
+    else:
+        raise ValueError("Gave method of {method} but only 'line_fit', 'plane_fit' are allowed")
+
+    new_sid_dset[:] = data_flat
+
+
+def rebin(im, binning=2):
+    """
+    rebin an image by the number of pixels in x and y direction given by binning
+
+    Parameter
+    ---------
+    image: numpy array in 2 dimensions or sidpy.Dataset of data_type 'Image'
+
+    Returns
+    -------
+    binned image as numpy array or sidpy.Dataset
+    """
+    if len(im.shape) == 2:
+        rebinned_image = np.array(im).reshape((im.shape[0]//binning,
+                                               binning, im.shape[1]//binning,
+                                               binning)).mean(axis=3).mean(1)
+        if isinstance(im, sidpy.Dataset):
+            rebinned_image = im.like_data(rebinned_image)
+            rebinned_image.title = 'rebinned_' + im.title
+            rebinned_image.data_type = 'image'
+            im_dims = im.get_image_dims(return_axis=True)
+
+            rebinned_image.set_dimension(0, sidpy.Dimension(np.arange(rebinned_image.shape[0])/im_dims[0].slope,
+                                                            name='x', units=im_dims[0].units,
+                                                            dimension_type=im_dims[0].dimension_type,
+                                                            quantity=im_dims[0].quantity))
+            rebinned_image.set_dimension(1, sidpy.Dimension(np.arange(rebinned_image.shape[1])/im_dims[1].slope,
+                                                            name='y', units=im_dims[1].units,
+                                                            dimension_type=im_dims[1].dimension_type,
+                                                            quantity=im_dims[1].quantity))
+            return rebinned_image
+    else:
+        raise TypeError('not a 2D image')
+
+
+def cart2pol(points):
+    """Cartesian to polar coordinate conversion
+
+    Parameters
+    ---------
+    points: float or numpy array
+        points to be converted (Nx2)
+
+    Returns
+    -------
+    rho: float or numpy array
+        distance
+    phi: float or numpy array
+        angle
+    """
+
+    rho = np.linalg.norm(points[:, 0:2], axis=1)
+    phi = np.arctan2(points[:, 1], points[:, 0])
+
+    return rho, phi
+
+
+def pol2cart(rho, phi):
+    """Polar to Cartesian coordinate conversion
+
+    Parameters
+    ----------
+    rho: float or numpy array
+        distance
+    phi: float or numpy array
+        angle
+
+    Returns
+    -------
+    x: float or numpy array
+        x coordinates of converted points(Nx2)
+    """
+
+    x = rho * np.cos(phi)
+    y = rho * np.sin(phi)
+    return x, y
+
+
+def xy2polar(points, rounding=1e-3):
+    """ Conversion from carthesian to polar coordinates
+
+    the angles and distances are sorted by r and then phi
+    The indices of this sort is also returned
+
+    Parameters
+    ----------
+    points: numpy array
+        number of points in axis 0 first two elements in axis 1 are x and y
+    rounding: int
+        optional rounding in significant digits
+
+    Returns
+    -------
+    r, phi, sorted_indices
+    """
+
+    r, phi = cart2pol(points)
+
+    r = (np.floor(r/rounding))*rounding  # Remove rounding error differences
+
+    sorted_indices = np.lexsort((phi, r))  # sort first by r and then by phi
+    r = r[sorted_indices]
+    phi = phi[sorted_indices]
+
+    return r, phi, sorted_indices
+
+
+def cartesian2polar(x, y, grid, r, t, order=3):
+    """Transform cartesian grid to polar grid
+
+    Used by warp
+    """
+
+    rr, tt = np.meshgrid(r, t)
+
+    new_x = rr*np.cos(tt)
+    new_y = rr*np.sin(tt)
+
+    ix = scipy.interpolate.interp1d(x, np.arange(len(x)))
+    iy = scipy.interpolate.interp1d(y, np.arange(len(y)))
+
+    new_ix = ix(new_x.ravel())
+    new_iy = iy(new_y.ravel())
+
+    return scipy.ndimage.map_coordinates(grid, np.array([new_ix, new_iy]),
+                                         order=order).reshape(new_x.shape)
+
+
+def warp(diff, center):
+    """Takes a diffraction pattern (as a sidpy dataset)and warps it to a polar grid"""
+
+    # Define original polar grid
+    nx = np.shape(diff)[0]
+    ny = np.shape(diff)[1]
+
+    x = np.linspace(1, nx, nx, endpoint=True)-center[0]
+    y = np.linspace(1, ny, ny, endpoint=True)-center[1]
+    z = diff
+
+    # Define new polar grid
+    nr = int(min([center[0], center[1], diff.shape[0]-center[0], diff.shape[1]-center[1]])-1)
+    nt = 360 * 3
+
+    r = np.linspace(1, nr, nr)
+    t = np.linspace(0., np.pi, nt, endpoint=False)
+
+    return cartesian2polar(x, y, z, r, t, order=3).T
+
+
+def inpaint_image(sid_dset, mask = None, channel = None):
+    """Inpaints a sparse image, given a mask.
+
+    Args:
+        sid_dset (_type_): sidpy Dataset 
+            with two dimensions being of spatial or reciprocal type
+        mask (np.ndarry) : mask [0,1] same shape as sid_dset. 
+            If providing a sidpy dataset and mask is in the metadata dict, 
+            then this entry is optional
+        channel (int): (optional) for multi-channel datasets, 
+            provide the channel to in-paint
+    """
+    if len(sid_dset.shape)==2:
+        image_data = np.array(sid_dset).squeeze()
+    elif len(sid_dset.shape)==3:
+        image_dims = []
+        selection = []
+        for dim, axis in sid_dset._axes.items():
+            if axis.dimension_type in [sidpy.DimensionType.SPATIAL, sidpy.DimensionType.RECIPROCAL]:
+                selection.append(slice(None))
+                image_dims.append(dim)
+            else:
+                if channel is None:
+                    channel=0
+                selection.append(slice(channel, channel+1))
+
+        image_data = np.array(sid_dset[tuple(selection)]).squeeze()
+    if mask is None:
+        mask_data = sid_dset.metadata["mask"]
+        mask = np.copy(mask_data)
+        mask[mask==1] = -1
+        mask[mask==0] = 1
+        mask[mask==-1] = 0
+
+    inpainted_data = inpaint.inpaint_biharmonic(image_data, mask)
+
+    #convert this into a sidpy dataset
+    data_set = sidpy.Dataset.from_array(inpainted_data, name='inpainted_image')
+    data_set.data_type = 'image'  # supported
+
+    data_set.units = sid_dset.units
+    data_set.quantity = sid_dset.quantity
+
+    data_set.set_dimension(0, sid_dset.get_dimension_by_number(image_dims[0])[0])
+    data_set.set_dimension(1, sid_dset.get_dimension_by_number(image_dims[1])[0])
+
+    data_set.metadata["mask"] = mask
+    return data_set