codebase / stimlib

# -*- coding: utf-8 -*-
"""
Created on Sun Mar 12 21:54:40 2017

@author: david
"""

import numpy as np
import scipy as sp
import scipy.ndimage
import progressbar
import glob

def st2(I, s=1, dtype=np.float32):   
    
    #calculates the 2D structure tensor for an image using a gaussian window with standard deviation s
    
    #calculate the gradient
    dI = np.gradient(I.astype(dtype))
      
    #calculate the dimensions of the tensor field
    field_dims = [dI[0].shape[0], dI[0].shape[1], 3]
    
    #allocate space for the tensor field
    Tg = np.zeros(field_dims, dtype=dtype)
    
    #calculate the gradient components of the tensor
    ti = 0
    for i in range(2):
        for j in range(i + 1):
            Tg[:, :, ti] = dI[j] * dI[i]
            ti = ti + 1
    
    #if the user does not want a blurred field
    if(s == 0):
        return Tg
        
    #blur the tensor field
    T = np.zeros(field_dims, dtype=dtype)
    
    for i in range(3):
        T[:, :, i] = scipy.ndimage.filters.gaussian_filter(Tg[:, :, i], [s, s])

    
    return T

def st3(I, s=1, v=[1, 1, 1], dtype=np.float32):
    #calculate the structure tensor field for the 3D input image I given the window size s and voxel size v
    #check the format for the window size
        
    v = np.array(v)
    print("\nCalculating gradient...")
    dI = np.gradient(I.astype(dtype), v[0], v[1], v[2])
    #calculate the dimensions of the tensor field
    field_dims = [dI[0].shape[0], dI[0].shape[1], dI[0].shape[2], 6]
    
    #allocate space for the tensor field
    Tg = np.zeros(field_dims, dtype=np.float32)
    
    #calculate the gradient components of the tensor
    ti = 0
    print("Calculating tensor components...")
    bar = progressbar.ProgressBar()
    bar.max_value = 6
    for i in range(3):
        for j in range(i + 1):
            Tg[:, :, :, ti] = dI[j] * dI[i]
            ti = ti + 1
            bar.update(ti)
    
    #blur the tensor field
    T = np.zeros(field_dims, dtype=np.float32)
    
    print("\nConvolving tensor field...")
    bar = progressbar.ProgressBar()
    bar.max_value = 6
    sigma = s / v
    print(sigma)
    for i in range(6):
        T[:, :, :, i] = scipy.ndimage.filters.gaussian_filter(Tg[:, :, :, i], sigma)
        bar.update(i+1)
        
    return T

def st(I, s=1):
    if I.ndim == 3:
        return st3(I, s)
    elif I.ndim == 2:
        return st2(I, s)
    else:
        print("Image must be 2D or 3D")
    return
        
   
    
def sym2mat(T):
    #Calculate the full symmetric matrix from a list of lower triangular elements.
    #The lower triangular components in the input field T are assumed to be the
    #   final dimension of the input matrix.
    
    #       | 1  2  4  7  |
    #       | 0  3  5  8  |
    #       | 0  0  6  9  |
    #       | 0  0  0  10 |
   
    in_s = T.shape
    
    #get the number of tensor elements
    n = in_s[T.ndim - 1]
    
    #calculate the dimension of the symmetric matrix
    d = int(0.5 * (np.sqrt(8. * n + 1.) - 1.))
    
    #calculate the dimensions of the output field
    out_s = list(in_s)[:-1] + [d] + [d]

    #allocate space for the output field
    R = np.zeros(out_s)
    
    ni = 0
    for i in range(d):
        for j in range(i + 1):
            R[..., i, j] = T[..., ni]
            if i != j:
                R[..., j, i] = T[..., ni]
            ni = ni + 1
    
    return R   

def vec(S, vector=0):
   
    if(S.ndim != 3):
        print("ERROR: a 2D slice is expected")
        return
      
    #convert the field to a full rank-2 tensor
    T = sym2mat(S);
    del(S)
    
    #calculate the eigenvectors and eigenvalues
    l, v = np.linalg.eig(T)
    
    #get the dimension of the tensor field
    d = T.shape[2]
    
    #allocate space for the vector field
    V = np.zeros([T.shape[0], T.shape[1], 3])

    #arrange the indices for each pixel from smallest to largest eigenvalue    
    idx = l.argsort()
    
    for di in range(d):
        b = idx[:, :, -1-vector] == di
        V[b, 0:d] = v[b, :, di]
    
    return V
   
def loadstack(filemask):
    #Load an image stack as a 3D grayscale array
    
    #get a list of all files matching the given mask
    files = [file for file in glob.glob(filemask)]
    
    #calculate the size of the output stack
    I = scipy.misc.imread(files[0])
    X = I.shape[0]
    Y = I.shape[1]
    Z = len(files)
    
    #allocate space for the image stack
    M = np.zeros([X, Y, Z]).astype('float32')
    
    #create a progress bar
    bar = progressbar.ProgressBar()
    bar.max_value = Z
    
    #for each file
    for z in range(Z):
        #load the file and save it to the image stack
        M[:, :, z] = scipy.misc.imread(files[z], flatten="True").astype('float32')
        bar.update(z+1)
    return M

