Skip to content
Snippets Groups Projects
Select Git revision
  • master
  • bepi
2 results

01_JupyterNotebook-checkpoint.ipynb

Blame
  • Code owners
    Assign users and groups as approvers for specific file changes. Learn more.
    ST3D_examples.py 4.48 KiB
    # -*- coding: utf-8 -*-
    """
    Created on Thu Mar 18 14:10:15 2021
    
    """
    
    import numpy as np
    import scipy.io
    import matplotlib.pyplot as plt
    #test for GPU
    flag_GPU = 1
    try:
        import cupy as cp
    except:
        flag_GPU = 0
    #import 3D structure tensor and utils (helper functions)
    if not(flag_GPU):
        from structure_tensor import eig_special_3d, structure_tensor_3d    #CPU version
    else:
        from structure_tensor.cp import eig_special_3d, structure_tensor_3d #GPU version      
    import utilsST_3D
    
    
    
    volume = scipy.io.loadmat('example_data_3D/multi_cube.mat')['vol']
    print(volume.shape)
    
    
    sigma = 0.5; # noise scale
    rho = 2; # integration scale
    S = structure_tensor_3d(volume, sigma, rho)
    val,vec = eig_special_3d(S)
    print(f'The volume has a shape {volume.shape}, i.e. {volume.size} voxels.')
    print(f'Structure tensor information is carried in a {S.shape} array.')
    print(f'Orientation information is carried in a {vec.shape} array.')
    
    
    utilsST_3D.show_vol_flow(volume, vec[0:2], s=10, double_arrow = True) 
    
    utilsST_3D.show_vol_orientation(volume, 
                                    vec, 
                                    coloring = utilsST_3D.fan_coloring)
    
    # Rotating sample:
    volume = np.transpose(volume, (1, 0, 2))
    S = structure_tensor_3d(volume, sigma, rho)
    val,vec = eig_special_3d(S)
    coloring_z = lambda v : utilsST_3D.fan_coloring(v[[2, 0, 1]]) # rotating the coloring function
    utilsST_3D.show_vol_orientation(volume, vec, coloring = coloring_z)
    
    # -- Investigating tensor-vector distance -- #
    #A tensor-vector distance can be computed for every voxel. Here, the distance is computed from the z-direction. In the rotated volume, the 90 deg. bundle (middle layer) is aligned with z-direction, and has the smallest distance.
    u = np.array([0,0,1])
    dist = utilsST_3D.tensor_vector_distance(S,u)
    utilsST_3D.show_vol(dist)
    
    # -- Visualise eigen vectors -- #
    
    #in the cartesian space
    plt.figure(figsize=(6,6))
    plt.subplot(2,2,1)
    plt.hist2d(vec[2,:].flatten(), vec[0,:].flatten(), bins=50)
    plt.axis('scaled'), plt.xlabel('x'), plt.ylabel('z'), 
    plt.subplot(2,2,2), 
    plt.hist2d(vec[1,:].flatten(), vec[0,:].flatten(), bins=50)
    plt.axis('scaled'), plt.xlabel('y'), plt.ylabel('z'), 
    plt.subplot(2,2,3)
    plt.hist2d(vec[2,:].flatten(), vec[1,:].flatten(), bins=50)
    plt.axis('equal'), plt.xlabel('x'), plt.ylabel('y'), 
    plt.show()
    
    #in the angular space (spherical histogram azimuth and elevation)
    nBin = (128,128)
    binVal, binC_az, binC_ele = utilsST_3D.histogramSphere(np.reshape(vec,[3,np.prod(vec.shape[1:])]), nBin)
    
    fig, ax = plt.subplots(1,1,figsize=(6,6))
    ax.imshow(binVal.T, cmap='gray',extent=[-np.pi,np.pi,-np.pi/2,np.pi/2])
    ax.set_xticks([-np.pi,-np.pi/2,0,np.pi/2,np.pi])
    ax.set_xticklabels(['-pi','-pi/2','0','pi/2','pi'])
    ax.set_yticks([-np.pi/2,0,np.pi/2])
    ax.set_yticklabels(['-pi/2','0','pi/2'])
    ax.set_title('Spherical histogram over angles')
    ax.set_xlabel('Azimuth')
    ax.set_ylabel('Inclination')
    
    # -- K-means like clustering of structure tensor using tensor-vector distance -- #
    #The tensor-vector distance allows tensor clustering using an approach similar to k-means. Here, every cluster is characterized by an orientation vector, and consists of the tensors which have a small distance to this orientation. The advantage of the approach is that it operates on tensors, and does not require their eigen decomposition. Only the computation of cluster orientations will require eigen decomposition.
    t = np.sqrt(1/2)
    u_clusters = np.array([[1,0,0],[0,0,1],[t,0,t],[-t,0,t]]).T # initalizing orientations close to desired solution   
    dist = utilsST_3D.tensor_vector_distance(S,u_clusters)
    assignment = np.argmin(dist,axis = -1)
    
    S_clusters = np.zeros((6,u_clusters.shape[1]))
    if flag_GPU:
        if type(S) == cp.core.core.ndarray:
            S = S.get()
    
    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 = eig_special_3d(S_clusters) # estimating new cluster orientation
        if flag_GPU:
            if type(val) == cp.core.core.ndarray:
                val = val.get()
            if type(vec) == cp.core.core.ndarray:
                vec = vec.get()
        print(f'Iter {r}: moved cluster centers for for {np.sqrt(np.sum((u_clusters-vec)**2))}')
        u_clusters = vec
        dist = utilsST_3D.tensor_vector_distance(S,u_clusters)
        assignment_new = np.argmin(dist,axis = -1)
        print(f'Iter {r}: moved {np.sum(abs(assignment-assignment_new))} voxels')
        assignment = assignment_new
    
    utilsST_3D.show_vol(assignment.reshape(volume.shape), cmap='jet')