structen.py 7.77 KB
# -*- coding: utf-8 -*-
"""
Created on Sun Mar 12 21:54:40 2017

@author: david
"""

import numpy
import scipy.ndimage
import progressbar
import glob

def st2(I, s=1, dtype=numpy.float32):

#calculates the 2D structure tensor for an image using a gaussian window with standard deviation s

#calculate the gradient
dI = numpy.gradient(I)

#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 = numpy.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

#blur the tensor field
T = numpy.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):
#calculate the structure tensor field for the 3D input image I given the window size s in 3D
#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])

print("\nCalculating gradient...")
dI = numpy.gradient(I)
#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 = numpy.zeros(field_dims, dtype=numpy.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 = numpy.zeros(field_dims, dtype=numpy.float32)

print("\nConvolving tensor field...")
bar = progressbar.ProgressBar()
bar.max_value = 6
for i in range(6):
T[:, :, :, i] = scipy.ndimage.filters.gaussian_filter(Tg[:, :, :, i], s)
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 * (numpy.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 = numpy.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 st2vec(S, vector='largest'):
#Create a color image from a 2D or 3D structure tensor slice

#convert the field to a full rank-2 tensor
T = sym2mat(S);
del(S)

#calculate the eigenvectors and eigenvalues
l, v = numpy.linalg.eig(T)

#get the dimension of the tensor field
d = T.shape[2]

#allocate space for the vector field
V = numpy.zeros([T.shape[0], T.shape[1], 3])

idx = l.argsort()

for di in range(d):
if vector == 'smallest':
b = idx[:, :, 0] == di
elif vector == 'largest':
b = idx[:, :, d-1] == di
else:
b = idx[:, :, 1] == 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 = numpy.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

def anisotropy(S):

Sf = sym2mat(S)

#calculate the eigenvectors and eigenvalues
l, v = numpy.linalg.eig(Sf)

#store the sorted eigenvalues
ls = numpy.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 = numpy.sqrt(3./2.) * numpy.sqrt(fa_num) / numpy.sqrt(fa_den)

return FA, Cl, Cp, Cs

def st2amira(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 = numpy.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 = numpy.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 = numpy.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 = numpy.ones(T.shape[0], T.shape[1])
image = C * numpy.expand_dims(A, 3)

filename = prefix + str(z).zfill(4) + ".bmp"
scipy.misc.imsave(filename, image)
bar.update(z + 1)