#calculate the anisotropy of a structure tensor given the tensor field S
def anisotropy3(S):

    Sf = sym2mat(S)
    
    #calculate the eigenvectors and eigenvalues
    l, v = np.linalg.eig(Sf)
    
    #store the sorted eigenvalues
    ls = np.sort(l)
    l0 = ls[:, :, 0]
    l1 = ls[:, :, 1]
    l2 = ls[:, :, 2]
    
    #calculate the linear anisotropy
    Cl = (l2 - l1)/(l2 + l1 + l0)
    
    #calculate the planar anisotropy
    Cp = 2 * (l1 - l0) / (l2 + l1 + l0)
    
    #calculate the spherical anisotropy
    Cs = 3 * l0 / (l2 + l1 + l0)
    
    #calculate the fractional anisotropy
    l_hat = (l0 + l1 + l2)/3
    fa_num = (l2 - l_hat) ** 2 + (l1 - l_hat) ** 2 + (l0 - l_hat) ** 2;
    fa_den = l0 ** 2 + l1 ** 2 + l2 ** 2
    FA = np.sqrt(3./2.) * np.sqrt(fa_num) / np.sqrt(fa_den)
    
    return FA, Cl, Cp, Cs

#calculate the fractional anisotropy
def fa(S):
    Sf = sym2mat(S)
    
    #calculate the eigenvectors and eigenvalues
    l, v = np.linalg.eig(Sf)
    
    #store the sorted eigenvalues
    ls = np.sort(l)
    l0 = ls[:, :, 0]
    l1 = ls[:, :, 1]
    
    #if this is a 2D tensor, calculate and return the coherence
    if(S.shape[2] == 3):
        C = ((l0 - l1) / (l0 + l1)) ** 2
        return C
        
    #if this is a 3D tensor    
    elif(S.shape[2] == 6):
        l2 = ls[:, :, 2]
        
        #calculate the fractional anisotropy
        l_hat = (l0 + l1 + l2)/3
        fa_num = (l2 - l_hat) ** 2 + (l1 - l_hat) ** 2 + (l0 - l_hat) ** 2;
        fa_den = l0 ** 2 + l1 ** 2 + l2 ** 2
        FA = np.sqrt(3./2.) * np.sqrt(fa_num) / np.sqrt(fa_den)
        return FA

#calculate the specified eigenvalue for the tensor field
def eigenval(S, ev):
    Sf = sym2mat(S)
    
     #calculate the eigenvectors and eigenvalues
    l, v = np.linalg.eig(Sf)
    
    #store the sorted eigenvalues
    ls = np.sort(l)
    evals = ls[:, :, ev]

    return evals

def amira(filename, T):
    #generates a tensor field that can be imported into Amira
    
    #   0    dx dx   ----> 0
    #   1    dx dy   ----> 1
    #   2    dy dy   ----> 3
    #   3    dx dz   ----> 2
    #   4    dy dz   ----> 4
    #   5    dz dz   ----> 5
    
    #swap the 2nd and 3rd tensor components
    A = np.copy(T)
    A[..., 3] = T[..., 2]
    A[..., 2] = T[..., 3]
    
    #roll the tensor axis so that it is the leading component
    #A = numpy.rollaxis(A, A.ndim - 1)
    A.tofile(filename)
    print("\n", A.shape)

def resample3(T, s=2):
    #resample a tensor field by an integer factor s
    #This function first convolves the field with a box filter and then
    #   re-samples to create a smaller field
    
    #check the format for the window size
    if type(s) is not list:
        s = [s] * 3
    elif len(s) == 1:
        s = s * 3
    elif len(s) == 2:
        s.insert(1, s[0])
    s = np.array(s)
    
    bar = progressbar.ProgressBar()
    bar.max_value = T.shape[3]
    
    #blur with a uniform box filter of size r
    for t in range(T.shape[3]):
        T[..., t] = scipy.ndimage.filters.uniform_filter(T[..., t], 2 * s)
        bar.update(t+1)
        
    #resample at a rate of r
    R = T[::s[0], ::s[1], ::s[2], :]
    return R

def color3(prefix, T, vector='largest', aniso=True):
    #Saves a stack of color images corresponding to the eigenvector and optionally scaled by anisotropy
    
    bar = progressbar.ProgressBar()
    bar.max_value = T.shape[2]
    
    #for each z-axis slice
    for z in range(T.shape[2]):
        S = T[:, :, z, :]                           #get the slice
        V = st2vec(S, vector='smallest')   #calculate the vector
        C = np.absolute(V)                       #calculate the absolute value
        
        if aniso == True:                              #if the image is scaled by anisotropy
            FA, Cl, Cp, Cs = anisotropy(S)          #calculate the anisotropy of the slice
            if vector == 'largest':
                A = Cl
            elif vector == 'smallest':
                A = Cp
        else:                                       #otherwise just scale by 1
            A = np.ones(T.shape[0], T.shape[1])
        image = C * np.expand_dims(A, 3)
        
        filename = prefix + str(z).zfill(4) + ".bmp"
        scipy.misc.imsave(filename, image)
        bar.update(z + 1)
        
