ACORBA Source code

#Load libraries
print("Loading modules\n##################################################")
from time import sleep, time
start_time = time()
import numpy as np
import os
import tifffile as tiff
from skimage.transform import resize
import matplotlib.pyplot as plt
import cv2
from skimage.morphology import skeletonize
import pandas as pd
from sklearn.cluster import KMeans
from utils import skeleton_endpoints, angle_to, plot_imgs, binary_cleaner, trad_micro_sandwich,trad_micro_through, Diff,looplenght_determinator
import PySimpleGUI as sg
import argparse
from pathlib import PurePath
import sys
import gc
import tifffile
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1' 
from tensorflow.autograph import set_verbosity
set_verbosity(0)
print("ALL LIBRARIES ARE IMPORTED, ACORBA WILL START NOW\n##################################################")

#retrieve user inputs from the GUI
parser = argparse.ArgumentParser(description='ACORBA')
parser.add_argument('--input_folder', type=str)
parser.add_argument('--exp_type', type=str)
parser.add_argument('--saveplot', type=str)
parser.add_argument('--normalization', type=str)
parser.add_argument('--prediction', type=str)
parser.add_argument('--binary_folder', type=str)
parser.add_argument('--rootplot', type=str)
parser.add_argument('--method', type=str)
parser.add_argument('--custom', type=str)
parser.add_argument('--smooth', type=str)
parser.add_argument('--superaccuracy', type=str)
parser.add_argument('--savesegmentation', type=str)
parser.add_argument('--tradmethod', type=str)
args = parser.parse_args()
print('user inputs')
rootfolder=args.input_folder
print('Root folder: '+rootfolder)
saveplot=args.saveplot
print('Save angle plot: '+saveplot)
rootplot=args.rootplot
print('Save root plot: '+rootplot)
save_segmentation=args.savesegmentation
print('Save raw segmentation: '+save_segmentation)
exp_type=args.exp_type
print('Experiment type: '+exp_type)
normalization=args.normalization
print('Normalization: '+normalization)
predictiont=args.prediction
print('Save prediction plot: '+predictiont)
seg_method=args.method
print('Segmentation method: '+seg_method)
DML_super_accuracy=args.superaccuracy
print('Super accuracy mode for scanner: '+DML_super_accuracy)
binary_folder=args.binary_folder
print('Use masks from: '+binary_folder)
custom_models=args.custom
print('Use models from: '+custom_models)
deactivate_smoothing=args.smooth
print('Deactivate smoothing: '+deactivate_smoothing)
trad_method=args.tradmethod
print("Traditional method: "+trad_method)

'''#for debugging
rootfolder="E:/Pic libraries/test bank/Vincent/test2 new model/test v1.2/micro sand/"
saveplot="False"
rootplot="False"
exp_type="Microscopy Sandwich"
normalization="True"
predictiont="None"
seg_method="Deep Machine Learning"
binary_folder=""
custom_models=""
deactivate_smoothing="False"
save_segmentation='True'
os.chdir('E:/ACORBA/v1.2')
###'''

#Set random seed to "The Answer to the Ultimate Question of Life, The Universe, and Everything."
np.random.seed(42)
size=256

if save_segmentation=="True":
    dir = rootfolder+"/Saved Segmentations/"
    if not os.path.exists(dir):
        os.mkdir(dir)

