Diff of /ConvNet_driver.py [000000] .. [271336]

Switch to side-by-side view

--- 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