import numpy as np

import nibabel as nib

from dipy.viz import window, actor

import dipy.data as dpd

import dipy.reconst.dti as dti

import matplotlib.pyplot as plt

import math

from sklearn.metrics.pairwise import cosine_similarity

from scipy.ndimage.filters import median_filter

from tqdm import tqdm

dmri = nib.load("Data/dmri.nii",mmap=False)

dmri_data =dmri.get_data()

# =============================================================================

# estimation des tenseurs

# =============================================================================

b=np.loadtxt(fname = "Data/gradient_directions_b-values.txt")

b_vectors=b[1:,0:3]

b_value=b[1,3]

#calcul de matrice B

B=np.zeros((b_vectors.shape[0],6))

B[:,0]=b_vectors[:,0]*b_vectors[:,0]

B[:,1]=b_vectors[:,0]*b_vectors[:,1]

B[:,2]=b_vectors[:,0]*b_vectors[:,2]

B[:,3]=b_vectors[:,1]*b_vectors[:,1]

B[:,4]=b_vectors[:,1]*b_vectors[:,2]

B[:,5]=b_vectors[:,2 ]*b_vectors[:,2]

ensemble_D=X = np.empty(shape=[dmri_data.shape[0], dmri_data.shape[1],dmri_data.shape[2],6])

for x in range (dmri_data.shape[0]) :

print(x)

for y in range (dmri_data.shape[1]) :

for z in range(dmri_data.shape[2]) :

S0=dmri_data[x,y,z,0]

if (S0>60) :

S=np.array(dmri_data[x,y,z,1:])

X=S/S0

X=np.log(X)

X=-X/b_value

D=np.linalg.inv(np.dot(B.T,B)).dot(B.T).dot(X)

ensemble_D[x,y,z]=D

else :

ensemble_D[x,y,z]=np.array([0, 0, 0, 0, 0, 0])

#enregister les tenseurs

img = nib.Nifti1Image(ensemble_D, None, header=dmri.header)

img.set_data_dtype(np.float32)

nib.save(img, filename="tensors.nii")

# =============================================================================

# SVD+ carte FA/ADC

# =============================================================================

evals=np.empty(shape=[dmri_data.shape[0], dmri_data.shape[1],dmri_data.shape[2],3])

evecs=np.empty(shape=[dmri_data.shape[0], dmri_data.shape[1],dmri_data.shape[2],3,3])

fa=np.zeros(shape=[dmri_data.shape[0], dmri_data.shape[1],dmri_data.shape[2]])

adc=np.empty(shape=[dmri_data.shape[0], dmri_data.shape[1],dmri_data.shape[2]])

for x in range (dmri_data.shape[0]) :

print(x)

for y in range (dmri_data.shape[1]) :

for z in range(dmri_data.shape[2]) :

img=ensemble_D[x,y,z]

D=np.zeros((3,3))

D[0,:]=img[0:3]

D[1,1:]=img[3:5]

D[1,0]=img[1]

D[2,0]=img[2]

D[2,1]=img[4]

D[2,2]=img[5]

if (dmri_data[x,y,z,0]<60) :

if (D[0,0]!=-np.inf and D[0,0]!=np.inf and math.isnan(D[0,0])==False) :

u,s,vh= np.linalg.svd(D, full_matrices=False)

evals[x,y,z]=s

evecs[x,y,z]=vh

fa[x,y,z]=0

adc[x,y,z]=0

else :

if (D[0,0]!=-np.inf and D[0,0]!=np.inf and math.isnan(D[0,0])==False) :

u,s,vh= np.linalg.svd(D, full_matrices=False)

evals[x,y,z]=s

evecs[x,y,z]=vh

fa[x,y,z]=np.sqrt(1/2)*np.sqrt(np.power(s[0]-s[1],2)+np.power(s[1]-s[2],2)+np.power(s[2]-s[0],2))/np.sqrt(s[0]*s[0]+s[1]*s[1]+s[2]*s[2])

