pyTEMlib 0.2024.2.1__py2.py3-none-any.whl → 0.2024.6.0__py2.py3-none-any.whl

pyTEMlib/image_tools.py

@@ -50,6 +50,11 @@ from sklearn.cluster import DBSCAN
 
 from collections import Counter
 
+# center diff function
+from skimage.filters import threshold_otsu, sobel
+from scipy.optimize import leastsq
+from sklearn.cluster import DBSCAN
+
 
 _SimpleITK_present = True
 try:
@@ -275,6 +280,67 @@ def diffractogram_spots(dset, spot_threshold, return_center=True, eps=0.1):
     return spots, center
 
 
+def center_diffractogram(dset, return_plot = True, histogram_factor = None, smoothing = 1, min_samples = 100):
+    try:
+        diff = np.array(dset).T.astype(np.float16)
+        diff[diff < 0] = 0
+        
+        if histogram_factor is not None:
+            hist, bins = np.histogram(np.ravel(diff), bins=256, range=(0, 1), density=True)
+            threshold = threshold_otsu(diff, hist = hist * histogram_factor)
+        else:
+            threshold = threshold_otsu(diff)
+        binary = (diff > threshold).astype(float)
+        smoothed_image = ndimage.gaussian_filter(binary, sigma=smoothing) # Smooth before edge detection
+        smooth_threshold = threshold_otsu(smoothed_image)
+        smooth_binary = (smoothed_image > smooth_threshold).astype(float)
+        # Find the edges using the Sobel operator
+        edges = sobel(smooth_binary)
+        edge_points = np.argwhere(edges)
+
+        # Use DBSCAN to cluster the edge points
+        db = DBSCAN(eps=10, min_samples=min_samples).fit(edge_points)
+        labels = db.labels_
+        if len(set(labels)) == 1:
+            raise ValueError("DBSCAN clustering resulted in only one group, check the parameters.")
+
+        # Get the largest group of edge points
+        unique, counts = np.unique(labels, return_counts=True)
+        counts = dict(zip(unique, counts))
+        largest_group = max(counts, key=counts.get)
+        edge_points = edge_points[labels == largest_group]
+
+        # Fit a circle to the diffraction ring
+        def calc_distance(c, x, y):
+            Ri = np.sqrt((x - c[0])**2 + (y - c[1])**2)
+            return Ri - Ri.mean()
+        x_m = np.mean(edge_points[:, 1])
+        y_m = np.mean(edge_points[:, 0])
+        center_guess = x_m, y_m
+        center, ier = leastsq(calc_distance, center_guess, args=(edge_points[:, 1], edge_points[:, 0]))
+        mean_radius = np.mean(calc_distance(center, edge_points[:, 1], edge_points[:, 0])) + np.sqrt((edge_points[:, 1] - center[0])**2 + (edge_points[:, 0] - center[1])**2).mean()
+    
+    finally:
+        if return_plot:
+            fig, ax = plt.subplots(1, 4, figsize=(10, 4))
+            ax[0].set_title('Diffractogram')
+            ax[0].imshow(dset.T, cmap='viridis')
+            ax[1].set_title('Otsu Binary Image')
+            ax[1].imshow(binary, cmap='gray')
+            ax[2].set_title('Smoothed Binary Image')
+            ax[2].imshow(smooth_binary, cmap='gray')
+            ax[3].set_title('Edge Detection and Fitting')
+            ax[3].imshow(edges, cmap='gray')
+            ax[3].scatter(center[0], center[1], c='r', s=10)
+            circle = plt.Circle(center, mean_radius, color='red', fill=False)
+            ax[3].add_artist(circle)
+            for axis in ax:
+                axis.axis('off')
+            fig.tight_layout()
+          
+    return center
+
+
 def adaptive_fourier_filter(dset, spots, low_pass=3, reflection_radius=0.3):
     """
     Use spots in diffractogram for a Fourier Filter
@@ -375,16 +441,10 @@ def complete_registration(main_dataset, storage_channel=None):
 
     rigid_registered_dataset = rigid_registration(main_dataset)
 
-    if storage_channel is not None:
-        registration_channel = ft.log_results(storage_channel, rigid_registered_dataset)
-
+    
     print('Non-Rigid_Registration')
 
     non_rigid_registered = demon_registration(rigid_registered_dataset)
-    if storage_channel is not None:
-        registration_channel = ft.log_results(storage_channel, non_rigid_registered)
-
-    non_rigid_registered.h5_dataset = registration_channel
     return non_rigid_registered, rigid_registered_dataset
 
 
@@ -473,7 +533,7 @@ def demon_registration(dataset, verbose=False):
 ###############################
 # Rigid Registration New 05/09/2020
 
-def rigid_registration(dataset):
+def rigid_registration(dataset, sub_pixel=True):
     """
     Rigid registration of image stack with pixel accuracy
 
