ST2D and ST3D

st2d.py

+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Thu Aug 29 11:30:17 2019
+
+@author: vand@dtu.dk
+"""
+
+import numpy as np
+import scipy.ndimage
+
+# STRUCTURE TENSOR 2D
+
+def structure_tensor(image, sigma, rho):
+    """ Structure tensor for 2D image data
+    Arguments:
+        image: a 2D array of size N = rows*columns
+        sigma: a noise scale, structures smaller than sigma will be 
+            removed by smoothing
+        rho: an integration scale giving the size over the neighborhood in 
+            which the orientation is to be analysed
+    Returns:
+        an array with shape (3,N) containing elements of structure tensor 
+            s_xx, s_yy, s_xy ordered acording to image.ravel()
+    Author: vand@dtu.dk, 2019
+    """
+    
+    # computing derivatives (scipy implementation truncates filter at 4 sigma)
+    image = image.astype(np.float);
+
+    Ix = scipy.ndimage.gaussian_filter(image, sigma, order=[1,0], mode='nearest')
+    Iy = scipy.ndimage.gaussian_filter(image, sigma, order=[0,1], mode='nearest')
+    
+    # integrating elements of structure tensor (scipy uses sequence of 1D)
+    Jxx = scipy.ndimage.gaussian_filter(Ix**2, rho, mode='nearest')
+    Jyy = scipy.ndimage.gaussian_filter(Iy**2, rho, mode='nearest')
+    Jxy = scipy.ndimage.gaussian_filter(Ix*Iy, rho, mode='nearest')
+    S = np.vstack((Jxx.ravel(), Jyy.ravel(), Jxy.ravel()));
+    return S
+
+def eig_special(S):
+    """ Eigensolution for symmetric real 2-by-2 matrices
+    Arguments:
+        S: an array with shape (3,N) containing structure tensor
+    Returns:
+        val: an array with shape (2,N) containing sorted eigenvalues
+        vec: an array with shape (2,N) containing eigenvector corresponding
+            to the smallest eigenvalue (the other is orthogonal to the first)
+    Author: vand@dtu.dk, 2019
+    """
+
+    val = 0.5*(S[0]+S[1]+np.outer(np.array([-1,1]), np.sqrt((S[0]-S[1])**2+4*S[2]**2)))
+    vec = np.vstack((-S[2],S[0]-val[0])) # y will be positive
+    aligned = S[2]==0 # dealing with diagonal matrices
+    vec[:,aligned] = 1-np.argsort(S[0:2,aligned], axis=0)
+    vec = vec/np.sqrt(vec[0]**2+vec[1]**2) # normalizing
+    return val, vec
+
+def solve_flow(S):
+    """ Solving 1D optic flow, returns LLS optimal x for flow along y axis
+        ( x is a solution to S[0]*x=S[2] )
+    Arguments:
+        S: an array with shape (3,N) containing 2D structure tensor
+    Returns:
+        x: an array with shape (1,N) containing x components of the flow
+    Author: vand@dtu.dk, 2019
+    """
+    aligned = S[0]==0 # 0 or inf solutions
+    x = np.zeros((1,S.shape[1]))
+    x[0,~aligned] = - S[2,~aligned]/S[0,~aligned]
+    return x # returning shape (1,N) array for consistancy with 3D case
+
+#% HELPING FUNCTIONS FOR VISUALIZATION
+
+def plot_orientations(ax, dim, vec, s = 5):
+    """ Helping function for adding orientation-quiver to the plot.
+    Arguments: plot axes, image shape, orientation, arrow spacing.
+    """
+    vx = vec[0].reshape(dim)
+    vy = vec[1].reshape(dim)
+    xmesh, ymesh = np.meshgrid(np.arange(dim[0]), np.arange(dim[1]), indexing='ij')
+    ax.quiver(ymesh[s//2::s,s//2::s],xmesh[s//2::s,s//2::s],vy[s//2::s,s//2::s],vx[s//2::s,s//2::s],color='r',angles='xy')
+    ax.quiver(ymesh[s//2::s,s//2::s],xmesh[s//2::s,s//2::s],-vy[s//2::s,s//2::s],-vx[s//2::s,s//2::s],color='r',angles='xy')
+
+def polar_histogram(ax, distribution, cmap = 'hsv'):
+    """ Helping function for producing polar histogram.
+    Arguments: plot axes, oriantation distribution, colormap.
+    """
+    N = distribution.size
+    bin_centers_full = (np.arange(2*N)+0.5)*np.pi/N # full circle (360 deg)
+    distribution_full = np.r_[distribution,distribution]/max(distribution) # full normalized distribution
+    x = np.r_[distribution_full*np.cos(bin_centers_full),0]
+    y = np.r_[distribution_full*np.sin(bin_centers_full),0]
+    triangles = np.array([(i, np.mod(i-1,2*N), 2*N) for i in range(2*N)]) # triangles[0] is symmetric over 0 degree
+    triangle_centers_full = (np.arange(2*N))*np.pi/N # a triangle covers area BETWEEN two bin_centers
+    triangle_colors = np.mod(triangle_centers_full, np.pi)/np.pi # from 0 to 1-(1/2N)
+    ax.tripcolor(y, x, triangles, facecolors=triangle_colors, cmap=cmap, vmin = 0.0, vmax = 1.0)
+    ax.set_aspect('equal')
+    ax.set_xlim([-1,1])
+    ax.set_ylim([1,-1])   

st2d_examples.py

+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Thu Aug 29 11:30:17 2019
+
+@author: vand@dtu.dk
+"""
+
+import numpy as np
+import skimage.io
+import matplotlib.pyplot as plt
+import st2d
+    
+#%% ST AND ORIENTATIONS - VISUALIZATION OPTIONS
+plt.close('all')
+
+filename = '../data2D/drawn_fibres_B.png';
+sigma = 0.5
+rho = 2
+
+image = skimage.io.imread(filename)
+S = st2d.structure_tensor(image, sigma, rho)
+val,vec = st2d.eig_special(S)
+    
+# visualization
+figsize = (10,5)
+fig, ax = plt.subplots(1, 5, figsize=figsize, sharex=True, sharey=True)
+ax[0].imshow(image,cmap=plt.cm.gray)
+st2d.plot_orientations(ax[0], image.shape, vec)
+ax[0].set_title('Orientation as arrows')
+orientation_st_rgba = plt.cm.hsv((np.arctan2(vec[1], vec[0])/np.pi).reshape(image.shape))
+ax[1].imshow(plt.cm.gray(image)*orientation_st_rgba)
+ax[1].set_title('Orientation as color on image')
+ax[2].imshow(orientation_st_rgba)
+ax[2].set_title('Orientation as color')
+anisotropy = (1-val[0]/val[1]).reshape(image.shape)
+ax[3].imshow(anisotropy)
+ax[3].set_title('Degree of anisotropy')
+ax[4].imshow(plt.cm.gray(anisotropy)*orientation_st_rgba)
+ax[4].set_title('Orientation and anisotropy')
+plt.show()    
+
+ #%% ST AND ORIENTATIONS - HISTOGRAMS OPTIONS
+
+filename = '../data2D/10X.png';
+sigma = 0.5
+rho = 15
+N = 180 # number of angle bins for orientation histogram
+
+# computation
+image = skimage.io.imread(filename)
+image = np.mean(image[:,:,0:3],axis=2).astype(np.uint8)
+S = st2d.structure_tensor(image, sigma, rho)
+val,vec = st2d.eig_special(S)
+angles = np.arctan2(vec[1], vec[0]) # angles from 0 to pi
+distribution = np.histogram(angles, bins=N, range=(0.0, np.pi))[0]
+    
+# visualization
+figsize = (10,5)
+fig, ax = plt.subplots(1, 2, figsize=figsize, sharex=True, sharey=True)
+ax[0].imshow(image,cmap=plt.cm.gray)
+ax[0].set_title('Input image')
+orientation_st_rgba = plt.cm.hsv((angles/np.pi).reshape(image.shape))
+ax[1].imshow(plt.cm.gray(image)*orientation_st_rgba)
+ax[1].set_title('Orientation as color on image')
+
+fig, ax = plt.subplots(1,2, figsize=figsize)
+bin_centers = (np.arange(N)+0.5)*np.pi/N # halp circle (180 deg)
+colors = plt.cm.hsv(bin_centers/np.pi)
+ax[0].bar(bin_centers, distribution, width = np.pi/N, color = colors)
+ax[0].set_xlabel('angle')
+ax[0].set_xlim([0,np.pi])
+ax[0].set_aspect(np.pi/ax[0].get_ylim()[1])
+ax[0].set_xticks([0,np.pi/2,np.pi])
+ax[0].set_xticklabels(['0','pi/2','pi'])
+ax[0].set_ylabel('count')
+ax[0].set_title('Histogram over angles')
+st2d.polar_histogram(ax[1], distribution)
+ax[1].set_title('Polar histogram')
+plt.show()
+
+ #%% ST AND ORIENTATIONS - HISTOGRAMS OPTIONS
+
+filename = '../data2D/drawn_field.png';
+sigma = 0.5
+rho = 15
+N = 90 # number of angle bins for orientation histogram
+
+# computation
+image = skimage.io.imread(filename)
+image = np.mean(image[:,:,0:3],axis=2).astype(np.uint8)
+S = st2d.structure_tensor(image, sigma, rho)
+val,vec = st2d.eig_special(S)
+angles = np.arctan2(vec[1], vec[0]) # angles from 0 to pi
+distribution = np.histogram(angles, bins=N, range=(0.0, np.pi))[0]
+distribution_weighted = np.histogram(angles, bins=N, range=(0.0, np.pi), weights = (image>175).ravel().astype(np.float))[0]
+    
+# visualization
+figsize = (10,5)
+fig, ax = plt.subplots(1, 2, figsize=figsize, sharex=True, sharey=True)
+ax[0].imshow(image,cmap=plt.cm.gray)
+ax[0].set_title('Input image')
+orientation_st_rgba = plt.cm.hsv((angles/np.pi).reshape(image.shape))
+ax[1].imshow(plt.cm.gray(image)*orientation_st_rgba)
+ax[1].set_title('Orientation as color on image')
+
+fig, ax = plt.subplots(2, 2, figsize=figsize)
+bin_centers = (np.arange(N)+0.5)*np.pi/N # halp circle (180 deg)
+colors = plt.cm.hsv(bin_centers/np.pi)
+ax[0][0].bar(bin_centers, distribution, width = np.pi/N, color = colors)
+ax[0][0].set_xlabel('angle')
+ax[0][0].set_xlim([0,np.pi])
+ax[0][0].set_aspect(np.pi/ax[0][0].get_ylim()[1])
+ax[0][0].set_xticks([0,np.pi/2,np.pi])
+ax[0][0].set_xticklabels(['0','pi/2','pi'])
+ax[0][0].set_ylabel('count')
+ax[0][0].set_title('Histogram over angles - all orientations')
+st2d.polar_histogram(ax[0][1], distribution)
+ax[0][1].set_title('Polar histogram - all orientations')
+ax[1][0].bar(bin_centers, distribution_weighted, width = np.pi/N, color = colors)
+ax[1][0].set_xlabel('angle')
+ax[1][0].set_xlim([0,np.pi])
+ax[1][0].set_aspect(np.pi/ax[1][0].get_ylim()[1])
+ax[1][0].set_xticks([0,np.pi/2,np.pi])
+ax[1][0].set_xticklabels(['0','pi/2','pi'])
+ax[1][0].set_ylabel('count')
+ax[1][0].set_title('Histogram over angles - orientations at fibres')
+st2d.polar_histogram(ax[1][1], distribution_weighted)
+ax[1][1].set_title('Polar histogram - orientations at fibres')
+plt.show()
+
+   #%% YET ANOTHER EXAMPLE 
+
+filename = '../data2D/OCT_im_org.png';
+sigma = 0.5
+rho = 5
+N = 180 # number of angle bins for orientation histogram
+
+# computation
+image = skimage.io.imread(filename)
+S = st2d.structure_tensor(image, sigma, rho)
+val,vec = st2d.eig_special(S)
+angles = np.arctan2(vec[1], vec[0]) # angles from 0 to pi
+distribution = np.histogram(angles, bins=N, range=(0.0, np.pi))[0]
+    
+# visualization
+figsize = (10,5)
+fig, ax = plt.subplots(1, 2, figsize=figsize, sharex=True, sharey=True)
+ax[0].imshow(image,cmap=plt.cm.gray)
+ax[0].set_title('Input image')
+orientation_st_rgba = plt.cm.hsv((angles/np.pi).reshape(image.shape))
+ax[1].imshow(plt.cm.gray(image)*orientation_st_rgba)
+ax[1].set_title('Orientation as color on image')
+
+fig, ax = plt.subplots(1,2, figsize=figsize)
+bin_centers = (np.arange(N)+0.5)*np.pi/N # halp circle (180 deg)
+colors = plt.cm.hsv(bin_centers/np.pi)
+ax[0].bar(bin_centers, distribution, width = np.pi/N, color = colors)
+ax[0].set_xlabel('angle')
+ax[0].set_xlim([0,np.pi])
+ax[0].set_aspect(np.pi/ax[0].get_ylim()[1])
+ax[0].set_xticks([0,np.pi/2,np.pi])
+ax[0].set_xticklabels(['0','pi/2','pi'])
+ax[0].set_ylabel('count')
+ax[0].set_title('Histogram over angles')
+st2d.polar_histogram(ax[1], distribution)
+ax[1].set_title('Polar histogram')
+plt.show()
+    
+#%% INVESTIGATING THE EFFECT OF RHO
+
+filename = '../data2D/short_fibres.png'
+image = skimage.io.imread(filename)
+image = np.mean(image[:,:,0:3],axis=2)
+image -= np.min(image)
+image /= np.max(image)
+s = 128 # quiver arrow spacing
+sigma = 0.5
+figsize = (10,5)
+
+rhos = [2,10,20,50]
+
+for k in range(4):
+    
+    # computation
+    rho = rhos[k] # changing integration radius
+    S = st2d.structure_tensor(image,sigma,rho)
+    val,vec = st2d.eig_special(S)
+        
+    # visualization
+    fig, ax = plt.subplots(1, 2, figsize=figsize, sharex=True, sharey=True)
+    ax[0].imshow(image,cmap=plt.cm.gray)
+    st2d.plot_orientations(ax[0], image.shape, vec, s = s)
+    ax[0].set_title(f'Rho = {rho}, arrows')
+    intensity_rgba = plt.cm.gray(image)
+    orientation_st_rgba = plt.cm.hsv((np.arctan2(vec[1], vec[0])/np.pi).reshape(image.shape))
+    ax[1].imshow((0.5+0.5*intensity_rgba)*orientation_st_rgba)
+    ax[1].set_title(f'Rho = {rho}, color')
+
+    #basis_filename = filename.split('/')[-1].split('.')[0]
+    #fig.savefig(basis_filename + '_rho_' + str(rho) + '.png')
+    plt.show()
+
+
+
+#%% INVESTIGATING THE EFFECT OF SCALING + RHO 
+
+filename = '../data2D/short_fibres.png'
+downsampling_range = 4
+figsize = (10,5)    
+
+for k in range(downsampling_range):
+
+    # downsampling and computation
+    scale = 2**k
+    f = 1/scale # image scale factor
+    s = 128//scale # quiver arrow spacing
+    sigma = 0.5 # would it make sense to scale this too?
+    rho = 50/scale # scaling the integration radius
+    image = skimage.io.imread(filename)
+    image = np.mean(image[:,:,0:3],axis=2)
+    image = skimage.transform.rescale(image,f,multichannel=False)
+    image -= np.min(image)
+    image /= np.max(image)
+    S = st2d.structure_tensor(image,sigma,rho)
+    val,vec = st2d.eig_special(S)
+        
+    # visualization
+    fig, ax = plt.subplots(1, 2, figsize=figsize, sharex=True, sharey=True)
+    ax[0].imshow(image,cmap=plt.cm.gray)
+    st2d.plot_orientations(ax[0], image.shape, vec, s = s)
+    ax[0].set_title(f'Downsample = {scale}, arrows')
+    intensity_rgba = plt.cm.gray(image)
+    orientation_st_rgba = plt.cm.hsv((np.arctan2(vec[1], vec[0])/np.pi).reshape(image.shape))
+    ax[1].imshow((0.5+0.5*intensity_rgba)*orientation_st_rgba)       
+    ax[1].set_title(f'Downsample = {scale}, color')
+ 
+    #basis_filename = filename.split('/')[-1].split('.')[0]
+    #fig.savefig(basis_filename + '_scale_' + str(scale) + '.png')
+    plt.show()
+    
+#%% COMPARING DOMINANT ORIENTATION AND OPTICAL FLOW
+image = skimage.io.imread('../data2D/drawn_fibres_B.png');
+
+# computing structure tensor, orientation and optical flow
+sigma = 0.5
+rho = 5
+S = st2d.structure_tensor(image,sigma,rho)
+val,vec = st2d.eig_special(S) # dominant orientation
+fx = st2d.solve_flow(S) # optical flow
+
+# visualization
+figsize = (10,10)
+fy = np.ones(image.shape)
+fig, ax = plt.subplots(2,2,figsize=figsize)
+
+ax[0][0].imshow(image,cmap=plt.cm.gray)
+st2d.plot_orientations(ax[0][0], image.shape, vec)
+ax[0][0].set_title('Orientation from structure tensor, arrows')
+ax[0][1].imshow(image,cmap=plt.cm.gray)
+st2d.plot_orientations(ax[0][1], image.shape, np.r_[fx,np.ones((1,image.size))])
+ax[0][1].set_title('Orientation from optical flow, arrows')
+intensity_rgba = plt.cm.gray(image)
+orientation_st_rgba = plt.cm.hsv((np.arctan2(vec[1],vec[0])/np.pi).reshape(image.shape))
+orientation_of_rgba = plt.cm.hsv((np.arctan2(1,fx)/np.pi).reshape(image.shape))
+ax[1][0].imshow(intensity_rgba*orientation_st_rgba)
+ax[1][0].set_title('Dominant orientation from structure tensor, color')
+ax[1][1].imshow(intensity_rgba*orientation_of_rgba)
+ax[1][1].set_title('Orientation from optical flow, color')
+plt.show()
\ No newline at end of file

