--- a
+++ b/demo.py
@@ -0,0 +1,267 @@
+import numpy as np
+import keras
+import argparse
+import os
+import tf_models
+import tensorflow as tf
+from keras.models import Sequential, Model
+from keras.layers import Dense, Conv3D, Dropout, Flatten, Input, concatenate, Reshape, Lambda, Permute
+from keras.layers.core import Dense, Dropout, Activation, Reshape
+from keras.layers.convolutional import Conv3D, Conv3DTranspose, UpSampling3D
+from keras.layers.pooling import AveragePooling3D
+from keras.layers import Input
+from keras.layers.merge import concatenate
+from keras.layers.normalization import BatchNormalization
+from tensorflow.contrib.keras.python.keras.backend import learning_phase
+
+from nibabel import load as load_nii
+from sklearn.preprocessing import scale
+import matplotlib.pyplot as plt
+
+
+
+# SAVE_PATH = 'unet3d_baseline.hdf5'
+# OFFSET_W = 16
+# OFFSET_H = 16
+# OFFSET_C = 4
+# HSIZE = 64
+# WSIZE = 64
+# CSIZE = 16
+# batches_h, batches_w, batches_c = (224-HSIZE)/OFFSET_H+1, (224-WSIZE)/OFFSET_W+1, (152 - CSIZE)/OFFSET_C+1
+
+
+def parse_inputs():
+    parser = argparse.ArgumentParser(description='Test different nets with 3D data.')
+    parser.add_argument('-r', '--root-path', dest='root_path', default='/mnt/disk1/dat/lchen63/brain/data/data2')
+    parser.add_argument('-m', '--model-path', dest='model_path',
+                        default='NoneDense-0')
+    parser.add_argument('-ow', '--offset-width', dest='offset_w', type=int, default=12)
+    parser.add_argument('-oh', '--offset-height', dest='offset_h', type=int, default=12)
+    parser.add_argument('-oc', '--offset-channel', dest='offset_c', nargs='+', type=int, default=12)
+    parser.add_argument('-ws', '--width-size', dest='wsize', type=int, default=38)
+    parser.add_argument('-hs', '--height-size', dest='hsize', type=int, default=38)
+    parser.add_argument('-cs', '--channel-size', dest='csize', type=int, default=38)
+    parser.add_argument('-ps', '--pred-size', dest='psize', type=int, default=12)
+    parser.add_argument('-gpu', '--gpu', dest='gpu', type=str, default='0')
+    parser.add_argument('-mn', '--model_name', dest='model_name', type=str, default='dense24')
+    parser.add_argument('-nc', '--correction', dest='correction', type=bool, default=True)
+    parser.add_argument('-sp', '--save_path', dest='save_path', type=str, default='/mnt/disk1/dat/lchen63/brain/data/result/')
+
+
+    return vars(parser.parse_args())
+
+
+options = parse_inputs()
+os.environ["CUDA_VISIBLE_DEVICES"] = options['gpu']
+
+
+
+def vox_preprocess(vox):
+    vox_shape = vox.shape
+    vox = np.reshape(vox, (-1, vox_shape[-1]))
+    vox = scale(vox, axis=0)
+    return np.reshape(vox, vox_shape)
+
+
+def one_hot(y, num_classees):
+    y_ = np.zeros([len(y), num_classees])
+    y_[np.arange(len(y)), y] = 1
+    return y_
+
+
+def dice_coef_np(y_true, y_pred, num_classes):
+    """
+
+    :param y_true: sparse labels
+    :param y_pred: sparse labels
+    :param num_classes: number of classes
+    :return:
+    """
+    y_true = y_true.astype(int)
+    y_pred = y_pred.astype(int)
+    y_true = y_true.flatten()
+    y_true = one_hot(y_true, num_classes)
+    y_pred = y_pred.flatten()
+    y_pred = one_hot(y_pred, num_classes)
+    intersection = np.sum(y_true * y_pred, axis=0)
+    return (2. * intersection) / (np.sum(y_true, axis=0) + np.sum(y_pred, axis=0))
+
+
+
+
+
+
+
+
+
+def norm(image):
+    image = np.squeeze(image)
+    image_nonzero = image[np.nonzero(image)]
+    return (image - image_nonzero.mean()) / image_nonzero.std()
+
+
+def vox_generator_test(all_files):
+
+    path = options['root_path']
+
+    while 1:
+        for file in all_files:
+            p = file
+            if options['correction']:
+                flair = load_nii(os.path.join(path, file, file + '_flair_corrected.nii.gz')).get_data()
+                t2 = load_nii(os.path.join(path, file, file + '_t2_corrected.nii.gz')).get_data()
+                t1 = load_nii(os.path.join(path, file, file + '_t1_corrected.nii.gz')).get_data()
+                t1ce = load_nii(os.path.join(path, file, file + '_t1ce_corrected.nii.gz')).get_data()
+            else:
+                flair = load_nii(os.path.join(path, p, p + '_flair.nii.gz')).get_data()
+
+                t2 = load_nii(os.path.join(path, p, p + '_t2.nii.gz')).get_data()
+
+                t1 = load_nii(os.path.join(path, p, p + '_t1.nii.gz')).get_data()
+
+                t1ce = load_nii(os.path.join(path, p, p + '_t1ce.nii.gz')).get_data()
+            data = np.array([flair, t2, t1, t1ce])
+            data = np.transpose(data, axes=[1, 2, 3, 0])
+
+            data_norm = np.array([norm(flair), norm(t2), norm(t1), norm(t1ce)])
+            data_norm = np.transpose(data_norm, axes=[1, 2, 3, 0])
+
+            labels = load_nii(os.path.join(path, p, p + '_seg.nii.gz')).get_data()
+
+            yield data, data_norm, labels
+
+
+
+def main():
+    test_files = []
+    with open('test.txt') as f:
+        for line in f:
+            test_files.append(line[:-1])
+
+    num_labels = 5
+    OFFSET_H = options['offset_h']
+    OFFSET_W = options['offset_w']
+    OFFSET_C = options['offset_c']
+    HSIZE = options['hsize']
+    WSIZE = options['wsize']
+    CSIZE = options['csize']
+    PSIZE = options['psize']
+    SAVE_PATH = options['model_path']
+    model_name = options['model_name']
+
+    OFFSET_PH = (HSIZE - PSIZE) / 2
+    OFFSET_PW = (WSIZE - PSIZE) / 2
+    OFFSET_PC = (CSIZE - PSIZE) / 2
+
+    batches_w = int(np.ceil((240 - WSIZE) / float(OFFSET_W))) + 1
+    batches_h = int(np.ceil((240 - HSIZE) / float(OFFSET_H))) + 1
+    batches_c = int(np.ceil((155 - CSIZE) / float(OFFSET_C))) + 1
+
+
+
+    flair_t2_node = tf.placeholder(dtype=tf.float32, shape=(None, HSIZE, WSIZE, CSIZE, 2))
+    t1_t1ce_node = tf.placeholder(dtype=tf.float32, shape=(None, HSIZE, WSIZE, CSIZE, 2))
+
+
+    if model_name == 'dense48':
+        
+        flair_t2_15, flair_t2_27 = tf_models.BraTS2ScaleDenseNetConcat_large(input=flair_t2_node, name='flair')
+        t1_t1ce_15, t1_t1ce_27 = tf_models.BraTS2ScaleDenseNetConcat_large(input=t1_t1ce_node, name='t1')
+    elif model_name == 'no_dense':
+
+        flair_t2_15, flair_t2_27 = tf_models.PlainCounterpart(input=flair_t2_node, name='flair')
+        t1_t1ce_15, t1_t1ce_27 = tf_models.PlainCounterpart(input=t1_t1ce_node, name='t1')
+
+    elif model_name == 'dense24':
+
+        flair_t2_15, flair_t2_27 = tf_models.BraTS2ScaleDenseNetConcat(input=flair_t2_node, name='flair')
+        t1_t1ce_15, t1_t1ce_27 = tf_models.BraTS2ScaleDenseNetConcat(input=t1_t1ce_node, name='t1')
+
+    elif model_name == 'dense24_nocorrection':
+
+        flair_t2_15, flair_t2_27 = tf_models.BraTS2ScaleDenseNetConcat(input=flair_t2_node, name='flair')
+        t1_t1ce_15, t1_t1ce_27 = tf_models.BraTS2ScaleDenseNetConcat(input=t1_t1ce_node, name='t1')
+
+    else:
+        print' No such model name '
+
+    t1_t1ce_15 = concatenate([t1_t1ce_15, flair_t2_15])
+    t1_t1ce_27 = concatenate([t1_t1ce_27, flair_t2_27])
+
+    t1_t1ce_15 = Conv3D(num_labels, kernel_size=1, strides=1, padding='same', name='t1_t1ce_15_cls')(t1_t1ce_15)
+    t1_t1ce_27 = Conv3D(num_labels, kernel_size=1, strides=1, padding='same', name='t1_t1ce_27_cls')(t1_t1ce_27)
+
+    t1_t1ce_score = t1_t1ce_15[:, 13:25, 13:25, 13:25, :] + \
+                    t1_t1ce_27[:, 13:25, 13:25, 13:25, :]
+
+
+    saver = tf.train.Saver()
+    data_gen_test = vox_generator_test(test_files)
+    dice_whole, dice_core, dice_et = [], [], []
+    with tf.Session() as sess:
+        saver.restore(sess, SAVE_PATH)
+        for i in range(len(test_files)):
+            print 'predicting %s' % test_files[i]
+            x, x_n, y = data_gen_test.next()
+            pred = np.zeros([240, 240, 155, 5])
+            for hi in range(batches_h):
+                offset_h = min(OFFSET_H * hi, 240 - HSIZE)
+                offset_ph = offset_h + OFFSET_PH
+                for wi in range(batches_w):
+                    offset_w = min(OFFSET_W * wi, 240 - WSIZE)
+                    offset_pw = offset_w + OFFSET_PW
+                    for ci in range(batches_c):
+                        offset_c = min(OFFSET_C * ci, 155 - CSIZE)
+                        offset_pc = offset_c + OFFSET_PC
+                        data = x[offset_h:offset_h + HSIZE, offset_w:offset_w + WSIZE, offset_c:offset_c + CSIZE, :]
+                        data_norm = x_n[offset_h:offset_h + HSIZE, offset_w:offset_w + WSIZE, offset_c:offset_c + CSIZE, :]
+                        data_norm = np.expand_dims(data_norm, 0)
+                        if not np.max(data) == 0 and np.min(data) == 0:
+                            score = sess.run(fetches=t1_t1ce_score,
+                                             feed_dict={flair_t2_node: data_norm[:, :, :, :, :2],
+                                                        t1_t1ce_node: data_norm[:, :, :, :, 2:],
+                                                        learning_phase(): 0}
+                                             )
+                            pred[offset_ph:offset_ph + PSIZE, offset_pw:offset_pw + PSIZE, offset_pc:offset_pc + PSIZE,
+                            :] += np.squeeze(score)
+
+            pred = np.argmax(pred, axis=-1)
+            pred = pred.astype(int)
+            print 'calculating dice...'
+            print  options['save_path'] +  test_files[i] +'_prediction'
+            np.save(options['save_path'] +  test_files[i] +'_prediction',pred)
+            whole_pred = (pred > 0).astype(int)
+            whole_gt = (y > 0).astype(int)
+            core_pred = (pred == 1).astype(int) + (pred == 4).astype(int)
+            core_gt = (y == 1).astype(int) + (y == 4).astype(int)
+            et_pred = (pred == 4).astype(int)
+            et_gt = (y == 4).astype(int)
+            dice_whole_batch = dice_coef_np(whole_gt, whole_pred, 2)
+            dice_core_batch = dice_coef_np(core_gt, core_pred, 2)
+            dice_et_batch = dice_coef_np(et_gt, et_pred, 2)
+            dice_whole.append(dice_whole_batch)
+            dice_core.append(dice_core_batch)
+            dice_et.append(dice_et_batch)
+            print dice_whole_batch
+            print dice_core_batch
+            print dice_et_batch
+
+        dice_whole = np.array(dice_whole)
+        dice_core = np.array(dice_core)
+        dice_et = np.array(dice_et)
+
+        print 'mean dice whole:'
+        print np.mean(dice_whole, axis=0)
+        print 'mean dice core:'
+        print np.mean(dice_core, axis=0)
+        print 'mean dice enhance:'
+        print np.mean(dice_et, axis=0)
+
+        np.save(model_name + '_dice_whole', dice_whole)
+        np.save(model_name + '_dice_core', dice_core)
+        np.save(model_name + '_dice_enhance', dice_et)
+        print 'pred saved'
+
+
+if __name__ == '__main__':
+    main()
\ No newline at end of file