@@ -529,9 +589,13 @@ def rigid_registration(dataset):
         selection[frame_dim[0]] = slice(i, i+1)
         moving = dataset[tuple(selection)].squeeze().compute()
         fft_moving = np.fft.fft2(moving)
-        image_product = fft_fixed * fft_moving.conj()
-        cc_image = np.fft.fftshift(np.fft.ifft2(image_product))
-        shift = np.array(ndimage.maximum_position(cc_image.real))-cc_image.shape[0]/2
+        if sub_pixel:
+            shift = skimage.registration.phase_cross_correlation(fft_fixed, fft_moving, upsample_factor=1000,
+                                                                 space='fourier')[0]
+        else:    
+            image_product = fft_fixed * fft_moving.conj()
+            cc_image = np.fft.fftshift(np.fft.ifft2(image_product))
+            shift = np.array(ndimage.maximum_position(cc_image.real))-cc_image.shape[0]/2
         fft_fixed = fft_moving
         relative_drift.append(shift)
     rig_reg, drift = rig_reg_drift(dataset, relative_drift)
@@ -542,9 +606,18 @@ def rigid_registration(dataset):
     rigid_registered.source = dataset.title
     rigid_registered.metadata = {'analysis': 'rigid sub-pixel registration', 'drift': drift,
                                  'input_crop': input_crop, 'input_shape': dataset.shape[1:]}
-    rigid_registered.set_dimension(0, dataset._axes[frame_dim[0]])
-    rigid_registered.set_dimension(1, dataset._axes[spatial_dim[0]][input_crop[0]:input_crop[1]])
-    rigid_registered.set_dimension(2, dataset._axes[spatial_dim[1]][input_crop[2]:input_crop[3]])
+    rigid_registered.set_dimension(0, sidpy.Dimension(np.arange(rigid_registered.shape[0]), 
+                                          name='frame', units='frame', quantity='time',
+                                          dimension_type='temporal'))
+    
+    array_x = dataset._axes[spatial_dim[0]][input_crop[0]:input_crop[1]].values
+    rigid_registered.set_dimension(1, sidpy.Dimension(array_x,
+                                          'x', units='nm', quantity='Length',
+                                          dimension_type='spatial'))
+    array_y = dataset._axes[spatial_dim[1]][input_crop[2]:input_crop[3]].values
+    rigid_registered.set_dimension(2, sidpy.Dimension(array_y,
+                                          'y', units='nm', quantity='Length',
+                                          dimension_type='spatial'))
     return rigid_registered.rechunk({0: 'auto', 1: -1, 2: -1})
 
 
@@ -589,6 +662,7 @@ def rig_reg_drift(dset, rel_drift):
     rig_reg = np.zeros([dset.shape[frame_dim[0]], dset.shape[spatial_dim[0]], dset.shape[spatial_dim[1]]])
 
     # absolute drift
+    print(rel_drift)
     drift = np.array(rel_drift).copy()
 
     drift[0] = [0, 0]
@@ -731,15 +805,27 @@ class LineSelector(matplotlib.widgets.PolygonSelector):
             self.line_verts[moved_point] = self.new_point
             self.set_linewidth()
 
-def get_profile(dataset, line):
+def get_profile(dataset, line, spline_order=-1):
+    """
+    This function extracts a line profile from a given dataset. The line profile is a representation of the data values 
+    along a specified line in the dataset. This function works for both image and spectral image data types.
+
+    Args:
+        dataset (sidpy.Dataset): The input dataset from which to extract the line profile.
+        line (list): A list specifying the line along which the profile should be extracted.
+        spline_order (int, optional): The order of the spline interpolation to use. Default is -1, which means no interpolation.
+
+    Returns:
+        profile_dataset (sidpy.Dataset): A new sidpy.Dataset containing the line profile.
+
+
+    """
     xv, yv = get_line_selection_points(line)
-    
-    
     if dataset.data_type.name == 'IMAGE':
         dataset.get_image_dims()
         xv /= (dataset.x[1] - dataset.x[0])
         yv /= (dataset.y[1] - dataset.y[0])
-        profile = scipy.ndimage.map_coordinates(np.array(dataset), [xv,yv])
+        profile = scipy.ndimage.map_coordinates(np.array(dataset), [xv, yv])
         
         profile_dataset = sidpy.Dataset.from_array(profile.sum(axis=0))
         profile_dataset.data_type='spectrum'
@@ -753,19 +839,21 @@ def get_profile(dataset, line):
 
     if dataset.data_type.name == 'SPECTRAL_IMAGE':
         spectral_axis = dataset.get_spectral_dims(return_axis=True)[0]
