HOZ/HOZfocus/autofocus.py

"""
Special autofocus function based on radia power density function with
radial exponential distribution fit

Luyang Han
2019.08.01  Inital version
2020.05.18  Add auto stig function
"""

import numpy as np
import tempfile, time
from imageio import imread, imsave
from numpy import fft
from scipy.optimize import curve_fit

from skimage import filters

def dif_of_gaussian(img, sigma = 10):
    # return difference of gaussian image
    # image will be automatically normalized.
    img = img / img.max()
    img_blur = filters.gaussian(img, sigma)
    return img - img_blur

def _func(x,I0,T,N):
    return I0*T**2*x/(1+x**2*T**2)**1.5+N*x

def sharpness_nv_g(img):
    # a sharpness definition of normalized variance
    img = dif_of_gaussian(img)
    mu = np.average(img)
    w,h = img.shape
    #sharp = np.sum((img-mu)**2)/w/h/mu
    sharp = np.sum((img-mu)**2)/w/h
    return sharp

def sharpness_nv(img):
    # a sharpness definition of normalized variance
    #img = dif_of_gaussian(img)
    mu = np.average(img)
    w,h = img.shape
    sharp = np.sum((img-mu)**2)/w/h/mu
    #sharp = np.sum((img-mu)**2)/w/h
    return sharp

def sharpness_psd(img):
    F2 = fft.fftshift(fft.fft2(img))
    psd2D = np.abs( F2 )**2
    y, x = np.indices(psd2D.shape)
    y0, x0 = np.array(psd2D.shape)/2
    r = ((x-x0)**2+(y-y0)**2)**0.5
    ind = np.argsort(r.flat)
    r_sorted = r.flat[ind]
    i_sorted = psd2D.flat[ind]
    r_int = r_sorted.astype(int)
    deltar = r_int[1:] - r_int[:-1]  # Assumes all radii represented
    rind = np.where(deltar)[0]       # location of changed radius
    nr = rind[1:] - rind[:-1]
    csim = np.cumsum(i_sorted, dtype=float)
    tbin = csim[rind[1:]] - csim[rind[:-1]]
    radial_prof = tbin / nr
    radial_prof_r = radial_prof * np.arange(len(radial_prof))
    radial_prof_r = radial_prof_r / np.max(radial_prof_r)
    #select_r = int(len(radial_prof)*0.8)
    select_r = int(len(radial_prof_r)*0.9)
    try:
        res = curve_fit(_func, np.arange(len(radial_prof_r))[0:select_r], radial_prof_r[0:select_r], 
                (radial_prof_r[0:select_r].max(), 1,0), 
                bounds=([0, 1/(len(radial_prof_r)), 0], [radial_prof_r[0:select_r].max()*5/(10/(len(radial_prof_r)/2)), 10, 2*radial_prof_r.max()/len(radial_prof_r)]))
        area_s = res[0][0]
        area_n = res[0][2]*len(radial_prof_r)**2/2
        snr = area_s / area_n
        #return np.abs(1/res[0][1])
        if snr > 0.005:
           return np.abs(1/res[0][1])
        else:
           return 0
    except:
        return 0

