Live radial profile Diff or FFT

'''
Code to compute radial-max profiles of either: 
    The FFT of an HRTEM image 
    Diffraction pattern
Profiles can be computed from either: 
    Camera live-view image in DigitalMicrograph
    In-situ camera data played back via the In-Situ Player
The result is a 2D display of these profiles over time
Since the data is computed as fast as possible, which could vary over time, 
the time axis in the resulting image may be somewhat non-linear, so it 
is not calibrated. 

WARNING: Due to a bug in GMS 3.5.0 and 3.5.1, this script will not run in those versions

Run the code with the live-view image, containing a rectangular ROI, front-most in GMS
    For diffraction data, be sure to center the ROI precisely over the diffraction pattern center.
    The code also works with an IS video played back with the IS player
    To stop calculation, delete the ROI. 
Profiles are computed as often as possible.

Lines of code between #XXXXXXXX... lines are specific to computing a radial-max profile
    All other lines of code are general, and can be re-used to produce other kinds of profiles from a live-view image
Code written by Ben Miller. Last Updated Oct 2024
'''
import time            
import numpy as np
import traceback
if not DM.IsScriptOnMainThread(): print('Scipy scripts cannot be run on Background Thread.'); exit()
import scipy
from scipy import ndimage
from scipy import signal
from scipy import fftpack
from scipy.ndimage.interpolation import geometric_transform
import sys 
from tkinter import *
sys.argv.extend(['-a', ' '])
from numpy.lib.stride_tricks import as_strided

#User editable variables are set here
#XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
Diff_not_FFT = False #Set to True to calculate profiles from a live diffraction pattern, rather than from the FFT of an HRTEM image.
median = 0 #set to 1 to apply median filter to the FFT prior to profile creation (slows calculation, especially for large input images)
initial_result_image_width = 200 #how many profiles can be displayed in the intial result window (window is automatically expanded as needed)
profile_result_length_ratio = 2 #Set to some integer 2^N, N=>0. Smaller N will make calculation slower. Default: 4
profile_angular_sampling_resolution = 256 #set how many samples are taken around the circumference of the radial profile
mask_center_lines = True #Set to True to mask the vertical and horizontal center lines (often useful for FFTs)
ask_about_diffraction = True #Set to True to always ask to confirm whether to process diffraction patterns or HRTEM images
FFT_bin = 2 #Set to an integer number >=1 to bin the diffraction pattern or diffractogram prior to computing the radial profile. Smaller N will make the calculation slower. 
#XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX


#XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
change = True
if ask_about_diffraction: 
    if Diff_not_FFT: 
        message = "Do you want to process a series of diffraction patterns? \nSelect OK to continue or cancel to instead process HRTEM images."
    else:    
        message = "Do you want to process a series of HRTEM images? \nSelect OK to continue or cancel to instead process diffraction patterns."
    change = not DM.OkCancelDialog(message)
if (change and ask_about_diffraction): Diff_not_FFT = not Diff_not_FFT 
#XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX

def DM_interaction_wait_dialog(instructions_str, title_str='DM_Interaction_Dialog',button1_str='Continue',button2_str='Cancel' ):
    '''
    Function to Create a tkinter dialog that enables the user to interact with GMS
    while the script is paused, and then resume the script after they are done
    Accepts:
        instructions_str    string giving instructions to the user about what actions are expected in GMS
        title_str            optional string for the dialog window title
        button1_str            optional string to change the continue button label text
        button1_str            optional string to change the cancel button label text
    Returns: Nothing
    '''
    root = Tk()
    root.attributes('-topmost', 'true')
    root.title(title_str)
    def go_on():
        root.destroy()
    def cancel():
        print("Script Aborted by User")
        root.destroy()
        sys.exit()
    Label(root, text=instructions_str).grid(row=0, sticky=W)
    Button(root, text=button1_str, command=go_on).grid(column=0, row=1, sticky=W, pady=4)
    Button(root, text=button2_str, command=cancel).grid(column=1, row=1, sticky=W, pady=4)
    mainloop()
    
    

