<reponame>summerblind/neurotrace
from keras.models import Sequential
from keras.layers.convolutional import Conv2D, ZeroPadding2D
from keras.layers.core import Dense
import nibabel as nib
import numpy as np
from keras.utils.np_utils import to_categorical
import pdb

def create_model(label_data,raw_data):

    input_shape = (9,)

    model = Sequential()
    model.add(Dense(64,activation='relu',input_shape=input_shape))
    model.add(Dense(4,activation='softmax'))

    model.compile(optimizer='adam',
    loss='categorical_crossentropy',
    metrics=['accuracy'])

    model.fit(raw_data, label_data, epochs=500, batch_size=8192)
    return model

def load_training_data():

    label_data = []
    raw_data = []

    for sidx in range(1,19):

        idx_str = str(sidx).zfill(2)
        in_path = '/mnt/hippocampus/mri_data/IBSR_nifti_stripped' + '/IBSR_' + idx_str + '/IBSR_' + idx_str + '_ana_strip.bc.nrm.nii'
        ref_path = '/mnt/hippocampus/mri_data/IBSR_nifti_stripped' + '/IBSR_' + idx_str + '/IBSR_' + idx_str + '_segTRI_fill_ana.nii'

        ref_data = nib.load(ref_path)
        ref_img = ref_data.get_data()
        ref_img = np.moveaxis(ref_img,1,2)
        ref_dim = ref_img.shape
        ref_array = ref_img.reshape(ref_dim[0]*ref_dim[1]*ref_dim[2])
        ref_array_idx = np.asarray(range(0,len(ref_array)))

        bg_idx = ref_array==0.0
        csf_idx = ref_array==1.0
        gm_idx = ref_array==2.0
        wm_idx = ref_array==3.0

        bg_total = np.sum(bg_idx)
        csf_total = np.sum(csf_idx)
        gm_total = np.sum(gm_idx)
        wm_total = np.sum(wm_idx)

        total = bg_total + csf_total + gm_total + wm_total

        min_total = np.min(np.asarray([bg_total,csf_total,gm_total,wm_total]))
        ni_data = nib.load(in_path)
        ni_img = ni_data.get_data()
        ni_img = np.moveaxis(ni_img,1,2)

        z_dim = ni_img.shape[2]
        y_dim = ni_img.shape[1]
        x_dim = ni_img.shape[0]
        img_3D_pad = np.pad(ni_img,((1,1),(1,1),(0,0)), 'constant')

        bg_perm_idx = np.random.permutation(ref_array_idx[bg_idx])[0:min_total]
        csf_perm_idx = np.random.permutation(ref_array_idx[csf_idx])[0:min_total]
        gm_perm_idx = np.random.permutation(ref_array_idx[gm_idx])[0:min_total]
        wm_perm_idx = np.random.permutation(ref_array_idx[wm_idx])[0:min_total]

        for idx in range(0,min_total):
            label_data.append(0.0)
            label_data.append(1.0)
            label_data.append(2.0)
            label_data.append(3.0)

            bg_perm_idx_3D = np.unravel_index(bg_perm_idx[idx],(x_dim,y_dim,z_dim))
            csf_perm_idx_3D = np.unravel_index(csf_perm_idx[idx],(x_dim,y_dim,z_dim))
            gm_perm_idx_3D = np.unravel_index(gm_perm_idx[idx],(x_dim,y_dim,z_dim))
            wm_perm_idx_3D = np.unravel_index(wm_perm_idx[idx],(x_dim,y_dim,z_dim))


            patch_bg = img_3D_pad[bg_perm_idx_3D[0]:bg_perm_idx_3D[0]+3,bg_perm_idx_3D[1]:bg_perm_idx_3D[1]+3,bg_perm_idx_3D[2]].reshape(9)
            patch_csf = img_3D_pad[csf_perm_idx_3D[0]:csf_perm_idx_3D[0]+3,csf_perm_idx_3D[1]:csf_perm_idx_3D[1]+3,csf_perm_idx_3D[2]].reshape(9)
            patch_gm = img_3D_pad[gm_perm_idx_3D[0]:gm_perm_idx_3D[0]+3,gm_perm_idx_3D[1]:gm_perm_idx_3D[1]+3,gm_perm_idx_3D[2]].reshape(9)
            patch_wm = img_3D_pad[wm_perm_idx_3D[0]:wm_perm_idx_3D[0]+3,wm_perm_idx_3D[1]:wm_perm_idx_3D[1]+3,wm_perm_idx_3D[2]].reshape(9)

            raw_data.append(patch_bg)
            raw_data.append(patch_csf)
            raw_data.append(patch_gm)
            raw_data.append(patch_wm)

    n_samps = len(raw_data)
    perm_idx = np.random.permutation(n_samps)
    label_data = np.asarray(label_data)
    raw_data = np.asarray(raw_data)

    label_data = label_data[perm_idx]
    raw_data = raw_data[perm_idx]
    label_dim = (label_data.shape[0],4)
    label_data = to_categorical(label_data)
    label_data = label_data.reshape(label_dim)

    return label_data, raw_data

label_data, raw_data = load_training_data()

model = create_model(label_data,raw_data)
model.save('dn.mdl')