def sharpness_psd2(img):
    # a modified sharpness measure. The sharpness is weighed against snr.
    F2 = fft.fftshift(fft.fft2(img))
    psd2D = np.abs( F2 )**2
    y, x = np.indices(psd2D.shape)
    y0, x0 = np.array(psd2D.shape)/2
    r = ((x-x0)**2+(y-y0)**2)**0.5
    ind = np.argsort(r.flat)
    r_sorted = r.flat[ind]
    i_sorted = psd2D.flat[ind]
    r_int = r_sorted.astype(int)
    deltar = r_int[1:] - r_int[:-1]  # Assumes all radii represented
    rind = np.where(deltar)[0]       # location of changed radius
    nr = rind[1:] - rind[:-1]
    csim = np.cumsum(i_sorted, dtype=float)
    tbin = csim[rind[1:]] - csim[rind[:-1]]
    radial_prof = tbin / nr
    radial_prof_r = radial_prof * np.arange(len(radial_prof))
    radial_prof_r = radial_prof_r / np.max(radial_prof_r)
    #select_r = int(len(radial_prof)*0.8)
    select_r = int(len(radial_prof_r)*0.9)
    try:
        res = curve_fit(_func, np.arange(len(radial_prof_r))[0:select_r], radial_prof_r[0:select_r], 
                (radial_prof_r[0:select_r].max(), 1,0), 
                bounds=([0, 1/(len(radial_prof_r)), 0], [radial_prof_r[0:select_r].max()*5/(10/(len(radial_prof_r)/2)), 10, 2*radial_prof_r.max()/len(radial_prof_r)]))
        area_s = res[0][0]
        area_n = res[0][2]*len(radial_prof_r)**2/2
        snr = area_s / area_n
        #return np.abs(1/res[0][1])
        if snr > 0.005:
           return np.abs(1/res[0][1]) * snr
        else:
           return 0
    except:
        return 0

    # if np.argmax(radial_prof_r) > len(radial_prof_r)*0.7:
    #     return 0
    # else:
    #     res = curve_fit(_func, np.arange(select_r), radial_prof_r[0:select_r], (radial_prof_r.max(), 1 , 0.1))
    #     return np.abs(1/res[0][1])


def sharpness_psd3(img):
    # yet another psd method, use overfit to quantify sharpness.
    sigma_width = np.pi/2
    F2 = fft.fftshift(fft.fft2(img))
    psd2D = np.abs( F2 )**2
    y, x = np.indices(psd2D.shape)
    y0, x0 = np.array(psd2D.shape)/2
    r = ((x-x0)**2+(y-y0)**2)**0.5
    ind = np.argsort(r.flat)
    r_sorted = r.flat[ind]
    i_sorted = psd2D.flat[ind]
    r_int = r_sorted.astype(int)
    deltar = r_int[1:] - r_int[:-1]  # Assumes all radii represented
    rind = np.where(deltar)[0]       # location of changed radius
    nr = rind[1:] - rind[:-1]
    csim = np.cumsum(i_sorted, dtype=float)
    tbin = csim[rind[1:]] - csim[rind[:-1]]
    radial_prof = tbin / nr
    radial_prof_r = radial_prof * np.arange(len(radial_prof))
    radial_prof_r = radial_prof_r / np.max(radial_prof_r)
    #select_r = int(len(radial_prof)*0.8)
    select_r = int(len(radial_prof_r)*0.8)
    try:
        res = curve_fit(_func, np.arange(len(radial_prof_r))[0:select_r], radial_prof_r[0:select_r], 
                (radial_prof_r[0:select_r].max(), 1,0), 
                bounds=([0, 1/(len(radial_prof_r)), 0], [radial_prof_r[0:select_r].max()*5/(10/(len(radial_prof_r)/2)), 10, 2*radial_prof_r.max()/len(radial_prof_r)]))
        area_s = res[0][0]
        area_n = res[0][2]*len(radial_prof_r)**2/2
        snr = area_s / area_n
        fit_y = _func(np.arange(len(radial_prof_r)), *res[0])
        half_len = len(radial_prof_r)//2
        extra_p = np.sum((radial_prof_r-fit_y)[0:half_len])
        p_factor = np.tanh(extra_p*sigma_width)+1
        sharp = np.abs(1/res[0][1])
        new_sharp = sharp*p_factor
        return new_sharp
    except:
        return 0

def gaussian(x, a, b, c, d):
    return a*np.exp(-(x-b)**2/(2*c**2)) + d

def gaussian2(x, a, b, c):
    return a*np.exp(-(x-b)**2/(2*c**2))


