Gatan | AMETEKSkip to Main Content
No options found

Live max. FFT

Python Script

Code to process the live-view image in DigitalMicrograph, producing an FFT-like display. A grid of small FFTs is computed from the image, and then rather than averaging the FFTs together, the maximum across all FFTs is taken pixel-by-pixel.

 

Preview

'''
Code to process the live-view image in GMS, producing an FFT-like display.
A grid of small FFTs is computed from the image, and then rather than averaging the FFTs together, 
the maximum across all FFTs is taken pixel-by-pixel

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
    The code also works with an IS video played back with the IS player
    To stop calculation, delete the ROI. 
FFTs are computed as often as possible.
Lines of code between #XXXXXXXX... lines are specific to computing a max FFT 
    All other lines of code are general, and can be re-used to produce other kinds of processed images from a live-view image
#Code written by Ben Miller. Last Updated Apr 2022
'''
import time 
import numpy as np
from numpy.lib.stride_tricks import as_strided
import traceback

#User editable variables are set here
#XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
FFTsize = 256 #FFT size. (Default 256) Smaller FFT will result in more noise and lower FFT resolution, larger FFT in fewer visible peaks.
spacing = 256 #FFT spacing. (Default 256) Choosing a spacing larger than the FFT size will result in regions of the image being excluded. Choosing smaller spacings will result in a longer time to compute each frame.
print_timing = True      # (Default True) Select whether to output the time it takes to compute each frame
#XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX

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(n) 
            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 process the data (numpy array) from the ROI area
    def ROI_process(self,image_data):
        #Create 4D Diffractogram Datacube (Like 4D STEM datacube)
        (w0,h0) = image_data.shape
        nw = (w0-FFTsize)//spacing+1
        nh = (h0-FFTsize)//spacing+1
        #Create (Virtual) 4D Datacube of Image Regions 
        shape = (FFTsize,FFTsize,nw,nh)
        strides = (image_data.strides[0],image_data.strides[1],image_data.strides[0]*spacing,image_data.strides[1]*spacing)
        image_datacube = np.transpose(as_strided(image_data,shape=shape,strides=strides),(2,3,0,1))
        #Create Hanning Window
        hanningf = np.hanning(FFTsize)
        hanningWindow2d = np.sqrt(np.outer(hanningf, hanningf)).astype('float32')
        #Compute FFTs
        if print_timing: print("Computing %s FFTs..." %(nw*nh))
        start_timef=time.perf_counter()
        datacube = np.log(np.fft.fftshift(np.abs(np.fft.fft2(hanningWindow2d*image_datacube))**2, axes=(2,3))).astype('float32')
        max_FFT = np.max(datacube, axis=(0,1))
        return max_FFT
    #XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX

    #Initialization Function
    def __init__(self,img):
        try:
            #Create an index that is incremented each time data is processed. 
            self.i = 0
            #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))
            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
            x_origin = self.imgref.GetDimensionOrigin(0); y_origin = self.imgref.GetDimensionOrigin(1)
            x_scale = self.imgref.GetDimensionScale(0);   y_scale = self.imgref.GetDimensionScale(1)
            #scale unit of microns causes problems for python in DM
            try: 
                x_unit = self.imgref.GetDimensionUnitString(0) ; y_unit = self.imgref.GetDimensionUnitString(1)
            except: 
                x_unit = self.imgref.GetDimensionUnitInfo(0)[0]; y_unit = self.imgref.GetDimensionUnitInfo(1)[0]
                if x_unit == 'micrometer': x_unit = 'um'        
                if y_unit == 'micrometer': y_unit = 'um'
            #XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
            #Create a new image to contain the results of processing
            self.result_image = DM.CreateImage(self.ROI_process(self.data.transpose()))
            #Set the calibration of the FFT result image (reciprocal space)
            binning = 1 # this is just here so I can use the equation below for the FFT calibration... it is not really used in this script
            self.result_image.SetDimensionCalibration(0,-1/(2*(x_scale*binning)),1/(x_scale*binning)/FFTsize,x_unit+"-1",0)
            self.result_image.SetDimensionCalibration(1,-1/(2*(y_scale*binning)),1/(y_scale*binning)/FFTsize,y_unit+"-1",0)
            #XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
            #Copy tags from original to processed image
            Tag_Copy(self.imgref, self.result_image,subPath='OriginalImageTags')
            #Get the numpy array of the result image so I can directly change the data values later
            self.result_data=self.result_image.GetNumArray()
            #Display the result image in GMS
            self.result_image.ShowImage()
            #Set the image name which will be displayed in the image window's title bar
            self.result_image.SetName("FFT Max of "+img.GetName())
            DM.Py_ScriptObject.__init__(self)
            self.stop = 0
        except: print(traceback.format_exc())
            
    #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()                                  
                #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)]
                #Process the data and place in the result array (note that the [:] here is essential)
                self.result_data[:]=self.ROI_process(self.data)
                #XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
                #Update the result image display
                self.result_image.UpdateImage()
                #end timing and output time to process this frame
                end=time.perf_counter()
                if print_timing: 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 Delete Image Listener
    def __del__(self):
        print("Listener Deleted")
        DM.Py_ScriptObject.__del__(self)
        
    #Function to end processing by deleting or unregistering listener
    def RemoveListeners(self):
        try: 
            if not self.stop:
                self.stop = 1
                DM.DoEvents()
                global listener
                #DM 3.5.2 and higher have new function for unregistering listeners. 
                #DM 3.4.3 and lower should delete the listener instead
                #DM 3.5.0 and 3.5.1 have a fatal flaw regarding listeners, so this script is not compatible with those versions
                if (get_DM_version()[1][0] == "4" and get_DM_version()[0] == "3"):
                    del listener
                else:
                    listener.UnregisterAllListeners()
                print("Live Processing Script Ended")
        except: print(traceback.format_exc())
        
    #Remove listeners if source image window is closed
    def HandleWindowClosedEvent(self, event_flags, window):
        print("Window Closed")
        self.RemoveListeners()
                
    #Remove listeners if the ROI is deleted
    def HandleROIRemovedEvent(self, img_disp_event_flags, img_disp, roi_change_flag, roi_disp_change_flags, roi):
        print("ROI Removed")
        self.RemoveListeners()
        
#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  
          
def Tag_Copy(image_source, image_dest, subPath = None ):
    '''
    Copy all tags between source and destination.
    If no destination subPath is provided, the destination tags will be replaced.
    '''
    #Copy Tags
    tg_source = image_source.GetTagGroup()
    tg_dest = image_dest.GetTagGroup()
    if ( subPath != None ):
        tg_dest.SetTagAsTagGroup(subPath,tg_source.Clone())
    else:
        tg_dest.DeleteAllTags()
        tg_dest.CopyTagsFrom(tg_source.Clone())

#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')) and get_DM_version()[0] == "3"):
    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.