from keras.models import Sequential
from keras.layers.convolutional import Conv2D, ZeroPadding2D
from keras.layers.core import Dense
import nibabel as nib
import numpy as np
from keras.utils.np_utils import to_categorical
import pdb

def create_model(label_data,raw_data):

    input_shape = (9,)

    model = Sequential()
    model.add(Dense(64,activation='relu',input_shape=input_shape))
    model.add(Dense(4,activation='softmax'))

    model.compile(optimizer='adam',
    loss='categorical_crossentropy',
    metrics=['accuracy'])

    model.fit(raw_data, label_data, epochs=500, batch_size=8192)
    return model

def load_training_data():

    label_data = []
    raw_data = []

    for sidx in range(1,19):

        idx_str = str(sidx).zfill(2)
        in_path = '/mnt/hippocampus/mri_data/IBSR_nifti_stripped' + '/IBSR_' + idx_str + '/IBSR_' + idx_str + '_ana_strip.bc.nrm.nii'
        ref_path = '/mnt/hippocampus/mri_data/IBSR_nifti_stripped' + '/IBSR_' + idx_str + '/IBSR_' + idx_str + '_segTRI_fill_ana.nii'

        ref_data = nib.load(ref_path)
        ref_img = ref_data.get_data()
        ref_img = np.moveaxis(ref_img,1,2)
        ref_dim = ref_img.shape
        ref_array = ref_img.reshape(ref_dim[0]*ref_dim[1]*ref_dim[2])
        ref_array_idx = np.asarray(range(0,len(ref_array)))

        bg_idx = ref_array==0.0
        csf_idx = ref_array==1.0
        gm_idx = ref_array==2.0
        wm_idx = ref_array==3.0

        bg_total = np.sum(bg_idx)
        csf_total = np.sum(csf_idx)
        gm_total = np.sum(gm_idx)
        wm_total = np.sum(wm_idx)

        total = bg_total + csf_total + gm_total + wm_total

        min_total = np.min(np.asarray([bg_total,csf_total,gm_total,wm_total]))
        ni_data = nib.load(in_path)
        ni_img = ni_data.get_data()
        ni_img = np.moveaxis(ni_img,1,2)

        z_dim = ni_img.shape[2]
        y_dim = ni_img.shape[1]
        x_dim = ni_img.shape[0]
        img_3D_pad = np.pad(ni_img,((1,1),(1,1),(0,0)), 'constant')

        bg_perm_idx = np.random.permutation(ref_array_idx[bg_idx])[0:min_total]
        csf_perm_idx = np.random.permutation(ref_array_idx[csf_idx])[0:min_total]
        gm_perm_idx = np.random.permutation(ref_array_idx[gm_idx])[0:min_total]
        wm_perm_idx = np.random.permutation(ref_array_idx[wm_idx])[0:min_total]

        for idx in range(0,min_total):
            label_data.append(0.0)
            label_data.append(1.0)
            label_data.append(2.0)
            label_data.append(3.0)

            bg_perm_idx_3D = np.unravel_index(bg_perm_idx[idx],(x_dim,y_dim,z_dim))
            csf_perm_idx_3D = np.unravel_index(csf_perm_idx[idx],(x_dim,y_dim,z_dim))
            gm_perm_idx_3D = np.unravel_index(gm_perm_idx[idx],(x_dim,y_dim,z_dim))
            wm_perm_idx_3D = np.unravel_index(wm_perm_idx[idx],(x_dim,y_dim,z_dim))


            patch_bg = img_3D_pad[bg_perm_idx_3D[0]:bg_perm_idx_3D[0]+3,bg_perm_idx_3D[1]:bg_perm_idx_3D[1]+3,bg_perm_idx_3D[2]].reshape(9)
            patch_csf = img_3D_pad[csf_perm_idx_3D[0]:csf_perm_idx_3D[0]+3,csf_perm_idx_3D[1]:csf_perm_idx_3D[1]+3,csf_perm_idx_3D[2]].reshape(9)
            patch_gm = img_3D_pad[gm_perm_idx_3D[0]:gm_perm_idx_3D[0]+3,gm_perm_idx_3D[1]:gm_perm_idx_3D[1]+3,gm_perm_idx_3D[2]].reshape(9)
            patch_wm = img_3D_pad[wm_perm_idx_3D[0]:wm_perm_idx_3D[0]+3,wm_perm_idx_3D[1]:wm_perm_idx_3D[1]+3,wm_perm_idx_3D[2]].reshape(9)

            raw_data.append(patch_bg)
            raw_data.append(patch_csf)
            raw_data.append(patch_gm)
            raw_data.append(patch_wm)

    n_samps = len(raw_data)
    perm_idx = np.random.permutation(n_samps)
    label_data = np.asarray(label_data)
    raw_data = np.asarray(raw_data)

    label_data = label_data[perm_idx]
    raw_data = raw_data[perm_idx]
    label_dim = (label_data.shape[0],4)
    label_data = to_categorical(label_data)
    label_data = label_data.reshape(label_dim)

    return label_data, raw_data

label_data, raw_data = load_training_data()

model = create_model(label_data,raw_data)
model.save('dn.mdl')