def autofocus(sem, search_step=10, wd_searh_factor=3, retry = 3, refine = True):
    delay = sem.GetValue("AP_FRAME_TIME")/1000*1.5
    buffer = tempfile.TemporaryFile(suffix='.bmp').name
    X = sem.GetValue("AP_RED_RASTER_POSN_X")
    Y = sem.GetValue("AP_RED_RASTER_POSN_Y")
    H = sem.GetValue("AP_RED_RASTER_H")
    W = sem.GetValue("AP_RED_RASTER_W")
    
    # local function to optimize
    def target(wd):
        sem.SetValue("AP_WD", wd)
        time.sleep(delay)
        sem.Grab(buffer,X,Y,W,H)
        img = imread(buffer)
        val = sharpness_nv_g(img)
        print("WD={:.3f},sharpness={}".format(wd*1000, val))
        return val

    def search(span, step):
        start_wd = sem.GetValue("AP_WD")
        wd_list = np.linspace(start_wd-span, start_wd+span, num = step)
        sem.SetValue("AP_WD", wd_list[0])
        time.sleep(0.3)
        sharpness_list = np.array([target(wd) for wd in wd_list])
        start_a = sharpness_list.max()
        start_b = wd_list[np.argmax(sharpness_list)]
        start_c = (wd_list.max() - wd_list.min())/2
        try:
            #res = curve_fit(gaussian, wd_list, sharpness_list, (start_a, start_b, start_c, start_d))
            res = curve_fit(gaussian2, wd_list, sharpness_list, (start_a, start_b, start_c))
            b_std = res[1][1,1]**0.5
            if res[0][1] + 0.1*span > wd_list[-1]:
                sem.SetValue("AP_WD", wd_list[-1])
                return "outside"
            elif res[0][1] - 0.1*span < wd_list[0]:
                sem.SetValue("AP_WD", wd_list[0])
                return "outside"
            else:
                sem.SetValue("AP_WD", wd_list[0])
                time.sleep(0.3)
                sem.SetValue("AP_WD", res[0][1])
                return "success"
        except:
            sem.SetValue("AP_WD", wd_list[0])
            time.sleep(0.3)
            sem.SetValue("AP_WD", wd_list[np.argmax(sharpness_list)])
            return "failed"

    factor = wd_searh_factor
    state = ""
    # first find a rough focus
    count = 0
    while state != "success" and count < retry:
        state = search(sem.GetValue("AP_WIDTH")*factor, search_step)
        print(state)
        if state == "failed":
            factor = factor * 1.5
            count += 1
        if state == "outside":
            factor = factor * 1.4
            count += 1

    if state == "success" and refine:
        search(sem.GetValue("AP_WIDTH"), search_step)
    
    return state