def strided_binning2D(array,binning=(1,1)):
    """
    Function to Bin 2D Data 
    Accepts:
        array        2D numpy array to be binned
        binning        2 element tuple of integer binning amounts (bin_x, bin_y)
    Returns: A binned 2D array
    """
    bin=np.flip(binning)
    nh = (array.shape[0]-bin[0])//bin[0]+1
    nw = (array.shape[1]-bin[1])//bin[1]+1
    strides = (array.strides[0],array.strides[1],array.strides[0]*bin[0],array.strides[1]*bin[1])
    strides = (array.strides[0],array.strides[1],array.strides[0]*bin[0],array.strides[1]*bin[1])
    shape = (bin[0],bin[1],nh,nw)
    virtual_datacube = as_strided(array,shape=shape,strides=strides)
    result = np.sum(virtual_datacube,axis=(0,1))
    return result
    
class CListen(DM.Py_ScriptObject):
        #Function to find an ROI placed on an image by the user, returning the ROI ID. 
    #If no ROI found, create a new one covering the entire image.
    def find_ROI(self,image):
        imageDisplay = image.GetImageDisplay(0)
        numROIs = imageDisplay.CountROIs()
        id = None
        for n in range(numROIs):
            roi = imageDisplay.GetROI(0) 
            if roi.IsRectangle():
                roi.SetVolatile(False)
                roi.SetResizable(False)  
                id = roi.GetID()
                break
        if id is None:
            #If No ROI is found, create one that covers the whole image. 
            print("\nRectangular ROI not found... using whole image")
            data_shape = image.GetNumArray().shape
            roi=DM.NewROI()
            roi.SetRectangle(0, 0, data_shape[0], data_shape[1])
            imageDisplay.AddROI(roi)
            roi.SetVolatile(False)
            roi.SetResizable(False) 
            id = roi.GetID()
        return id
    
    #XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX    
    #Function to Determine Result Profile Length
    def profile_length(self,sx,sf):
        return int(sx/sf/FFT_bin)
    #Function to Set Result Image Calibration
    def calibrate_result(self,scale,sx,sf,ustr,diff_not_fft=False):
        if diff_not_fft:
            diff_scale = scale*sf*FFT_bin
            unit_str = unit_str
        else:
            diff_scale = FFT_bin*sf/scale/sx/2
            unit_str = ustr+"-1"
        return(diff_scale,unit_str)
    #Funtion to convert cartesian-coordinate image to polar-coordinate image
    def topolar(self, img,  r_size, theta_size, order=1):
        sx, sy = img.shape
        max_radius = int(sx/2)
        #define transform
        def transform(coords):
            theta = 2.0*np.pi*coords[1] / (theta_size - 1.)
            radius = max_radius * coords[0] / r_size
            i = int(sx/2) - radius*np.sin(theta)
            j = radius*np.cos(theta) + int(sx/2)
            return i,j
        #perform transform
        polar = geometric_transform(img, transform, output_shape=(r_size,theta_size), order=order,mode='constant',cval=1.0,prefilter=False)    
        return polar
        
        #Function to calculate radial profile of a diffraction pattern or of FFT from image
    def radial_profile(self, image_o, profile_ang_res,length_ratio,do_median, FFT=False):    
        if FFT: image_o = np.absolute(scipy.fftpack.fftshift(np.fft.fft2(image_o)))
        #determine profile size
        sx, sy = image_o.shape
        if mask_center_lines: 
            image_o[:,sy//2-1:sy//2-1+1] = 0
            image_o[sx//2-1:sx//2-1+1,:] = 0
        image_o = strided_binning2D(image_o,binning=(FFT_bin,FFT_bin))
        sx, sy = image_o.shape
        #Median-Filter FFT to remove single-pixel outliers
        if do_median: image_o_median = scipy.ndimage.median_filter(image_o, size=3)
        else: image_o_median = image_o
        if mask_center_lines: 
            image_o_median[:,sx//2] = 0
            image_o_median[sx//2,:] = 0
        profile_size = int(sx/length_ratio)
        #convert FFT image to polar coordinates
        polar_im = self.topolar(image_o_median, profile_size, profile_ang_res, order=1)
        #compute radial mean and maximum profiles
        radial_max=np.amax(polar_im,1)
        radial_mean=np.mean(polar_im,1)
        #median-filter the radial mean profile to smooth this further
        radial_mean_median = scipy.signal.medfilt(radial_mean)
        #radial profile is radial-max minus radial-mean
        radial_profile = np.atleast_2d(radial_max-radial_mean_median)    
        return radial_profile
        
    #XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX  
    
    #Initialization Function
    def __init__(self, img):
        try:
            #Create an index that is incremented each time data is processed. (This index is not utilized in this script.)
            self.i = 0
            self.name = "ImageO"
            #get the original image and assign it to self.imgref
            self.imgref = img
            #Get the data from the region within an ROI                                     
            self.roi = DM.GetROIFromID(self.find_ROI(self.imgref))
            if Diff_not_FFT:
                self.roi.SetVolatile(True)
                self.roi.SetResizable(True) 
                DM_interaction_wait_dialog("Move the ROI in your dataset so it is exactly centered on the diffraction pattern center, and includes as much of the pattern as you want to analyze.", title_str='DM_Interaction_Dialog',button1_str='Continue',button2_str='Cancel' )
                self.roi = DM.GetROIFromID(self.find_ROI(self.imgref))
            val, val2, val3, val4 = self.roi.GetRectangle()
            self.data = self.imgref.GetNumArray()[int(val):int(val3),int(val2):int(val4)]
            #get the shape and calibration of the original image
            (input_sizex, input_sizey) = self.data.shape
            origin, x_scale_orig, scale_unit_orig = self.imgref.GetDimensionCalibration(1, 0)
            if scale_unit_orig == b'\xb5m': scale_unit_orig = 'um' #scale unit of microns causes problems for python in DM
            #XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
            #Set the size and calibration of the result image
            (self.scale,self.unit_string) = self.calibrate_result(x_scale_orig,input_sizex,profile_result_length_ratio,scale_unit_orig)
            r_img_size=self.profile_length(input_sizex,profile_result_length_ratio)
            #XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
            #create empty set for result images
            self.result_images = {}
            #create 1st result image and set calibration
            self.result_images[self.name] = DM.CreateImage(np.zeros((r_img_size,initial_result_image_width)))
            self.result_images[self.name].SetDimensionCalibration(1,0,self.scale,self.unit_string,0)
            self.result_images[self.name].ShowImage()       
            #get numpy array from result image
            self.result_array = self.result_images[self.name].GetNumArray()
            #Set the image name which will be displayed in the image window's title bar
            self.result_images[self.name].SetName("Radial Max Profiles of "+img.GetName())
            DM.Py_ScriptObject.__init__(self)
            self.stop = 0
        except: print(traceback.format_exc())
    
    #Function to end Image Listener
    def __del__(self):
        print("Script Ended")
        DM.Py_ScriptObject.__del__(self)
        
    #This function is run each time the image changes
    def HandleDataChangedEvent(self, flags, image):
        try:
            if not self.stop:
                #start timing
                start=time.perf_counter()
                (result_sizey, result_sizex) = self.result_array.shape
                #if the result image is nearly full, make it 2x larger
                if self.i > result_sizex-2:
                    #create a new numpy array 2x larger
                    self.result_array_temp = np.append(self.result_array, np.zeros_like(self.result_array), axis=1)
                    #close the old results image in DM
                    DM.DeleteImage(self.result_images[self.name])
                    #create a new results image and calibrate it
                    self.name="Image{0}".format(self.i)
                    self.result_images[self.name] = DM.CreateImage(np.copy((self.result_array_temp)))
                    self.result_images[self.name].SetDimensionCalibration(1,0,self.scale,self.unit_string,0)
                    #Set the image name which will be displayed in the image window's title bar
                    self.result_images[self.name].SetName("Radial Max Profiles of "+self.imgref.GetName())
                    #display new result image in DM
                    self.result_images[self.name].ShowImage()    
                    #get numpy array from new result image
                    self.result_array = (self.result_images[self.name].GetNumArray())
                #XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
                #Get an (updated) ROI position                                                                     
                val, val2, val3, val4 = self.roi.GetRectangle()
                #Get the data from the ROI area as a numpy array                           
                self.data = self.imgref.GetNumArray()[int(val):int(val3),int(val2):int(val4)]
                #compute radial profile from the diffraction pattern or radial FFT profile from the image, and place this profile into results image
                self.result_array[:,self.i] = self.radial_profile(self.data,profile_angular_sampling_resolution,profile_result_length_ratio, median, FFT=(not Diff_not_FFT))
                #XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
                #Update the result image display
                self.result_images[self.name].UpdateImage()
                #end timing and output time to process this frame
                end=time.perf_counter()
                print("Processed Image "+str(self.i)+"  Processing Time= "+str(end-start))
                #Increment an index each time data is processed. 
                self.i = self.i+1    
        except: print(traceback.format_exc())
        
    #Function to end script if source image window is closed
    def HandleWindowClosedEvent(self, event_flags, window):
        if not self.stop:
            self.stop = 1
            global listener
            print("Window Closed")
            DM.DoEvents()
            #Unregister all listeners
            listener.UnregisterAllListeners()
            print("Live Processing Script Ended")
   
            
    #Function to end script if the ROI is deleted
    def HandleROIRemovedEvent(self, img_disp_event_flags, img_disp, roi_change_flag, roi_disp_change_flags, roi):
        if not self.stop:
            self.stop = 1
            global listener
            print("ROI Removed")
            DM.DoEvents()
            #Unregister all listeners
            listener.UnregisterAllListeners()
            print("Live Processing Script Ended")
            
#Function to get the currently used version of DigitalMicrograph 
def get_DM_version():
    #No Python script command exists to get the DM version, 
    #so we first run a DM script to put the values in the global tags
    dm_script = ('number minor, major, bugVersion\n'
        'GetApplicationVersion(major, minor, bugVersion)\n'
        'GetPersistentTagGroup().TagGroupSetTagAsLong("Python_Temp:DM_Version_Major",major)\n'
        'GetPersistentTagGroup().TagGroupSetTagAsLong("Python_Temp:DM_Version_Minor",minor)\n'
        'GetPersistentTagGroup().TagGroupSetTagAsLong("Python_Temp:DM_Version_bugVersion",bugVersion)')
    DM.ExecuteScriptString(dm_script)
    #Now get the information stored in the global tags by the DM script
    version = [0,0,0]
    _,version[0] = DM.GetPersistentTagGroup().GetTagAsString("Python_Temp:DM_Version_Major")
    _,version[1] = DM.GetPersistentTagGroup().GetTagAsString("Python_Temp:DM_Version_Minor")
    _,version[2] = DM.GetPersistentTagGroup().GetTagAsString("Python_Temp:DM_Version_bugVersion")
    return version  
          
#Main Code Starts Here

#Check that we are not running 3.5.0 or 3.5.1 which have a known bug affecting this script.
if ((get_DM_version()[1] == '51') or (get_DM_version()[1] == '50')):
    DM.OkDialog("Due to a bug in DigitalMicrograph 3.5.0 and 3.5.1, this script would cause DM to crash in those versions. \n\nScript Aborted.")
    exit()

#Get front image in GMS        
img1 = DM.GetFrontImage()
#Get the image window, so we can check if it gets closed
imageDoc = DM.GetFrontImageDocument()
imDocWin = imageDoc.GetWindow()
#Get the image display, for the ROI-removed listener
imageDisplay = img1.GetImageDisplay(0)

#Listeners are started here
#initiate the image listener
listener = CListen(img1)
#check if the source window closes
WindowClosedListenerID = listener.WindowHandleWindowClosedEvent(imDocWin, 'pythonplugin')
#check if the ROI has been deleted
ROIRemovedListenerID = listener.ImageDisplayHandleROIRemovedEvent(imageDisplay,'pythonplugin')
#check if the source image changes
DataChangedListenerID = listener.ImageHandleDataChangedEvent(img1, 'pythonplugin')
#IDs are not used in this script, but could be used to unregister individual listeners in DM 3.5.2 and higher.