adc[x,y,z]=(s[0]+s[1]+s[2])/3

print(ensemble_D[30:30])

def afficher_tenseurs(fa_,evec,eva) :

window.record(ren, out_path='tensor.png', size=(1200, 1200))

afficher_tenseurs(fa, evecs,evals)

def enregistrer_cartes(fa_, adc_) :

nib.save(img, filename="adc.nii")

fa=median_filter(fa, size=2)

enregistrer_cartes(fa, adc)

# =============================================================================

# tractographie

# =============================================================================

fa_threshold1=0.80

fa_threshold2=0.40

streamlines=nib.streamlines.array_sequence.ArraySequence()

#choix mask

mask1=np.zeros(shape=[dmri_data.shape[0], dmri_data.shape[1],dmri_data.shape[2]])

mask2=np.zeros(shape=[dmri_data.shape[0], dmri_data.shape[1],dmri_data.shape[2]])

for x in range (dmri_data.shape[0]) :

for y in range (dmri_data.shape[1]) :

for z in range(dmri_data.shape[2]) :

if (fa[x,y,z]>fa_threshold1) :

mask1[x,y,z]=255

if (fa[x,y,z]>fa_threshold2) :

mask2[x,y,z]=255

plt.imshow(fa[:,:,30], cmap="gray")

plt.figure()

plt.imshow(mask1[:,:,30], cmap="gray")

plt.figure()

plt.imshow(fa[:,63,:], cmap="gray")

plt.figure()

plt.imshow(mask1[:,63,:], cmap="gray")

plt.figure()

plt.imshow(fa[63,:,:], cmap="gray")

plt.figure()

plt.imshow(mask1[63,:,:], cmap="gray")

#generation des seeds

taille=100

seeds=np.empty(shape=(taille,3))

cmpt=0

while (cmpt<taille):

random_index_x = np.random.randint(0, dmri_data.shape[0])

random_index_y = np.random.randint(0, dmri_data.shape[1])

random_index_z = np.random.randint(0, dmri_data.shape[2])

if(mask1[random_index_x,random_index_y, random_index_z]==255) :#position valide dans le mask

seeds[cmpt]=[random_index_x,random_index_y, random_index_z]

cmpt+=1

#tracking

delta_step=0.1

def get_closest_dir(evec, direction) :

return evec[:,np.argmax(sim)], angle

def interpolation (dir_prec, coord):

return retour, angle

streamlines = nib.streamlines.array_sequence.ArraySequence()

for i in tqdm(range(taille)) :

l=seeds[i]

l=np.vstack([l, seeds[i]])

l=np.int32(l)

dir_prec=evecs[l[1][0],l[1][1],l[1][2]][:,0]

angle=0

conditions=l[-1][0]<dmri_data.shape[0]-1 and l[-1][1]<dmri_data.shape[1]-1 and l[-1][2]<dmri_data.shape[2]-1

conditions=conditions and mask2[l[-1][0],l[-1][1], l[-1][2]]==255

conditions= conditions and angle<40

while (conditions):

x=np.int32(l[-1][0])

y=np.int32(l[-1][1])

z=np.int32(l[-1][2])

directions=evecs[x,y,z]

direction_princ, angle= interpolation(dir_prec, l[-1])

p=l[-1]+ direction_princ*delta_step

dir_prec=direction_princ

l=np.vstack([l, p])

conditions=l[-1][0]<dmri_data.shape[0]-1 and l[-1][1]<dmri_data.shape[1]-1 and l[-1][2]<dmri_data.shape[2]-1

conditions=conditions and mask2[np.int32(l[-1][0]),np.int32(l[-1][1]), np.int32(l[-1][2])]==255

conditions= conditions and angle<40

streamlines.append(l)

#print(streamlines)

tractogramme=nib.streamlines.tractogram.Tractogram(streamlines, affine_to_rasmm=dmri.affine)

file=nib.streamlines.trk.TrkFile(tractogramme)

file.save("fibre.trk")

#