#Load DML models and weights
if len(binary_folder)==0:
    if len(custom_models)>0:
        from keras.models import model_from_json
        print('Custom models and weights loading:')
        model_weights=os.listdir(custom_models)
        custom_models_list=[]
        custom_weights_list=[]
        for i in model_weights:
            if i.endswith(".json"):
                custom_models_list.append(i)
            if i.endswith('.h5'):
                custom_weights_list.append(i)
        ans_surface = 'surface'
        ans_tip = 'tip'
        for title in custom_models_list:
            if ans_surface in title:
                json_file = open(custom_models+"/"+title, 'r')
                loaded_model_json = json_file.read()
                json_file.close()
                loaded_model = model_from_json(loaded_model_json)
                print('root surface model loaded')
            elif ans_tip in title:
                json_file = open(custom_models+"/"+title, 'r')
                loaded_model_json = json_file.read()
                json_file.close()
                loaded_model_tip = model_from_json(loaded_model_json)
                print('root tip model loaded')
        for title in custom_weights_list:
            if ans_surface in title:
                loaded_model.load_weights(custom_models+"/"+title)
                print('root surface weights loaded')
            elif ans_tip in title:
                loaded_model_tip.load_weights(custom_models+"/"+title)
                print('root tip weights loaded')
    else:
        from keras_unet.models import satellite_unet
        print("Default models and weights loading")
        loaded_model = satellite_unet(input_shape=(size, size, 1))
            #load root tip model
        loaded_model_tip = satellite_unet(input_shape=(size, size, 1))
            # load weights into new model
        if exp_type=='Microscopy Through':
            loaded_model.load_weights("models/model_through_256_sat_jaccard_update16062021_54.h5")
            loaded_model_tip.load_weights("models/model_through_TIP2_256_sat_jaccard_update17062021_176.h5")
        else:
            loaded_model.load_weights("models/sandwichsurface_256_jaccard_update24062021_145.h5")
            loaded_model_tip.load_weights("models/microsandtip_256_sat_jaccard_update24062021_71.h5")
    print("Models and weights are loaded")
     
# Predict root masks

rootlist_temp=os.listdir(rootfolder)
rootlist=[]
for i in rootlist_temp:
    if i.endswith(".tif"):
        rootlist.append(i)
del rootlist_temp
    
print("The following files will be processed")
print(rootlist)
#Start for loop for all the file in the rootlist
angletotal=[]
inc_stack2=0
rootlist2=[]

#Setting the progress bar total length by opening every image lenght info and summing their timeframe numbers
looplenght=looplenght_determinator(list_image=rootlist,folder=rootfolder)