-        profile = np.zeros([xv.shape[1], 2, len(spectral_axis)])
-        data =np.array(dataset)
-
-        for index_x in range(xv.shape[1]):
-            for  index_y  in range(xv.shape[0]):
-                x = xv[index_y, index_x]
-                y = yv[index_y, index_x]
-                profile[index_x, 0] +=data[int(x),int(y)]
+        if spline_order > -1:
+            xv, yv, zv = get_line_selection_points_interpolated(line, z_length=dataset.shape[2])
+            profile = scipy.ndimage.map_coordinates(np.array(dataset), [xv, yv, zv], order=spline_order)
+            profile = profile.sum(axis=0)
+            profile = np.stack([profile, profile], axis=1)
+            start = xv[0, 0, 0]
+        else:
+            profile = get_line_profile(np.array(dataset), xv, yv, len(spectral_axis))
+            start = xv[0, 0]
+        print(profile.shape)
         profile_dataset = sidpy.Dataset.from_array(profile)
         profile_dataset.data_type='spectral_image'
         profile_dataset.units = dataset.units
         profile_dataset.quantity = dataset.quantity
-        profile_dataset.set_dimension(0, sidpy.Dimension(np.linspace(xv[0,0], xv[-1,-1], profile_dataset.shape[0]), 
+        profile_dataset.set_dimension(0, sidpy.Dimension(np.arange(profile_dataset.shape[0])+start, 
                                                   name='x', units=dataset.x.units, quantity=dataset.x.quantity,
                                                   dimension_type='spatial'))
         profile_dataset.set_dimension(1, sidpy.Dimension([0, 1], 
@@ -776,6 +864,42 @@ def get_profile(dataset, line):
     return profile_dataset
 
 
+
+def get_line_selection_points_interpolated(line, z_length=1):
+    
+    start_point = line.line_verts[3]
+    right_point = line.line_verts[0]
+    low_point = line.line_verts[2]
+
+    if start_point[0] > right_point[0]:
+        start_point = line.line_verts[0]
+        right_point = line.line_verts[3]
+        low_point = line.line_verts[1]
+    m = (right_point[1] - start_point[1]) / (right_point[0] - start_point[0])
+    length_x = int(abs(start_point[0]-right_point[0]))
+    length_v = int(np.linalg.norm(start_point-right_point))
+    
+    linewidth = int(abs(start_point[1]-low_point[1]))
+    x = np.linspace(0,length_x, length_v)
+    y = np.linspace(0,linewidth, line.line_width)
+    if z_length > 1:
+        z = np.linspace(0, z_length, z_length)
+        xv, yv, zv = np.meshgrid(x, y, np.arange(z_length))
+        x = np.atleast_2d(x).repeat(z_length, axis=0).T
+        y = np.atleast_2d(y).repeat(z_length, axis=0).T
+    else:
+        xv, yv = np.meshgrid(x, y)
+    
+    
+    yv = yv + x*m + start_point[1] 
+    xv = (xv.swapaxes(0,1) -y*m ).swapaxes(0,1) + start_point[0]
+
+    if z_length > 1:
+        return xv, yv, zv
+    else:   
+        return xv, yv
+
+
 def get_line_selection_points(line):
     
     start_point = line.line_verts[3]
@@ -801,6 +925,16 @@ def get_line_selection_points(line):
     return xx, yy
 
 
+def get_line_profile(data, xv, yv, z_length):
+    profile = np.zeros([len(xv[0]), 2, z_length])
+    for index_x in range(xv.shape[1]):
+        for  index_y  in range(xv.shape[0]):
+            x = int(xv[index_y, index_x])
+            y = int(yv[index_y, index_x])
+            if x< data.shape[0] and x>0 and y < data.shape[1] and y>0:
+                profile[index_x, 0] +=data[x, y]
+    return profile
+     
 
 def histogram_plot(image_tags):
     """interactive histogram"""
@@ -1010,9 +1144,8 @@ def cartesian2polar(x, y, grid, r, t, order=3):
     return ndimage.map_coordinates(grid, np.array([new_ix, new_iy]), order=order).reshape(new_x.shape)
 
 
-def warp(diff):
-    """Takes a centered diffraction pattern (as a sidpy dataset)and warps it to a polar grid"""
-    """Centered diff can be produced with it.diffractogram_spots(return_center = True)"""
+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]
@@ -1020,20 +1153,19 @@ def warp(diff):
 
     # Define center pixel
     pix2nm = np.gradient(diff.u.values)[0]
-    center_pixel = [abs(min(diff.u.values)), abs(min(diff.v.values))]//pix2nm
 
-    x = np.linspace(1, nx, nx, endpoint=True)-center_pixel[0]
-    y = np.linspace(1, ny, ny, endpoint=True)-center_pixel[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_pixel[0], center_pixel[1], diff.shape[0]-center_pixel[0], diff.shape[1]-center_pixel[1]])-1)
-    nt = 360*3
+    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)
+    return cartesian2polar(x, y, z, r, t, order=3).T
 
 
 def calculate_ctf(wavelength, cs, defocus, k):