--- a +++ b/ConvNet_driver.py @@ -0,0 +1,170 @@ +# -*- coding: utf-8 -*- +""" +Created on Wed Nov 02 21:35:59 2016 + +@author: seeker105 +""" +import os.path +import sys +from ConvNet import LeNet +import json +import SimpleITK as sitk +import pylab +from skimage import color +from sklearn.utils import shuffle +from scipy.ndimage.interpolation import rotate +from keras.optimizers import SGD +from keras.utils import np_utils +from keras.callbacks import LearningRateScheduler, ModelCheckpoint +import numpy as np +from Brain_pipeline import Pipeline +import Brain_pipeline +import Metrics +from glob import glob +import model_test + +''' Script to drive loading, training, testing and saving the brain MRI + First, we load all the images, and process them through the Pipeline, + and get the pre-processed images as output. + Then, we train the model from ConvNet or load the weights into it. + We divide the training and test data using train_test_split. + ''' +def show_segmented_image(orig_img, pred_img): + ''' + Show the prediction over the original image + INPUT: + 1)orig_img: the test image, which was used as input + 2)pred_img: the prediction output + OUTPUT: + segmented image rendering + ''' + #define the colours of the labels + red = [10, 0, 0] #label 1 + yellow = [10, 10, 0] #label 2 + green = [0, 10, 0] #label 3 + blue = [0, 0, 10] #label 4 + #convert original image to rgb + gray_im = color.gray2rgb(orig_img) + #color the tumor voxels + gray_im[pred_img == 1] = red + gray_im[pred_img == 2] = yellow + gray_im[pred_img == 3] = green + gray_im[pred_img == 4] = blue + pylab.imshow(gray_im) + + +def step_decay(epochs): + init_rate = 0.003 + fin_rate = 0.00003 + total_epochs = 24 + print 'ep: {}'.format(epochs) + if epochs<25: + lrate = init_rate - (init_rate - fin_rate)/total_epochs * float(epochs) + else: lrate = 0.00003 + print 'lrate: {}'.format(model.optimizer.lr.get_value()) + return lrate + +pth_train = 'D:/New folder/BRATS2015_Training/train_slices/' +pth_test = 'D:/New folder/BRATS2015_Training/test_slices/' +x = Pipeline(pth_train, pth_test) #pass the images through the preprocessing steps + +#build the model +model = LeNet.build_Pereira(33, 33, 4, 5) + +#callback +change_lr = LearningRateScheduler(step_decay) + +#initialize the optimizer and model +opt = SGD(lr = 0.003, momentum=0.9, decay= 0, nesterov = True) +model.compile(loss = 'categorical_crossentropy', optimizer=opt, metrics = ['accuracy']) + + +#load training patches +X_patches, Y_labels, mu, sigma = x.training_patches([180000, 67500, 67500, 67500, 67500]) +tmp = rotate(X_patches, 90, (2, 3)) +tmp = np.append(tmp, rotate(X_patches, -90, (2, 3)), axis=0) +tmp = np.append(tmp, rotate(X_patches, 180, (2, 3)), axis=0) +X_patches = np.append(X_patches, tmp, axis=0) +Y_labels = np.hstack(Y_labels) +for i in xrange(2): + Y_labels = np.append(Y_labels, Y_labels, axis=0) +# Labels should be in categorical array form 1x5 +Y_labels = np_utils.to_categorical(Y_labels, 5) +X_patches, Y_labels = shuffle(X_patches, Y_labels, random_state=0) + +#save model after each epoch +os.mkdir(r'D:\New folder\Pereira_model_checkpoints') +checkpointer = ModelCheckpoint(filepath='D:/New folder/Pereira_model_checkpoints/weights.{epoch:02d}-{val_loss:.2f}.keras2.hdf5',monitor = 'val_loss', verbose=1) +#fit model and shuffle training data +hist = model.fit(X_patches[:200000], Y_labels[:200000], nb_epoch=25, batch_size=128, verbose=1, validation_split=0.1, callbacks = [change_lr, checkpointer]) + +#save model +sv_pth = 'D:/New Folder/Pereira_model_checkpoints/model_weights' +m = '{}.json'.format(sv_pth) +w = '{}.hdf5'.format(sv_pth) +model.save_weights(w) +json_strng = model.to_json() +with open(m, 'w') as f: + json.dump(json_strng, f) + + +#test all the test image slices +test_im = x.test_im.swapaxes(0,1) +gt = test_im[4] +test_im = test_im[:4].swapaxes(0, 1) +predicted_images, params = model_test.test_slices(test_im[158:159], gt[158:159], model, mu, sigma) + +'''test_pths = zip(*x.pathnames_test) +#show a segmented slice +tst = test_pths[0]#random.choice(test_pths) +test_arr = [sitk.GetArrayFromImage(sitk.ReadImage(i)) for i in tst] +final_pth = os.path.dirname(os.path.dirname(tst[0])) + '/' + os.path.splitext(os.path.splitext(os.path.basename(tst[0]))[0])[0] + '_processed_predicted_70.mha' +slice_arr = [test_arr[j][70] for j in xrange(4)] +patches = Brain_pipeline.test_patches(slice_arr) +pred = model.predict_classes(patches) +pred = Brain_pipeline.reconstruct_labels(pred) +show_segmented_image(test_arr[0][70], pred) +sitk.WriteImage(sitk.GetImageFromArray(np.array(pred.astype(float))), final_pth) + + +#evaluate metrics +DSC_arr = [] #stores DSC +DSC_core_arr = [] #stores list of core DSCs +PPV_arr = [] +acc_arr = [] + +#use for getting orignal brain image and prediction label slices +# use for: + #overlay images + #segmentation vs orig label. it's in test_paths + #with/without nyul + #4 sequences after nyul. for original ones, redefine paths + #ok. now we gotta see metrics brother +pred_pth = [] +t1c_pth = [] +pred_arr = [] + +for i in xrange(len(test_pths)): + tst = test_pths[i] + test_arr = [sitk.GetArrayFromImage(sitk.ReadImage(j)) for j in tst] + #take slices + slice_arr = [test_arr[j][70] for j in xrange(4)] + #read original slice label + orig = test_arr[4][70] + patches = Brain_pipeline.test_patches(slice_arr) + pred = model.predict_classes(patches) + pred = Brain_pipeline.reconstruct_labels(pred) + acc_arr.append(Metrics.accuracy(pred, orig)) + DSC_arr.append(Metrics.DSC(pred, orig, 2)) + DSC_core_arr.append(Metrics.DSC_core_tumor(pred, orig)) + PPV_arr.append(Metrics.PPV(pred, orig)) + print 'acc: {}'.format(acc_arr[i]) + print 'DSC: {}'.format(DSC_arr[i]) + print 'DSC_core: {}'.format(DSC_core_arr[i]) + print 'PPV : {}'.format(PPV_arr[i]) + sys.stdout.flush() + final_pth = os.path.dirname(tst[4]) + '/' + os.path.splitext(os.path.basename(tst[0]))[0] + '_predicted_70.mha' + pred_pth.append(final_pth) + pred_arr.append(pred) + t1c_pth.append([flp for flp in glob(os.path.dirname(tst[2]) + '/*.mha') if 'n4' not in flp]) + ''' \ No newline at end of file