codebase / stimlib

Commit 7ce7cd7b6f828d6d357c67e2fdf697f201ac3896

Authored by David Mayerich
1 parent fddfdc02

added structure tensor functions to python/

Showing 1 changed file with 232 additions and 0 deletions   Show diff stats
Inline Side-by-side
python/structen.py 0 → 100644
  Show/Hide comments   View file @7ce7cd7
  1 +# -*- coding: utf-8 -*-
  2 +"""
  3 +Created on Sun Mar 12 21:54:40 2017
  4 +
  5 +@author: david
  6 +"""
  7 +
  8 +import numpy
  9 +import scipy.ndimage
  10 +import progressbar
  11 +import glob
  12 +
  13 +def st2(I, s=1, dtype=numpy.float32):
  14 +
  15 +
  16 + #calculates the 2D structure tensor for an image using a gaussian window with standard deviation s
  17 +
  18 + #calculate the gradient
  19 + dI = numpy.gradient(I)
  20 +
  21 + #calculate the dimensions of the tensor field
  22 + field_dims = [dI[0].shape[0], dI[0].shape[1], 3]
  23 +
  24 + #allocate space for the tensor field
  25 + Tg = numpy.zeros(field_dims, dtype=dtype)
  26 +
  27 + #calculate the gradient components of the tensor
  28 + ti = 0
  29 + for i in range(2):
  30 + for j in range(i + 1):
  31 + Tg[:, :, ti] = dI[j] * dI[i]
  32 + ti = ti + 1
  33 +
  34 + #blur the tensor field
  35 + T = numpy.zeros(field_dims, dtype=dtype)
  36 +
  37 + for i in range(3):
  38 + T[:, :, i] = scipy.ndimage.filters.gaussian_filter(Tg[:, :, i], [s, s])
  39 +
  40 +
  41 + return T
  42 +
  43 +def st3(I, s=1):
  44 + #calculate the structure tensor field for the 3D input image I given the window size s in 3D
  45 + #check the format for the window size
  46 + if type(s) is not list:
  47 + s = [s] * 3
  48 + elif len(s) == 1:
  49 + s = s * 3
  50 + elif len(s) == 2:
  51 + s.insert(1, s[0])
  52 +
  53 + print("\nCalculating gradient...")
  54 + dI = numpy.gradient(I)
  55 + #calculate the dimensions of the tensor field
  56 + field_dims = [dI[0].shape[0], dI[0].shape[1], dI[0].shape[2], 6]
  57 +
  58 + #allocate space for the tensor field
  59 + Tg = numpy.zeros(field_dims, dtype=numpy.float32)
  60 +
  61 + #calculate the gradient components of the tensor
  62 + ti = 0
  63 + print("Calculating tensor components...")
  64 + bar = progressbar.ProgressBar()
  65 + bar.max_value = 6
  66 + for i in range(3):
  67 + for j in range(i + 1):
  68 + Tg[:, :, :, ti] = dI[j] * dI[i]
  69 + ti = ti + 1
  70 + bar.update(ti)
  71 +
  72 + #blur the tensor field
  73 + T = numpy.zeros(field_dims, dtype=numpy.float32)
  74 +
  75 + print("\nConvolving tensor field...")
  76 + bar = progressbar.ProgressBar()
  77 + bar.max_value = 6
  78 + for i in range(6):
  79 + T[:, :, :, i] = scipy.ndimage.filters.gaussian_filter(Tg[:, :, :, i], s)
  80 + bar.update(i+1)
  81 +
  82 + return T
  83 +
  84 +def st(I, s=1):
  85 + if I.ndim == 3:
  86 + return st3(I, s)
  87 + elif I.ndim == 2:
  88 + return st2(I, s)
  89 + else:
  90 + print("Image must be 2D or 3D")
  91 + return
  92 +
  93 +
  94 +
  95 +def sym2mat(T):
  96 + #Calculate the full symmetric matrix from a list of lower triangular elements.
  97 + #The lower triangular components in the input field T are assumed to be the
  98 + # final dimension of the input matrix.
  99 +
  100 + # | 1 2 4 7 |
  101 + # | 0 3 5 8 |
  102 + # | 0 0 6 9 |
  103 + # | 0 0 0 10 |
  104 +
  105 + in_s = T.shape
  106 +
  107 + #get the number of tensor elements
  108 + n = in_s[T.ndim - 1]
  109 +
  110 + #calculate the dimension of the symmetric matrix
  111 + d = int(0.5 * (numpy.sqrt(8. * n + 1.) - 1.))
  112 +
  113 + #calculate the dimensions of the output field
  114 + out_s = list(in_s)[:-1] + [d] + [d]
  115 +
  116 + #allocate space for the output field
  117 + R = numpy.zeros(out_s)
  118 +
  119 + ni = 0
  120 + for i in range(d):
  121 + for j in range(i + 1):
  122 + R[..., i, j] = T[..., ni]
  123 + if i != j:
  124 + R[..., j, i] = T[..., ni]
  125 + ni = ni + 1
  126 +
  127 + return R
  128 +
  129 +def st2vec(S, vector='largest'):
  130 + #Create a color image from a 2D or 3D structure tensor slice
  131 +
  132 + #convert the field to a full rank-2 tensor
  133 + T = sym2mat(S);
  134 + del(S)
  135 +
  136 + #calculate the eigenvectors and eigenvalues
  137 + l, v = numpy.linalg.eig(T)
  138 +
  139 + #get the dimension of the tensor field
  140 + d = T.shape[2]
  141 +
  142 + #allocate space for the vector field
  143 + V = numpy.zeros([T.shape[0], T.shape[1], 3])
  144 +
  145 + idx = l.argsort()
  146 +
  147 + for di in range(d):
  148 + if vector == 'smallest':
  149 + b = idx[:, :, 0] == di
  150 + elif vector == 'largest':
  151 + b = idx[:, :, d-1] == di
  152 + else:
  153 + b = idx[:, :, 1] == di
  154 + V[b, 0:d] = v[b, :, di]
  155 +
  156 + return V
  157 +
  158 +def loadstack(filemask):
  159 + #Load an image stack as a 3D grayscale array
  160 +
  161 + #get a list of all files matching the given mask
  162 + files = [file for file in glob.glob(filemask)]
  163 +
  164 + #calculate the size of the output stack
  165 + I = scipy.misc.imread(files[0])
  166 + X = I.shape[0]
  167 + Y = I.shape[1]
  168 + Z = len(files)
  169 +
  170 + #allocate space for the image stack
  171 + M = numpy.zeros([X, Y, Z]).astype('float32')
  172 +
  173 + #create a progress bar
  174 + bar = progressbar.ProgressBar()
  175 + bar.max_value = Z
  176 +
  177 + #for each file
  178 + for z in range(Z):
  179 + #load the file and save it to the image stack
  180 + M[:, :, z] = scipy.misc.imread(files[z], flatten="True").astype('float32')
  181 + bar.update(z+1)
  182 + return M
  183 +
  184 +def anisotropy(S):
  185 +
  186 + Sf = sym2mat(S)
  187 +
  188 + #calculate the eigenvectors and eigenvalues
  189 + l, v = numpy.linalg.eig(Sf)
  190 +
  191 + #store the sorted eigenvalues
  192 + ls = numpy.sort(l)
  193 + l0 = ls[:, :, 0]
  194 + l1 = ls[:, :, 1]
  195 + l2 = ls[:, :, 2]
  196 +
  197 + #calculate the linear anisotropy
  198 + Cl = (l2 - l1)/(l2 + l1 + l0)
  199 +
  200 + #calculate the planar anisotropy
  201 + Cp = 2 * (l1 - l0) / (l2 + l1 + l0)
  202 +
  203 + #calculate the spherical anisotropy
  204 + Cs = 3 * l0 / (l2 + l1 + l0)
  205 +
  206 + #calculate the fractional anisotropy
  207 + l_hat = (l0 + l1 + l2)/3
  208 + fa_num = (l2 - l_hat) ** 2 + (l1 - l_hat) ** 2 + (l0 - l_hat) ** 2;
  209 + fa_den = l0 ** 2 + l1 ** 2 + l2 ** 2
  210 + FA = numpy.sqrt(3./2.) * numpy.sqrt(fa_num) / numpy.sqrt(fa_den)
  211 +
  212 + return FA, Cl, Cp, Cs
  213 +
  214 +def st2amira(filename, T):
  215 + #generates a tensor field that can be imported into Amira
  216 +
  217 + # 0 dx dx ----> 0
  218 + # 1 dx dy ----> 1
  219 + # 2 dy dy ----> 3
  220 + # 3 dx dz ----> 2
  221 + # 4 dy dz ----> 4
  222 + # 5 dz dz ----> 5
  223 +
  224 + #swap the 2nd and 3rd tensor components
  225 + temp = T[..., 3]
  226 + T[..., 3] = T[..., 2]
  227 + T[..., 2] = temp;
  228 +
  229 + #roll the tensor axis so that it is the leading component
  230 + A = numpy.rollaxis(T, T.ndim - 1)
  231 + A.tofile(filename)
  232 + print(A.shape)
0 233 \ No newline at end of file
... ...