def st2stack(T, outfile, **kwargs):
    eigenvector = False                                             #initialize the colormap flags to false
    aniso_color = False
    aniso_alpha = False
    #image = False
    aniso_pwr = 1
    cimage_pwr = 1
    aimage_pwr = 1
    anisostretch = 1                                                #set the contrast stretch parameter
    alpha_channel = False
    alpha_image = False
    color_image = False
    
    for k,v in kwargs.items():                                  #for each argument
        if(k == "ev"):                                               #if the user wants a colormap based on an eigenvector
            eigenvector = True                                     #set the eigenvector flag to true
            ev = v                                                 #save the desired eigenvector
        if(k == "aniso"):                                            #if the user wants to factor in anisotropy
            aniso = True                                      #set the anisotropy flag to true
            aniso_channel = v                                      #save the anisotropy channel
        if(k == "aniso_color"):
            aniso_color = v
        if(k == "aniso_alpha"):
            aniso_alpha = v
        if(k == "apwr"):                                              #if the user wants to amplify the anisotropy
            aniso_pwr = v                                                 #set the anisotropy exponent
        if(k == "cipwr"):                                            #if the user specifies the image power
            cimage_pwr = v
        if(k == "aipwr"):
            aimage_pwr = v
        if(k == "alphaimage"):
            Ia = v
            alpha_image = True
        if(k == "colorimage"):
            Ic = v
            color_image = True
        if(k == "anisostretch"):
            anisostretch = v
        if(k == "alpha"):
            alpha_channel = v
             
    bar = progressbar.ProgressBar()
    bar.max_value = T.shape[2]
    for i in range(0, T.shape[2]):
    #for i in range(0, 50):
    
        if(alpha_image or alpha_channel):
            img = np.ones([T.shape[0], T.shape[1], 4])
        else:
            img = np.ones([T.shape[0], T.shape[1], 3])
        if(eigenvector):
            V = st2vec(T[:, :, i], ev)                         #get the vector field for slice i corresponding to eigenvector ev
            img[:, :, 0:3] = V                                      #update the image with the vector field information
        if(aniso):                                             #if the user is requesting anisotropy be incorporated into the image
            FA, Cl, Cp, Cs = anisotropy(T[:, :, i])        #calculate the anisotropy of the tensors in slice i
            if(aniso_channel == "fa"):
                A = FA
            elif(aniso_channel == "l"):
                A = Cl
            elif(aniso_channel == "p"):
                A = Cp
            else:
                A = Cs
            if(aniso_alpha):
                print("rendering anisotropy to the alpha channel")
                img[:, :, 3] = A ** aniso_pwr * anisostretch
            if(aniso_color):
                print("rendering anisotropy to the color channel")
                img[:, :, 0:3] = img[:, :, 0:3] * np.expand_dims(A ** aniso_pwr, 3) * anisostretch               
        if(alpha_image):
            img[:, :, 3] = Ia[:, :, i]/255 ** aimage_pwr
        if(color_image):
            img[:, :, 0:3] = img[:, :, 0:3] * (np.expand_dims(Ic[:, :, i], 3)/255) ** cimage_pwr    #multiply the vector field by the image intensity
        #outname = outfile + str(i).zfill(3) + ".bmp"                    #get the file name to be saved
        outname = outfile.replace("*", str(i).zfill(3))
        
        sp.misc.imsave(outname, np.ndarray.astype(np.abs(img)*255, "uint8"))                              #save the output image
        bar.update(i+1)
        

#this function takes a 3D image and turns it into a stack of color images based on the structure tensor
def img2stack(I, outfile, **kwargs):
    
    vs = [1, 1, 1]                                                  #set the default voxel size to 1
    w = 5
    
    for k,v in kwargs.items():                                  #for each argument
        if(k == "voxelsize"):                                       #if the voxel size is specified
            if(len(v) == 1):                                        #if the user just specifies one value
                vs = [v] * 3                                        #assume that the voxels are isotropic, create a list of 3 v's
            elif(len(v) == 2):                                      #if the user specifies two values
                vs[0] = v[0]                                        #assume that the voxels are isotropic along (x, y) and anisotropic along z
                vs[1] = v[0]
                vs[2] = v[1]
            elif(len(v) == 3):
                vs = v
        if(k == "window"):                                          #if the user specifies a window size
            w = v
            
    T = st3(I, w, vs)                                       #calculate the structure tensor
    
    st2stack(T, outfile, **kwargs)
    
def stack2stack(infile_mask, outfile, **kwargs):
     
    I = loadstack(infile_mask)                            #load the file mask
    for k,v in kwargs.items():                                  #for each argument
        if(k == "ipwr"):
            img = I
    
    img2stack(I, outfile, image=img, **kwargs)                                #call the function to convert the image to an output ST stack