for inc_root in rootlist:
    sg.theme('Light Brown 8')
    print("Processing file "+inc_root)
    test= tiff.imread(rootfolder+'/'+inc_root)
    test_shape=test.shape
    tradlist=[]
    print("Converting stack to an array and resizing/padding to 256x256")
    X_test = np.zeros((len(test),size, size), dtype=np.float32)
    try:
        n=0
        while n<len(test):
            sg.OneLineProgressMeter(inc_root+' import, padding, resizing', n+1, len(test), 'key')
            im=(test[n]/255).astype('uint8')
            tradlist.append(im)
            if im.shape[0]!=im.shape[1]:
                desired_size = size
                old_size = im.shape # old_size is in (height, width) format
                ratio = float(desired_size)/max(old_size)
                new_size = tuple([int(x*ratio) for x in old_size])
                img_t = cv2.resize(im, (new_size[1], new_size[0]))
                delta_w = desired_size - new_size[1]
                delta_h = desired_size - new_size[0]
                top, bottom = delta_h//2, delta_h-(delta_h//2)
                left, right = delta_w//2, delta_w-(delta_w//2)
                color = np.average(im)
                img_t = cv2.copyMakeBorder(img_t, top, bottom, left, right, cv2.BORDER_CONSTANT,value=color)
                X_test[n]=img_t/255
            else:
                X_test[n] = resize(im, (size, size), mode='constant', preserve_range=True)/255
            n=n+1
        if save_segmentation=="True":    
            path=rootfolder+'//Saved Segmentations/'+inc_root+'_'+'original_resized256.tif'
            tifffile.imwrite(path, X_test)
        tradlist=np.array(tradlist)
        #Predictions
        if seg_method=="Deep Machine Learning":
            X_test=np.expand_dims(X_test,3)
            print("Prediction of root surface by deep machine learning")
            preds_test = loaded_model.predict(X_test, verbose=2,batch_size=1)
            prediction =(preds_test > 0.5).astype(np.uint8)
            print("Prediction of root tip by deep machine learning")
            preds_test = loaded_model_tip.predict(X_test, verbose=2,batch_size=1)
            prediction_tip =(preds_test > 0.5).astype(np.uint8)
            if predictiont =="First":
                print("Exporting first timeframe prediction")
                plot_imgs(org_imgs=X_test, folder=args.input_folder, inc=inc_root, pred_imgs=prediction,nm_img_to_plot=1,color="green",kind="root") #Need to be in the exp directory
                plot_imgs(org_imgs=X_test,folder=args.input_folder, inc=inc_root, pred_imgs=prediction_tip,nm_img_to_plot=1,color="red",kind="tip") #SAME
            elif predictiont =='All':
                print("Exporting all timeframe predictions")
                plot_imgs(org_imgs=X_test, folder=args.input_folder, inc=inc_root, pred_imgs=prediction,nm_img_to_plot=len(X_test),color="green",kind="root") #Need to be in the exp directory
                plot_imgs(org_imgs=X_test,folder=args.input_folder, inc=inc_root, pred_imgs=prediction_tip,nm_img_to_plot=len(X_test),color="red",kind="tip") #SAME
        elif seg_method=='Own masks':
            prediction=tiff.imread(binary_folder+'/surface/'+inc_root)
            prediction_tip=tiff.imread(binary_folder+'/tip/'+inc_root)
            if np.max(prediction.shape)>256:
                prediction= resize(prediction, (len(prediction),size, size), mode='constant', preserve_range=True)
                prediction_tip= resize(prediction_tip, (len(prediction_tip),size, size), mode='constant', preserve_range=True)
            if np.max(prediction>1):
                prediction=prediction/np.max(prediction)
            if np.max(prediction_tip>1):
                prediction_tip=prediction_tip/np.max(prediction_tip)
        else:
            print('Traditionnal segmentation of root surface')
            prediction=[]
            if exp_type=='Microscopy Through':
                for img_pred in X_test:
                    prediction.append(trad_micro_through(img_pred))
            else:
                for img_pred in tradlist:
                    prediction.append(trad_micro_sandwich(img_pred))
            prediction=np.array(prediction)
            X_test=np.expand_dims(X_test,3)
            print("Prediction of root tip by deep machine learning")
            preds_test = loaded_model_tip.predict(X_test, verbose=2,batch_size=1)
            prediction_tip =(preds_test > 0.5).astype(np.uint8)
            if predictiont =="First":
                print("Exporting first timeframe prediction")
                plot_imgs(org_imgs=X_test, folder=rootfolder+'/', inc=inc_root, pred_imgs=prediction,nm_img_to_plot=1,color="green",kind="root") #Need to be in the exp directory
                plot_imgs(org_imgs=X_test,folder=rootfolder+'/', inc=inc_root, pred_imgs=prediction_tip,nm_img_to_plot=1,color="red",kind="tip") #SAME
            elif predictiont =='All':
                print("Exporting all timeframe predictions")
                plot_imgs(org_imgs=X_test, folder=rootfolder+'/', inc=inc_root, pred_imgs=prediction,nm_img_to_plot=len(X_test),color="green",kind="root") #Need to be in the exp directory
                plot_imgs(org_imgs=X_test,folder=rootfolder+'/', inc=inc_root, pred_imgs=prediction_tip,nm_img_to_plot=len(X_test),color="red",kind="tip") #SAME
        if save_segmentation=="True":    
            path=rootfolder+'//Saved Segmentations/'+inc_root+'_'+'root_tip_256.tif'
            tifffile.imwrite(path, prediction_tip)
            path=rootfolder+'//Saved Segmentations/'+inc_root+'_'+'root_surface_256.tif'
            tifffile.imwrite(path, prediction)
        #Angle calculation
        print("Calculation of the angles")
        inc_stack=0
        slopetip=[]
        inter=[]
        angleslist2=[]
        Rs=[]
        flipping=0
        while inc_stack<len(X_test):
            if not sg.OneLineProgressMeter('Calculation in progress', inc_stack2, looplenght, 'single'):
                sys.exit("User stop!!!!!!!!!!!!!!")
            ol=prediction[inc_stack]#import binary mask
            oltip=prediction_tip[inc_stack]
            
            ############Binary preprocessessing#############
            #Fill holes
            kernel = np.ones((2,2),np.uint8)
            oltip=cv2.dilate(oltip,kernel,iterations=5)
            oltip=cv2.erode(oltip,kernel,iterations=5)
            oltip=cv2.dilate(oltip,kernel,iterations=2)
            oltip=cv2.erode(oltip,kernel,iterations=5)
            ol=cv2.dilate(ol,kernel,iterations=5)
            ol=cv2.erode(ol,kernel,iterations=5)
            ol=cv2.dilate(ol,kernel,iterations=2)
            ol=cv2.erode(ol,kernel,iterations=5)
            
            #Flip the tip on the left side if on the right side (horizontal flip)
            if inc_stack==0 and np.sum(ol[:,250])<2:
                flipping=1
            if flipping==1:
                ol=cv2.flip(ol,1)
                oltip=cv2.flip(oltip,1)
                
            #Cleaning
            oltip=binary_cleaner(oltip)
            ol=binary_cleaner(ol)
            
            #Find contour and clean it
            contour=cv2.morphologyEx(oltip, cv2.MORPH_GRADIENT, kernel)
            perimeter=np.where(contour==1)
            contour_root=cv2.morphologyEx(ol, cv2.MORPH_GRADIENT, kernel)
            perimeter_root=np.where(contour_root==1)
            
            #Define area of the root tip on the root surface prediction
            M = cv2.moments(oltip)
            # calculate x,y coordinate of center
            cX = int(M["m10"] / M["m00"])
            cY = int(M["m01"] / M["m00"])
            circle=cv2.circle(ol.copy(), (cX, cY), 40, 2, -1)
            circle[circle==2]=0
            ol2=ol.copy()
            ol2=ol2-circle
            contour_tip=cv2.morphologyEx(ol2, cv2.MORPH_GRADIENT, kernel)
            perimeter_tip=np.where(contour_tip==1)
            contours, hierarchy = cv2.findContours(ol2, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
            (xd,yd),radius = cv2.minEnclosingCircle(np.stack((perimeter_tip[1],perimeter_tip[0]),axis=-1))
            center = (int(xd),int(yd))
            radius = int(radius)
            mask2=cv2.drawContours(ol2.copy(), contours[0],-1, color=3)
            rr=cv2.circle(mask2,center,radius,4,2)
            peri_circle=np.where(cv2.circle(mask2,center,radius,4,2)==4)
            coord_circle=np.stack((peri_circle[1],peri_circle[0]),axis=-1)
            coord_peri=np.stack((perimeter_tip[1],perimeter_tip[0]),axis=-1)
            intcp=list(set(map(tuple,coord_peri)).intersection(set(map(tuple,coord_circle))))
            intercp=list(zip(*intcp))
            intercpx=np.array(intercp[0])
            intercpy=np.array(intercp[1])
            intercp2=np.stack((intercpx,intercpy),-1)
    
            km = KMeans(
            n_clusters=2, init='random',
            n_init=10, max_iter=300, 
            tol=1e-04, random_state=0
            )   
            y_km = km.fit_predict(intercp2)
            g1=len(np.where(y_km==0)[0])
            g2=len(np.where(y_km==1)[0])
            indd=y_km.argsort()
            intercp2=intercp2[indd]
            if g1<g2:
                tipx=intercp2[0][0]
                tipy=intercp2[0][1]
            else:
                m=len(intercp2)
                tipx=intercp2[m-1][0]
                tipy=intercp2[m-1][1]
                
            #Skeletonize root tip
            skel = skeletonize(ol2,method='lee')
            skel=np.array(skel)
            skel=np.uint8(skel)
            cfit=np.where(skel==1)
            endpoints_skel=skeleton_endpoints(skel)
            p1temp=(endpoints_skel[1][0],endpoints_skel[0][0])
            p2temp=(endpoints_skel[1][1],endpoints_skel[0][1])
            dtip_P1=np.sqrt((p1temp[0]-tipx)**2+(p1temp[1]-tipy)**2)
            dtip_P2=np.sqrt((p2temp[0]-tipx)**2+(p2temp[1]-tipy)**2)
            if dtip_P1<dtip_P2:
                origin=p1temp
                end=p2temp
            else:
                origin=p2temp
                end=p1temp
            pp=0
            
            points = cfit[1],cfit[0]
            pp=0
            #lenghtskel=len(skelperi[1])
            newskelx=[]
            newskely=[]
            while pp<len(cfit[1])-1:
                if pp==0:
                    listpoint=(cfit[1],cfit[0])
                    start=origin
                    indextormx=np.where(listpoint[0]==start[0])
                    indextormy=np.where(listpoint[1]==start[1])
                    indextorm=int(np.intersect1d(indextormx, indextormy)[0])
                    listpoint=[np.delete(listpoint[0],indextorm),np.delete(listpoint[1],indextorm)]
                    newskelx.append(start[0])
                    newskely.append(start[1])
                else:
                    listpoint=[np.delete(listpoint[0],closestindex),np.delete(listpoint[1],closestindex)]
                    
                uu=0
                distance=[]
                while uu<len(listpoint[1]):
                    distance.append(np.sqrt((start[0] - listpoint[0][uu]) ** 2 + (start[1] - listpoint[1][uu]) ** 2))
                    uu=uu+1
                distancear=np.array(distance)
                closestindex=int(np.where(distance==min(distance))[0])
                newskelx.append(listpoint[0][closestindex])
                newskely.append(listpoint[1][closestindex])
                pp=pp+1
            idx=int(len(newskelx)/2)    
    
            #Fit root tip skeleton (avoid branches)
            splines=np.polyfit(newskelx,newskely,1)
            poly1d = np.poly1d(splines)
            xnew=newskelx
            ynew= poly1d(xnew)
            Tipmax_x=[origin[0],newskelx[idx]]
            Tipmax_y=[origin[1],ynew[idx]]
            
            #adapt fitting line x to either go left of right (<90° or >90°) (only one first timeframe, avoid root tip detection jumping one side to another during analysis)
            if inc_stack==0:
                if origin[0]<newskelx[idx]:
                    xnew=np.linspace(0,newskelx[idx],500)
                    direction=0
                else:
                    xnew=np.linspace(newskelx[idx],size,500)
                    direction=1
            else:
                if direction==0:
                    xnew=np.linspace(0,newskelx[idx],500)
                else:
                    xnew=np.linspace(newskelx[idx],size,500)
            ynew= poly1d(xnew)
            
            #Find the closest pair of coordinates between perimeter and middle model
            listx_fit=xnew.tolist()
            listy_fit=ynew.tolist()
            listx_mask=perimeter_tip[1].tolist()
            listy_mask=perimeter_tip[0].tolist()
            i=0
            distance=[]
            points=[]
            lenghtfit=len(listx_fit)
            lenghtmask=len(listx_mask)
            while i < lenghtmask:
                totestx_mask=listx_mask[i]
                totesty_mask=listy_mask[i]
                yy=0
                while yy < lenghtfit:    
                    totestx_fit=listx_fit[yy]
                    totesty_fit=listy_fit[yy]
                    distance.append(np.sqrt((totestx_fit - totestx_mask) ** 2 + (totesty_fit - totesty_mask) ** 2))
                    points.append([totestx_fit, totesty_fit])
                    yy=yy+1
                i=i+1
            closest_d=min(distance)
            closest_index=distance.index(closest_d)
            closest_coord=points[closest_index]
            Tipmax_x[0]=closest_coord[0]
            Tipmax_y[0]=closest_coord[1]
    
    
            #Calculate the angle
            angle2=angle_to((Tipmax_x[1],Tipmax_y[1]),(Tipmax_x[0],Tipmax_y[0]), clockwise=False)
            angle2=angle2-180
            angleslist2.append(angle2)
            
            #Plot
            plt.clf()
            plt.ioff()
            plt.scatter(perimeter_root[1],perimeter_root[0],color="gainsboro",s=0.5)
            plt.scatter(perimeter_tip[1],perimeter_tip[0],color="red",s=0.5)
            plt.plot(xnew,ynew,linewidth=2,color='orange')
            plt.plot(newskelx,newskely)
            plt.xlim([0, size])
            plt.ylim([size, 0])
            #plt.scatter(Tipmax_x[0],Tipmax_y[0],color="blue")
            plt.scatter(Tipmax_x[1],Tipmax_y[1],color="purple")
            #plt.plot(x_s, poly1d_fntip(x_s),'r--', color="red",linewidth=1.0)
            plt.axis("off")
            stamp=r'File: '+inc_root+' /// Timeframe: '+str(inc_stack+1)+' /// Angle: '+str(angle2)+'°'
            plt.text(1, 1, stamp, fontsize=6)
            plt.show(block=False)
            if rootplot=='True':
                plt.savefig(rootfolder+'/'+inc_root+'_'+'rootplot_'+str(inc_stack+1)+'.png')
            
            #Increment timeframe loops
            inc_stack=inc_stack+1
            inc_stack2=inc_stack2+1
            sleep(0.02)
        angletotal.append((angleslist2))
        rootlist2.append(inc_root)
    except Exception as e:
        print('Oups something went wrong!')
        print(e)
        inc_stack2=inc_stack2+len(test)-inc_stack
        angletotal.append((angleslist2))
        rootlist2.append(inc_root)
        
        
listremoved=Diff(rootlist,rootlist2)
angletotalf=np.array(angletotal,dtype=object)       
angletotalf2=angletotalf.copy()
angletoplot=angletotalf2.copy()
if normalization=='True':
    i=0
    angletotalf2norm=angletotalf.copy()
    while i<len(angletotalf):
        angletotalf2norm[i]=np.array(angletotalf2[i])-angletotalf2[i][0]
        i=i+1
    angletoplot=angletotalf2norm.copy()
if saveplot=='True':
    i=0
    while i<len(angletoplot):
        plt.clf()
        plt.plot(np.linspace(0,len(angletoplot[i]),len(angletoplot[i])),angletoplot[i],color='orange')
        plt.xlabel('Timeframes')
        if normalization=='True':
            plt.ylabel('Relative bending angle (°)')
        else:
            plt.ylabel('Bending angle (°)')
        #plt.ylim(np.min(angletotalf2)-10,np.max(angletotalf2)+10)
        print(rootfolder+'/'+rootlist2[i]+'_angle.png')
        if normalization=='True':
            plt.savefig(rootfolder+'/'+rootlist2[i]+'_angle (normalized).png')
        else:
            plt.savefig(rootfolder+'/'+rootlist2[i]+'_angle (raw).png')
        plt.show(block=False)
        i=i+1
        

#Export output as a xlsx file containing all the angles
print('writing output')

outangles=np.array(angletotalf2)
len_list=[]
for zz in outangles:
    len_list.append(len(zz))
len_list=np.array(len_list)
len_dif=np.unique(len_list)
if len(len_dif)>1:
    #outangles=np.transpose(outangles)
    d = dict(enumerate(outangles.flatten(), 1))
    df = pd.DataFrame.from_dict(d, orient='index')
    df=df.transpose()
    df.columns=rootlist2
else:
    outangles=np.transpose(outangles)
    df = pd.DataFrame(outangles)
    df.columns=rootlist2
    
if normalization=='True':
    outanglenorm=np.array(angletotalf2norm)
    if len(len_dif)>1:
        #outangles=np.transpose(outangles)
        d2 = dict(enumerate(outanglenorm.flatten(), 1))
        df2 = pd.DataFrame.from_dict(d2, orient='index')
        df2=df2.transpose()
        df2.columns=rootlist2
    else:
        outanglenorm=np.transpose(outanglenorm)
        df2 = pd.DataFrame(outanglenorm)
        df2.columns=rootlist2
#
pathout = PurePath(rootfolder)
foldername=pathout.name
writer = pd.ExcelWriter(rootfolder+'/'+foldername+'_output.xlsx')
    # write dataframe to excel
if normalization=='True':
    df.to_excel(writer,'Raw data')
    df2.to_excel(writer,'Normalized data')
else:
    df.to_excel(writer,'Raw data')
# save the excel
writer.save()
print("I'm done")
if len(listremoved)==0:
    print("everything went well")
else:
    print("The following files raised errors and were not or only partially analyzed: ",listremoved)
print("it took me",time() - start_time, "seconds to do my job!")

gc.collect()