structen.py 14.5 KB
``````# -*- 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 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)
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

#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
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):