def autofocusstig(sem, search_step=10, iter_n=2, wd_searh_factor=3, stig_range=2, iter_range_damping = 0.8):
    delay = sem.GetValue("AP_FRAME_TIME")/1000*1.7
    buffer = tempfile.TemporaryFile(suffix='.bmp').name
    X = sem.GetValue("AP_RED_RASTER_POSN_X")
    Y = sem.GetValue("AP_RED_RASTER_POSN_Y")
    H = sem.GetValue("AP_RED_RASTER_H")
    W = sem.GetValue("AP_RED_RASTER_W")
    
    # maximum iteration 5 times.
    iter_n = np.clip(iter_n, 1, 5)

    # local function to optimize
    def target(wd):
        sem.SetValue("AP_WD", wd)
        time.sleep(delay)
        sem.Grab(buffer,X,Y,W,H)
        img = imread(buffer)
        val = sharpness_nv_g(img)
        print("WD={:.3f},sharpness={}".format(wd*1000, val))
        return val

    def target_stig(target_val, stig_type = "X"):
        if stig_type == "X":
            AP_NAME = "AP_STIG_X"
        elif stig_type == "Y":
            AP_NAME = "AP_STIG_Y"
        sem.SetValue(AP_NAME, target_val)
        time.sleep(delay)
        sem.Grab(buffer,X,Y,W,H)
        img = imread(buffer)
        val = sharpness_nv_g(img)
        print(f"Stig {stig_type}={target_val:.3f},sharpness={val}")
        return val

    def search_wd(span, step):
        start_wd = sem.GetValue("AP_WD")
        wd_list = np.linspace(start_wd-span, start_wd+span, num = step)
        sem.SetValue("AP_WD", wd_list[0])
        time.sleep(0.3)
        sharpness_list = np.array([target(wd) for wd in wd_list])
        start_a = sharpness_list.max()
        start_b = wd_list[np.argmax(sharpness_list)]
        start_c = (wd_list.max() - wd_list.min())/2
        try:
            #res = curve_fit(gaussian, wd_list, sharpness_list, (start_a, start_b, start_c, start_d))
            res = curve_fit(gaussian2, wd_list, sharpness_list, (start_a, start_b, start_c))
            #b_std = res[1][1,1]**0.5
            if res[0][1] + 0.1*span > wd_list[-1]:
                sem.SetValue("AP_WD", wd_list[-1])
                return "outside"
            elif res[0][1] - 0.1*span < wd_list[0]:
                sem.SetValue("AP_WD", wd_list[0])
                return "outside"
            else:
                sem.SetValue("AP_WD", wd_list[0])
                time.sleep(0.3)
                sem.SetValue("AP_WD", res[0][1])
                return "success"
        except:
            sem.SetValue("AP_WD", wd_list[0])
            time.sleep(0.3)
            sem.SetValue("AP_WD", wd_list[np.argmax(sharpness_list)])
            return "failed"

    def search_stig(span, step, stig_type = "X"):
        if stig_type == "X":
            AP_NAME = "AP_STIG_X"
        elif stig_type == "Y":
            AP_NAME = "AP_STIG_Y"
        start_stig = sem.GetValue(AP_NAME)
        stig_list = np.linspace(start_stig-span, start_stig+span, num = step)
        sem.SetValue(AP_NAME, stig_list[0])
        time.sleep(0.3)
        sharpness_list = np.array([target_stig(val, stig_type) for val in stig_list])
        start_a = sharpness_list.max()
        start_b = stig_list[np.argmax(sharpness_list)]
        start_c = (stig_list.max() - stig_list.min())/2
        try:
            res = curve_fit(gaussian2, stig_list, sharpness_list, (start_a, start_b, start_c))
            b_std = res[1][1,1]**0.5
            if res[0][1]+b_std > stig_list[-1]:
                sem.SetValue(AP_NAME, stig_list[-1])
                return "outside"
            elif res[0][1]-b_std < stig_list[0]:
                sem.SetValue(AP_NAME, stig_list[0])
                return "outside"
            else:
                sem.SetValue(AP_NAME, stig_list[0])
                time.sleep(0.3)
                sem.SetValue(AP_NAME, res[0][1])
                return "success"
        except:
            sem.SetValue(AP_NAME, stig_list[np.argmax(sharpness_list)])
            print("cannot fit.")
            return "failed"

    # now start the main function
    default_wd_range = wd_searh_factor * sem.GetValue("AP_WIDTH")
    state = ""
    # first find a rough focus. This must be successful in the first place
    count = 0
    retry = 3
    while state != "success" and count < retry:
        state = search_wd(default_wd_range, search_step)
        print(state)
        if state == "failed":
            default_wd_range = default_wd_range * 1.5
            count += 1
        if state == "outside":
            default_wd_range = default_wd_range * 1.4
            count += 1
    
    if state == "success":

    
    
    # while state != "success":
    #     state = search_wd(default_wd_range, search_step)
    #     print(state)
    #     if state == "failed":
    #         default_wd_range = default_wd_range * 1.5
    #     elif state == "outside":
    #         pass
    
        wd_range = wd_searh_factor * sem.GetValue("AP_WIDTH") * iter_range_damping
        
        for i in range(iter_n):
            current_wd_range = wd_range * iter_range_damping**i
            state = current_stig_range = stig_range * iter_range_damping**i
            state = search_stig(current_stig_range, search_step, "X")
            state = search_stig(current_stig_range, search_step, "Y")
            state = search_wd(current_wd_range, search_step)

    return state