st3d.py

+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Thu Aug 29 11:30:17 2019
+
+@author: vand@dtu.dk
+"""
+
+import numpy as np
+import scipy.io
+import scipy.ndimage
+import matplotlib.pyplot as plt
+
+#% STRUCTURE TENSOR 3D
+
+def structure_tensor(volume, sigma, rho):
+    """ Structure tensor for 3D image data
+    Arguments:
+        volume: a 3D array of size N = slices(z)*rows(y)*columns(x)
+        sigma: a noise scale, structures smaller than sigma will be 
+            removed by smoothing
+        rho: an integration scale giving the size over the neighborhood in 
+            which the orientation is to be analysed
+    Returns:
+        an array with shape (6,N) containing elements of structure tensor 
+            s_xx, s_yy, s_zz, s_xy, s_xz, s_yz ordered acording to 
+            volume.ravel(). 
+    Author: vand@dtu.dk, 2019
+    """    # computing derivatives (scipy implementation truncates filter at 4 sigma)
+    volume = volume.astype(np.float);
+    
+    Vx = scipy.ndimage.gaussian_filter(volume, sigma, order=[0,0,1], mode='nearest')
+    Vy = scipy.ndimage.gaussian_filter(volume, sigma, order=[0,1,0], mode='nearest')
+    Vz = scipy.ndimage.gaussian_filter(volume, sigma, order=[1,0,0], mode='nearest')
+  
+    # integrating elements of structure tensor (scipy uses sequence of 1D)
+    Jxx = scipy.ndimage.gaussian_filter(Vx**2, rho, mode='nearest')
+    Jyy = scipy.ndimage.gaussian_filter(Vy**2, rho, mode='nearest')
+    Jzz = scipy.ndimage.gaussian_filter(Vz**2, rho, mode='nearest')
+    Jxy = scipy.ndimage.gaussian_filter(Vx*Vy, rho, mode='nearest')
+    Jxz = scipy.ndimage.gaussian_filter(Vx*Vz, rho, mode='nearest')
+    Jyz = scipy.ndimage.gaussian_filter(Vy*Vz, rho, mode='nearest')
+    S = np.vstack((Jxx.ravel(), Jyy.ravel(), Jzz.ravel(), Jxy.ravel(),\
+                   Jxz.ravel(), Jyz.ravel()));
+    return S
+
+def eig_special(S, full=False):
+    """ Eigensolution for symmetric real 3-by-3 matrices
+    Arguments:
+        S: an array with shape (6,N) containing structure tensor
+        full: a flag indicating that all three eigenvalues should be returned
+    Returns:
+        val: an array with shape (3,N) containing sorted eigenvalues
+        vec: an array with shape (3,N) containing eigenvector corresponding to 
+            the smallest eigenvalue. If full, vec has shape (6,N) and contains 
+            all three eigenvectors 
+    More:        
+        An analytic solution of eigenvalue problem for real symmetric matrix,
+        using an affine transformation and a trigonometric solution of third
+        order polynomial. See https://en.wikipedia.org/wiki/Eigenvalue_algorithm
+        which refers to Smith's algorithm https://dl.acm.org/citation.cfm?id=366316
+    Author: vand@dtu.dk, 2019
+    """    
+    # TODO -- deal with special cases, decide treatment of full (i.e. maybe return 2 for full)
+    # computing eigenvalues
+    s = S[3]**2 + S[4]**2 + S[5]**2 # off-diagonal elements
+    q = (1/3)*(S[0]+S[1]+S[2]) # mean of on-diagonal elements
+    p = np.sqrt((1/6)*(np.sum((S[0:3] - q)**2, axis=0) + 2*s)) # case p==0 treated below 
+    p_inv = np.zeros(p.shape)
+    p_inv[p!=0] = 1/p[p!=0] # to avoid division by 0
+    B = p_inv * (S - np.outer(np.array([1,1,1,0,0,0]),q))  # B represents a 3-by-3 matrix, A = pB+2I   
+    d = B[0]*B[1]*B[2] + 2*B[3]*B[4]*B[5] - B[3]**2*B[2]\
+            - B[4]**2*B[1] - B[5]**2*B[0] # determinant of B
+    phi = np.arccos(np.minimum(np.maximum(d/2,-1),1))/3 # min-max to ensure -1 <= d/2 <= 1 
+    val = q + 2*p*np.cos(phi.reshape((1,-1))+np.array([[2*np.pi/3],[4*np.pi/3],[0]])) # ordered eigenvalues
+
+    # computing eigenvectors -- either only one or all three
+    if full:
+        l = val
+    else:
+        l=val[0]
+            
+    u = S[4]*S[5]-(S[2]-l)*S[3]
+    v = S[3]*S[5]-(S[1]-l)*S[4]
+    w = S[3]*S[4]-(S[0]-l)*S[5]
+    vec = np.vstack((u*v, u*w, v*w)) # contains one or three vectors
+   
+    # normalizing -- depends on number of vectors
+    if full: # vec is [x1 x2 x3 y1 y2 y3 z1 z2 z3]
+        vec = vec[[0,3,6,1,4,7,2,5,8]] # vec is [v1, v2, v3]
+        l = np.sqrt(np.vstack((np.sum(vec[0:3]**2,axis=0), np.sum(vec[3:6]**2,\
+                axis=0), np.sum(vec[6:]**2, axis=0))))
+        vec = vec/l[[0,0,0,1,1,1,2,2,2]] # division by 0 should not occur
+    else: # vec is [x1 y1 z1] = v1
+        vec = vec/np.sqrt(np.sum(vec**2, axis=0));
+    return val,vec
+
+def solve_flow(S): 
+    """ Solving 2D optic flow, returns LLS optimal x and y for flow along z axis
+        (A solution of a 2x2 linear system.)
+    Arguments:
+        S: an array with shape (6,N) containing 3D structure tensor
+    Returns:
+        xy: an array with shape (2,N) containing x and y components of the flow
+    Author: vand@dtu.dk, 2019
+    """
+    d = S[0]*S[1]-S[3]**2 # denominator
+    aligned = d==0 # 0 or inf solutions
+    n = np.vstack((S[3]*S[5]-S[1]*S[4],S[3]*S[4]-S[0]*S[5]))
+    xy = np.zeros((2,S.shape[1]))
+    xy[:,~aligned] = n[:,~aligned]/d[~aligned]
+    return xy 
+
+def tensor_vector_distance(S, u):
+    """ Caclulating pairwise distance between tensors and vectors
+    Arguments:
+        S: an array with shape (6,N) containing tensor
+        v: an array with shape (M,3) containing vectors
+    Returns:
+        v: an array with shape (N,M) containing pairwise distances
+    Author: vand@dtu.dk, 2019
+    """
+    dist = np.dot(S[0:3].T,u**2) + 2*np.dot(S[3:].T,u[[0,0,1]]*u[[1,2,2]])
+    return dist
+    
+#% INTERACTIVE VISUALIZATION FUNCTIONS - DOES NOT WORK WITH INLINE FIGURES
+
+def arrow_navigation(event,z,Z):
+    if event.key == "up":
+        z = min(z+1,Z-1)
+    elif event.key == 'down':
+        z = max(z-1,0)
+    elif event.key == 'right':
+        z = min(z+10,Z-1)
+    elif event.key == 'left':
+        z = max(z-10,0)
+    elif event.key == 'pagedown':
+        z = min(z+50,Z+1)
+    elif event.key == 'pageup':
+        z = max(z-50,0)
+    return z
+
+def show_vol(V,cmap='gray'): 
+    """
+    Shows volumetric data and colored orientation for interactive inspection.
+    @author: vand at dtu dot dk
+    """
+    def update_drawing():
+        ax.images[0].set_array(V[z])
+        ax.set_title(f'slice z={z}')
+        fig.canvas.draw()
+ 
+    def key_press(event):
+        nonlocal z
+        z = arrow_navigation(event,z,Z)
+        update_drawing()
+        
+    Z = V.shape[0]
+    z = (Z-1)//2
+    fig, ax = plt.subplots()
+    vmin = np.min(V)
+    vmax = np.max(V)
+    ax.imshow(V[z], cmap=cmap, vmin=vmin, vmax=vmax)
+    ax.set_title(f'slice z={z}')
+    fig.canvas.mpl_connect('key_press_event', key_press)
+
+def show_vol_flow(V, fxy, s=5, double_arrow = False): 
+    """
+    Shows volumetric data and xy optical flow for interactive inspection.
+    Arguments:
+         V: volume
+         fxy: flow in x and y direction
+         s: spacing of quiver arrows
+    @author: vand at dtu dot dk
+    """
+    def update_drawing():
+        ax.images[0].set_array(V[z])
+        ax.collections[0].U = fxy[0,z,s//2::s,s//2::s].ravel()
+        ax.collections[0].V = fxy[1,z,s//2::s,s//2::s].ravel()
+        if double_arrow:
+            ax.collections[1].U = -fxy[0,z,s//2::s,s//2::s].ravel()
+            ax.collections[1].V = -fxy[1,z,s//2::s,s//2::s].ravel()
+        ax.set_title(f'slice z={z}')
+        fig.canvas.draw()
+ 
+    def key_press(event):
+        nonlocal z
+        z = arrow_navigation(event,z,Z)
+        update_drawing()
+        
+    Z = V.shape[2]
+    z = (Z-1)//2
+    xmesh, ymesh = np.meshgrid(np.arange(V.shape[1]), np.arange(V.shape[2]), indexing='ij')
+            # TODO: figure out exactly why this ij later needs 'xy' 
+    fig, ax = plt.subplots()
+    ax.imshow(V[z],cmap='gray')
+    ax.quiver(ymesh[s//2::s,s//2::s], xmesh[s//2::s,s//2::s],
+              fxy[0,z,s//2::s,s//2::s], fxy[1,z,s//2::s,s//2::s],
+              color='r', angles='xy')
+    if double_arrow:
+        ax.quiver(ymesh[s//2::s,s//2::s], xmesh[s//2::s,s//2::s],
+          -fxy[0,z,s//2::s,s//2::s], -fxy[1,z,s//2::s,s//2::s],
+          color='r', angles='xy')
+    ax.set_title(f'slice z={z}')
+    fig.canvas.mpl_connect('key_press_event', key_press)
+    
+def fan_coloring(vec):
+    """
+    Fan-based colors for orientations in xy plane
+    Arguments:
+        vec: an array with shape (3,N) containing orientations
+    Returns:
+        rgba: an array with shape (4,N) containing rgba colors
+     @author:vand@dtu.dk
+    """
+    h = (vec[2]**2).reshape((vec.shape[1],1)) # no discontinuity and less gray
+    s = np.mod(np.arctan(vec[0]/vec[1]),np.pi) # hue angle
+    hue = plt.cm.hsv(s/np.pi)
+    rgba = hue*(1-h) + 0.5*h
+    rgba[:,3] = 1 # fixing alpha value
+    return rgba
+    
+def show_vol_orientation(V, vec, 
+                         coloring = lambda v : np.c_[abs(v).T,np.ones((v.shape[1],1))], 
+                         blending = lambda g,c : 0.5*(g+c)): 
+    """
+    Shows volumetric data and colored orientation for interactive inspection.
+    @author: vand at dtu dot dk
+    """
+    rgba = coloring(vec).reshape(V.shape+(4,))
+    
+    def update_drawing():
+        ax.images[0].set_array(blending(plt.cm.gray(V[z]), rgba[z]))
+        ax.set_title(f'slice z={z}')
+        fig.canvas.draw()
+ 
+    def key_press(event):
+        nonlocal z
+        z = arrow_navigation(event,z,Z)
+        update_drawing()
+        
+    Z = V.shape[0]
+    z = (Z-1)//2
+    fig, ax = plt.subplots()
+    ax.imshow(blending(plt.cm.gray(V[z]), rgba[z]))
+    ax.set_title(f'slice z={z}')
+    fig.canvas.mpl_connect('key_press_event', key_press)
+    
\ No newline at end of file

st3d_examples.py

+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Thu Aug 29 11:30:17 2019
+
+@author: vand@dtu.dk
+"""
+
+import numpy as np
+import scipy.io
+import scipy.ndimage
+import matplotlib.pyplot as plt
+import st3d
+
+# reading in the data
+volume = scipy.io.loadmat('../testing_data_3D/multi_cube.mat')['vol']
+
+# computing structure tensor and orientations
+sigma = 0.5;
+rho = 2;
+S = st3d.structure_tensor(volume, sigma, rho)
+val,vec = st3d.eig_special(S)
+    
+# visualization - use arrow keys to inspect the volume
+st3d.show_vol_flow(volume, vec[0:2].reshape((2,)+volume.shape), s=10, double_arrow = True)
+st3d.show_vol_orientation(volume, vec, coloring = st3d.fan_coloring)
+
+    
+#%% AND FOR A ROTATED SAMPLE...
+volume = np.transpose(volume, (1, 0, 2))
+S = st3d.structure_tensor(volume, sigma, rho)
+val,vec = st3d.eig_special(S)
+coloring_z = lambda v : st3d.fan_coloring(v[[2, 0, 1]]) # rotating coloring
+st3d.show_vol_orientation(volume, vec, coloring = coloring_z)
+
+#%% DISTANCE FROM A DIRECTION
+u = np.array([0,0,1])
+dist = st3d.tensor_vector_distance(S,u)
+st3d.show_vol(dist.reshape(volume.shape))
+ 
+# #%% FLOW MAKES SENSE ONLY WITH ROUGHLY UD FIBERS
+# fxy = solve_flow_2D(S).reshape((2,)+volume.shape)
+
+#%% K-MEANS LIKE CLUSTERING, INITILIZED BY 4 DIRECTIONS
+t = np.sqrt(1/2)
+u_clusters = np.array([[1,0,0],[0,0,1],[t,0,t],[-t,0,t]]).T # initalizing close to desired solution   
+dist = st3d.tensor_vector_distance(S,u_clusters)
+assignment = np.argmin(dist,axis = 1)
+ 
+S_clusters = np.zeros((6,u_clusters.shape[1]))
+for r in range(10): # iterations of k-means
+    for i in range(u_clusters.shape[1]): # collecting ST for all voxels in a cluster
+        S_clusters[:,i] = np.mean(S[:,assignment==i],axis=1)
+    val,vec = st3d.eig_special(S_clusters) # estimating new cluster orientation
+    print(f'Iter {r}: moved for {np.sum(u_clusters-vec)}')
+    u_clusters = vec
+    dist = st3d.tensor_vector_distance(S,u_clusters)
+    assignment = np.argmin(dist,axis = 1)
+
+st3d.show_vol(assignment.reshape(volume.shape), cmap='jet')
+
+
+   
+