def main(): ###======================== HYPER-PARAMETERS ============================### parser = argparse.ArgumentParser() parser.add_argument('--task', type=str, default='all', help='all, necrotic, edema, enhance') args = parser.parse_args() task = args.task ## Create folder to save trained model and result images save_dir = "checkpoint" tl.files.exists_or_mkdir(save_dir) tl.files.exists_or_mkdir("samples/{}".format(task)) ###======================== LOAD DATA ===================================### ## by importing this, you can load a training set and a validation set. # you will get X_train_input, X_train_target, X_dev_input and X_dev_target # there are 4 labels in targets: # Label 0: background # Label 1: necrotic and non-enhancing tumor # Label 2: edema # Label 4: enhancing tumor X_test, y_test = load_image() X_test = X_test[91][np.newaxis, :, :, :] y_test = y_test[:, :, :, np.newaxis][91][np.newaxis, :, :, :] if task == 'all': y_test = (y_test > 0).astype(int) elif task == 'necrotic': y_test = (y_test == 1).astype(int) elif task == 'edema': y_test = (y_test == 2).astype(int) elif task == 'enhance': y_test = (y_test == 4).astype(int) else: exit("Unknown task %s" % task) ###======================== SHOW DATA ===================================### if not os.path.exists('outputs/all'): os.makedirs('outputs/all') X = np.asarray(X_test[0]) y = np.asarray(y_test[0]) vis_imgs_with_pred(X, y, y, "outputs/{}/run_input.png".format(task)) ###======================== TRAIN ===================================### nw, nh, nz = X.shape with tf.device('/cpu:0'): sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) with tf.device('/cpu:0'): # <- remove it if you train on CPU or other GPU ###======================== DEFINE MODEL =======================### ## nz is 4 as we input all Flair, T1, T1c and T2. t_image = tf.placeholder('float32', [1, nw, nh, nz], name='input_image') ## labels are either 0 or 1 t_seg = tf.placeholder('float32', [1, nw, nh, 1], name='target_segment') ## test inference net_test = model.u_net(t_image, is_train=False, reuse=False, n_out=1) ###======================== DEFINE LOSS =========================### ## test losses test_out_seg = net_test.outputs test_dice_loss = 1 - tl.cost.dice_coe(test_out_seg, t_seg, axis=(0, 1, 2, 3)) # , 'jaccard', epsilon=1e-5) test_iou_loss = tl.cost.iou_coe(test_out_seg, t_seg, axis=(0, 1, 2, 3)) test_dice_hard = tl.cost.dice_hard_coe(test_out_seg, t_seg, axis=(0, 1, 2, 3)) ###======================== LOAD MODEL ==============================### tl.layers.initialize_global_variables(sess) ## load existing model if possible tl.files.load_and_assign_npz(sess=sess, name=save_dir + '/u_net_{}.npz'.format(task), network=net_test) ###======================== EVALUATION ==========================### total_dice, total_iou, total_dice_hard, n_batch = 0, 0, 0, 0 for batch in tl.iterate.minibatches(inputs=X_test, targets=y_test, batch_size=1, shuffle=True): b_images, b_labels = batch _dice, _iou, _diceh, out = sess.run([test_dice_loss, test_iou_loss, test_dice_hard, net_test.outputs], {t_image: b_images, t_seg: b_labels}) total_dice += _dice total_iou += _iou total_dice_hard += _diceh n_batch += 1 vis_imgs_with_pred(b_images[0], b_labels[0], out[0], "outputs/{}/test_{}.png".format(task, 0)) print(" **" + " " * 17 + "test 1-dice: %f hard-dice: %f iou: %f (2d no distortion)" % (total_dice / n_batch, total_dice_hard / n_batch, total_iou / n_batch)) print(" task: {}".format(task))
def main(): ## Create folder to save trained model and result images save_dir = "checkpoint" tl.files.exists_or_mkdir(save_dir) tl.files.exists_or_mkdir("samples") ###======================== HYPER-PARAMETERS ============================### batch_size = 10 lr = 0.0001 # lr_decay = 0.5 # decay_every = 100 beta1 = 0.9 n_epoch = 100 print_freq_step = 100 height = 160 width = 160 ###======================== LOAD DATA ===================================### ## by importing this, you can load a training set and a validation set. train_path = '/DB/rhome/qyzheng/Desktop/qyzheng/source/renji_data/from_senior/0_cv_train.csv' val_path = '/DB/rhome/qyzheng/Desktop/qyzheng/source/renji_data/from_senior/0_cv_val.csv' train = [] val = [] train_size = 0 val_size = 0 with open(train_path, 'r') as f: reader = csv.reader(f) for i in reader: train.append(i) train_size += 1 with open(val_path, 'r') as f: reader = csv.reader(f) for i in reader: val.append(i) val_size += 1 train_epoch = train_size / batch_size val_epoch = val_size / batch_size ###======================== SHOW DATA ===================================### sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) ###======================== DEFIINE MODEL =======================### ## nz is 3 as we input ADC, b0, b1000 t_image = tf.placeholder('float32', [batch_size, height, width, 3], name='input_image') ## labels are either 0 or 1 t_seg = tf.placeholder('float32', [batch_size, height, width, 1], name='target_segment') ## train inference net = model.u_net(t_image, is_train=True, reuse=False, n_out=1) ## test inference net_test = model.u_net(t_image, is_train=False, reuse=True, n_out=1) ###======================== DEFINE LOSS =========================### ## train losses out_seg = net.outputs dice_loss = 1 - tl.cost.dice_coe(out_seg, t_seg, axis=[0,1,2,3])#, 'jaccard', epsilon=1e-5) iou_loss = tl.cost.iou_coe(out_seg, t_seg, axis=[0,1,2,3]) dice_hard = tl.cost.dice_hard_coe(out_seg, t_seg, axis=[0,1,2,3]) loss = dice_loss ## test losses test_out_seg = net_test.outputs test_dice_loss = 1 - tl.cost.dice_coe(test_out_seg, t_seg, axis=[0,1,2,3])#, 'jaccard', epsilon=1e-5) test_iou_loss = tl.cost.iou_coe(test_out_seg, t_seg, axis=[0,1,2,3]) test_dice_hard = tl.cost.dice_hard_coe(test_out_seg, t_seg, axis=[0,1,2,3]) ###======================== DEFINE TRAIN OPTS =======================### t_vars = tl.layers.get_variables_with_name('u_net', True, True) with tf.variable_scope('learning_rate'): lr_v = tf.Variable(lr, trainable=False) train_op = tf.train.AdamOptimizer(lr_v, beta1=beta1).minimize(loss, var_list=t_vars) ###======================== LOAD MODEL ==============================### tl.layers.initialize_global_variables(sess) ## load existing model if possible # tl.files.load_and_assign_npz(sess=sess, name=save_dir+'/u_net.npz', network=net) ###======================== TRAINING ================================### for epoch in range(0, n_epoch+1): epoch_time = time.time() ## update decay learning rate at the beginning of a epoch # if epoch !=0 and (epoch % decay_every == 0): # new_lr_decay = lr_decay ** (epoch // decay_every) # sess.run(tf.assign(lr_v, lr * new_lr_decay)) # log = " ** new learning rate: %f" % (lr * new_lr_decay) # print(log) # elif epoch == 0: # sess.run(tf.assign(lr_v, lr)) # log = " ** init lr: %f decay_every_epoch: %d, lr_decay: %f" % (lr, decay_every, lr_decay) # print(log) total_dice, total_iou, total_dice_hard, n_batch = 0, 0, 0, 0 shuffle(train) shuffle(val) for i in range(train_epoch): images, labels = next_batch(train, batch_size, height, width, i) step_time = time.time() ## update network _, _dice, _iou, _diceh, out = sess.run([train_op, dice_loss, iou_loss, dice_hard, net.outputs], {t_image: images, t_seg: labels}) total_dice += _dice; total_iou += _iou; total_dice_hard += _diceh n_batch += 1 ## you can show the predition here: # vis_imgs2(b_images[0], b_labels[0], out[0], "samples/{}/_tmp.png".format(task)) # exit() # if _dice == 1: # DEBUG # print("DEBUG") # vis_imgs2(b_images[0], b_labels[0], out[0], "samples/{}/_debug.png".format(task)) if n_batch % print_freq_step == 0: print("Epoch %d step %d 1-dice: %f hard-dice: %f iou: %f took %fs (2d with distortion)" % (epoch, n_batch, _dice, _diceh, _iou, time.time()-step_time)) ## check model fail if np.isnan(_dice): exit(" ** NaN loss found during training, stop training") if np.isnan(out).any(): exit(" ** NaN found in output images during training, stop training") print(" ** Epoch [%d/%d] train 1-dice: %f hard-dice: %f iou: %f took %fs (2d with distortion)" % (epoch, n_epoch, total_dice/n_batch, total_dice_hard/n_batch, total_iou/n_batch, time.time()-epoch_time)) ''' ## save a predition of training set for i in range(batch_size): if np.max(b_images[i]) > 0: vis_imgs2(b_images[i], b_labels[i], out[i], "samples/{}/train_{}.png".format(task, epoch)) break elif i == batch_size-1: vis_imgs2(b_images[i], b_labels[i], out[i], "samples/{}/train_{}.png".format(task, epoch)) ''' ###======================== VALIDATION ==========================### total_dice, total_iou, total_dice_hard, n_batch = 0, 0, 0, 0 for i in range(val_epoch): val_images, val_labels = next_batch(val, batch_size, height, width, i) _dice, _iou, _diceh, out = sess.run([test_dice_loss, test_iou_loss, test_dice_hard, net_test.outputs], {t_image: val_images, t_seg: val_labels}) total_dice += _dice; total_iou += _iou; total_dice_hard += _diceh n_batch += 1 print(" **"+" "*17+"test 1-dice: %f hard-dice: %f iou: %f (2d no distortion)" % (total_dice/n_batch, total_dice_hard/n_batch, total_iou/n_batch)) ''' ## save a predition of test set for i in range(batch_size): if np.max(b_images[i]) > 0: vis_imgs2(b_images[i], b_labels[i], out[i], "samples/{}/test_{}.png".format(task, epoch)) break elif i == batch_size-1: vis_imgs2(b_images[i], b_labels[i], out[i], "samples/{}/test_{}.png".format(task, epoch)) ''' ###======================== SAVE MODEL ==========================### tl.files.save_npz(net.all_params, name=save_dir+'/u_net.npz', sess=sess)
def main(): ###======================== HYPER-PARAMETERS ============================### parser = argparse.ArgumentParser() parser.add_argument('--task', type=str, default='all', help='all, necrotic, edema, enhance') parser.add_argument('--batch_size', type=int, default=10) parser.add_argument('--lr', type=int, default=0.0001) parser.add_argument('--beta1', type=int, default=0.9) parser.add_argument('--n_epoch', type=int, default=100) parser.add_argument('--print_interval', type=int, default=100) args = parser.parse_args() task = args.task batch_size = args.batch_size lr = args.lr # lr_decay = 0.5 # decay_every = 100 beta1 = args.beta1 n_epoch = args.n_epoch print_freq_step = args.print_interval ## Create folder to save trained model and result images save_dir = "checkpoint" tl.files.exists_or_mkdir(save_dir) tl.files.exists_or_mkdir("samples/{}".format(task)) ###======================== LOAD DATA ===================================### ## by importing this, you can load a training set and a validation set. # you will get X_train_input, X_train_target, X_dev_input and X_dev_target # there are 4 labels in targets: # Label 0: background # Label 1: necrotic and non-enhancing tumor # Label 2: edema # Label 4: enhancing tumor from dataset import Dataset dataset = Dataset() X_train = dataset.X_train_input y_train = dataset.X_train_target[:, :, :, np.newaxis] X_test = dataset.X_dev_input y_test = dataset.X_dev_target[:, :, :, np.newaxis] if task == 'all': y_train = (y_train > 0).astype(int) y_test = (y_test > 0).astype(int) elif task == 'necrotic': y_train = (y_train == 1).astype(int) y_test = (y_test == 1).astype(int) elif task == 'edema': y_train = (y_train == 2).astype(int) y_test = (y_test == 2).astype(int) elif task == 'enhance': y_train = (y_train == 4).astype(int) y_test = (y_test == 4).astype(int) else: exit("Unknown task %s" % task) ###======================== SHOW DATA ===================================### X = np.asarray(X_train[80]) y = np.asarray(y_train[80]) show_image_sample(X, y, task) ###======================== TRAIN ===================================### nw, nh, nz = X.shape with tf.device('/gpu:1'): sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) with tf.device( '/gpu:0'): # <- remove it if you train on CPU or other GPU ###======================== DEFINE MODEL =======================### ## nz is 4 as we input all Flair, T1, T1c and T2. t_image = tf.placeholder('float32', [batch_size, nw, nh, nz], name='input_image') ## labels are either 0 or 1 t_seg = tf.placeholder('float32', [batch_size, nw, nh, 1], name='target_segment') ## train inference net = model.u_net(t_image, is_train=True, reuse=False, n_out=1) ## test inference net_test = model.u_net(t_image, is_train=False, reuse=True, n_out=1) ###======================== DEFINE LOSS =========================### ## train losses out_seg = net.outputs dice_loss = 1 - tl.cost.dice_coe(out_seg, t_seg, axis=( 0, 1, 2, 3)) # , 'jaccard', epsilon=1e-5) iou_loss = tl.cost.iou_coe(out_seg, t_seg, axis=(0, 1, 2, 3)) dice_hard = tl.cost.dice_hard_coe(out_seg, t_seg, axis=(0, 1, 2, 3)) loss = dice_loss ## test losses test_out_seg = net_test.outputs test_dice_loss = 1 - tl.cost.dice_coe( test_out_seg, t_seg, axis=(0, 1, 2, 3)) # , 'jaccard', epsilon=1e-5) test_iou_loss = tl.cost.iou_coe(test_out_seg, t_seg, axis=(0, 1, 2, 3)) test_dice_hard = tl.cost.dice_hard_coe(test_out_seg, t_seg, axis=(0, 1, 2, 3)) ###======================== DEFINE TRAIN OPTS =======================### t_vars = tl.layers.get_variables_with_name('u_net', True, True) with tf.device('/gpu:0'): with tf.variable_scope('learning_rate'): lr_v = tf.Variable(lr, trainable=False) train_op = tf.train.AdamOptimizer(lr_v, beta1=beta1).minimize( loss, var_list=t_vars) ###======================== LOAD MODEL ==============================### tl.layers.initialize_global_variables(sess) ## load existing model if possible tl.files.load_and_assign_npz(sess=sess, name=save_dir + '/u_net_{}.npz'.format(task), network=net) ###======================== TRAINING ================================### for epoch in range(0, n_epoch + 1): epoch_time = time.time() ## update decay learning rate at the beginning of a epoch # if epoch !=0 and (epoch % decay_every == 0): # new_lr_decay = lr_decay ** (epoch // decay_every) # sess.run(tf.assign(lr_v, lr * new_lr_decay)) # log = " ** new learning rate: %f" % (lr * new_lr_decay) # print(log) # elif epoch == 0: # sess.run(tf.assign(lr_v, lr)) # log = " ** init lr: %f decay_every_epoch: %d, lr_decay: %f" % (lr, decay_every, lr_decay) # print(log) total_dice, total_iou, total_dice_hard, n_batch = 0, 0, 0, 0 for batch in tl.iterate.minibatches(inputs=X_train, targets=y_train, batch_size=batch_size, shuffle=True): images, labels = batch step_time = time.time() ## data augmentation for a batch of Flair, T1, T1c, T2 images # and label maps synchronously. data = tl.prepro.threading_data( [ _ for _ in zip(images[:, :, :, 0, np.newaxis], images[:, :, :, 1, np.newaxis], images[:, :, :, 2, np.newaxis], images[:, :, :, 3, np.newaxis], labels) ], fn=distort_imgs) # (10, 5, 240, 240, 1) b_images = data[:, 0:4, :, :, :] # (10, 4, 240, 240, 1) b_labels = data[:, 4, :, :, :] b_images = b_images.transpose((0, 2, 3, 1, 4)) b_images.shape = (batch_size, nw, nh, nz) ## update network _, _dice, _iou, _diceh, out = sess.run( [train_op, dice_loss, iou_loss, dice_hard, net.outputs], { t_image: b_images, t_seg: b_labels }) total_dice += _dice total_iou += _iou total_dice_hard += _diceh n_batch += 1 ## you can show the prediction here: # vis_imgs2(b_images[0], b_labels[0], out[0], "samples/{}/_tmp.png".format(task)) # exit() # if _dice == 1: # DEBUG # print("DEBUG") # vis_imgs2(b_images[0], b_labels[0], out[0], "samples/{}/_debug.png".format(task)) if n_batch % print_freq_step == 0: print( "Epoch %d step %d 1-dice: %f hard-dice: %f iou: %f took %fs (2d with distortion)" % (epoch, n_batch, _dice, _diceh, _iou, time.time() - step_time)) ## check model fail if np.isnan(_dice): exit(" ** NaN loss found during training, stop training") if np.isnan(out).any(): exit( " ** NaN found in output images during training, stop training" ) print( " ** Epoch [%d/%d] train 1-dice: %f hard-dice: %f iou: %f took %fs (2d with distortion)" % (epoch, n_epoch, total_dice / n_batch, total_dice_hard / n_batch, total_iou / n_batch, time.time() - epoch_time)) ## save a prediction of training set for i in range(batch_size): if np.max(b_images[i]) > 0: vis_imgs_with_pred( b_images[i], b_labels[i], out[i], "samples/{}/train_{}.png".format(task, epoch)) break elif i == batch_size - 1: vis_imgs_with_pred( b_images[i], b_labels[i], out[i], "samples/{}/train_{}.png".format(task, epoch)) ###======================== EVALUATION ==========================### total_dice, total_iou, total_dice_hard, n_batch = 0, 0, 0, 0 for batch in tl.iterate.minibatches(inputs=X_test, targets=y_test, batch_size=batch_size, shuffle=True): b_images, b_labels = batch _dice, _iou, _diceh, out = sess.run([ test_dice_loss, test_iou_loss, test_dice_hard, net_test.outputs ], { t_image: b_images, t_seg: b_labels }) total_dice += _dice total_iou += _iou total_dice_hard += _diceh n_batch += 1 print(" **" + " " * 17 + "test 1-dice: %f hard-dice: %f iou: %f (2d no distortion)" % (total_dice / n_batch, total_dice_hard / n_batch, total_iou / n_batch)) print(" task: {}".format(task)) ## save a prediction of test set for i in range(batch_size): if np.max(b_images[i]) > 0: vis_imgs_with_pred( b_images[i], b_labels[i], out[i], "samples/{}/test_{}.png".format(task, epoch)) break elif i == batch_size - 1: vis_imgs_with_pred( b_images[i], b_labels[i], out[i], "samples/{}/test_{}.png".format(task, epoch)) ###======================== SAVE MODEL ==========================### tl.files.save_npz(net.all_params, name=save_dir + '/u_net_{}.npz'.format(task), sess=sess)
def main(task='all'): ## Create folder to save trained model and result images save_dir = "checkpoint" tl.files.exists_or_mkdir(save_dir) tl.files.exists_or_mkdir("samples/{}".format(task)) ###======================== LOAD DATA ===================================### ## by importing this, you can load a training set and a validation set. # you will get X_train_input, X_train_target, X_dev_input and X_dev_target # there are 4 labels in targets: # Label 0: background # Label 1: necrotic and non-enhancing tumor # Label 2: edema # Label 4: enhancing tumor X_train, y_train = load_data('Arrays/') y_train = y_train.astype(int) ###======================== HYPER-PARAMETERS ============================### batch_size = 10 lr = 0.0001 # lr_decay = 0.5 # decay_every = 100 beta1 = 0.9 n_epoch = 100 print_freq_step = 100 ###======================== SHOW DATA ===================================### # show one slice X = np.asarray(X_train[80]) y = np.asarray(y_train[80]) # print(X.shape, X.min(), X.max()) # (240, 240, 4) -0.380588 2.62761 # print(y.shape, y.min(), y.max()) # (240, 240, 1) 0 1 nw, nh, nz = X.shape vis_imgs(X, y, 'samples/{}/_train_im.png'.format(task)) # show data augumentation results for i in range(10): x_flair, label = distort_imgs([X[:,:, 0, np.newaxis], y])#[:,:,np.newaxis]]) # print(x_flair.shape, x_t1.shape, x_t1ce.shape, x_t2.shape, label.shape) # (240, 240, 1) (240, 240, 1) (240, 240, 1) (240, 240, 1) (240, 240, 1) X_dis = x_flair # print(X_dis.shape, X_dis.min(), X_dis.max()) # (240, 240, 4) -0.380588233471 2.62376139209 vis_imgs(X_dis, label, 'samples/{}/_train_im_aug{}.png'.format(task, i)) with tf.device('/cpu:0'): sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) with tf.device('/gpu:0'): #<- remove it if you train on CPU or other GPU ###======================== DEFIINE MODEL =======================### ## nz is 4 as we input all Flair, T1, T1c and T2. t_image = tf.placeholder('float32', [batch_size, nw, nh, nz], name='input_image') ## labels are either 0 or 1 t_seg = tf.placeholder('float32', [batch_size, nw, nh, 1], name='target_segment') ## train inference net = model.u_net(t_image, is_train=True, reuse=False, n_out=1) ## test inference net_test = model.u_net(t_image, is_train=False, reuse=True, n_out=1) ###======================== DEFINE LOSS =========================### ## train losses out_seg = net.outputs dice_loss = 1 - tl.cost.dice_coe(out_seg, t_seg, axis=[0,1,2,3])#, 'jaccard', epsilon=1e-5) iou_loss = tl.cost.iou_coe(out_seg, t_seg, axis=[0,1,2,3]) dice_hard = tl.cost.dice_hard_coe(out_seg, t_seg, axis=[0,1,2,3]) loss = dice_loss ## test losses test_out_seg = net_test.outputs test_dice_loss = 1 - tl.cost.dice_coe(test_out_seg, t_seg, axis=[0,1,2,3])#, 'jaccard', epsilon=1e-5) test_iou_loss = tl.cost.iou_coe(test_out_seg, t_seg, axis=[0,1,2,3]) test_dice_hard = tl.cost.dice_hard_coe(test_out_seg, t_seg, axis=[0,1,2,3]) ###======================== DEFINE TRAIN OPTS =======================### t_vars = tl.layers.get_variables_with_name('u_net', True, True) with tf.device('/gpu:0'): with tf.variable_scope('learning_rate'): lr_v = tf.Variable(lr, trainable=False) train_op = tf.train.AdamOptimizer(lr_v, beta1=beta1).minimize(loss, var_list=t_vars) ###======================== LOAD MODEL ==============================### tl.layers.initialize_global_variables(sess) ## load existing model if possible tl.files.load_and_assign_npz(sess=sess, name=save_dir+'/u_net_{}.npz'.format(task), network=net) ###======================== TRAINING ================================### for epoch in range(0, n_epoch+1): epoch_time = time.time() ## update decay learning rate at the beginning of a epoch # if epoch !=0 and (epoch % decay_every == 0): # new_lr_decay = lr_decay ** (epoch // decay_every) # sess.run(tf.assign(lr_v, lr * new_lr_decay)) # log = " ** new learning rate: %f" % (lr * new_lr_decay) # print(log) # elif epoch == 0: # sess.run(tf.assign(lr_v, lr)) # log = " ** init lr: %f decay_every_epoch: %d, lr_decay: %f" % (lr, decay_every, lr_decay) # print(log) total_dice, total_iou, total_dice_hard, n_batch = 0, 0, 0, 0 for batch in tl.iterate.minibatches(inputs=X_train, targets=y_train, batch_size=batch_size, shuffle=True): images, labels = batch step_time = time.time() ## data augumentation for a batch of Flair, T1, T1c, T2 images # and label maps synchronously. data = tl.prepro.threading_data([_ for _ in zip(images[:,:,:,0, np.newaxis], labels)], fn=distort_imgs) # (10, 5, 240, 240, 1) print(data.shape) b_images = data[:,0:1,:,:,:] # (10, 4, 240, 240, 1) b_labels = data[:,1,:,:,:] b_images = b_images.transpose((0,2,3,1,4)) b_images.shape = (batch_size, nw, nh, nz) ## update network _, _dice, _iou, _diceh, out = sess.run([train_op, dice_loss, iou_loss, dice_hard, net.outputs], {t_image: b_images, t_seg: b_labels}) total_dice += _dice; total_iou += _iou; total_dice_hard += _diceh n_batch += 1 ## you can show the predition here: # vis_imgs2(b_images[0], b_labels[0], out[0], "samples/{}/_tmp.png".format(task)) # exit() # if _dice == 1: # DEBUG # print("DEBUG") # vis_imgs2(b_images[0], b_labels[0], out[0], "samples/{}/_debug.png".format(task)) if n_batch % print_freq_step == 0: print("Epoch %d step %d 1-dice: %f hard-dice: %f iou: %f took %fs (2d with distortion)" % (epoch, n_batch, _dice, _diceh, _iou, time.time()-step_time)) ## check model fail if np.isnan(_dice): exit(" ** NaN loss found during training, stop training") if np.isnan(out).any(): exit(" ** NaN found in output images during training, stop training") print(" ** Epoch [%d/%d] train 1-dice: %f hard-dice: %f iou: %f took %fs (2d with distortion)" % (epoch, n_epoch, total_dice/n_batch, total_dice_hard/n_batch, total_iou/n_batch, time.time()-epoch_time)) ## save a predition of training set for i in range(batch_size): if np.max(b_images[i]) > 0: vis_imgs2(b_images[i], b_labels[i], out[i], "samples/{}/train_{}.png".format(task, epoch)) break elif i == batch_size-1: vis_imgs2(b_images[i], b_labels[i], out[i], "samples/{}/train_{}.png".format(task, epoch)) ###======================== EVALUATION ==========================### # total_dice, total_iou, total_dice_hard, n_batch = 0, 0, 0, 0 # for batch in tl.iterate.minibatches(inputs=X_test, targets=y_test, # batch_size=batch_size, shuffle=True): # b_images, b_labels = batch # _dice, _iou, _diceh, out = sess.run([test_dice_loss, # test_iou_loss, test_dice_hard, net_test.outputs], # {t_image: b_images, t_seg: b_labels}) # total_dice += _dice; total_iou += _iou; total_dice_hard += _diceh # n_batch += 1 # # print(" **"+" "*17+"test 1-dice: %f hard-dice: %f iou: %f (2d no distortion)" % # (total_dice/n_batch, total_dice_hard/n_batch, total_iou/n_batch)) # print(" task: {}".format(task)) # ## save a predition of test set # for i in range(batch_size): # if np.max(b_images[i]) > 0: # vis_imgs2(b_images[i], b_labels[i], out[i], "samples/{}/test_{}.png".format(task, epoch)) # break # elif i == batch_size-1: # vis_imgs2(b_images[i], b_labels[i], out[i], "samples/{}/test_{}.png".format(task, epoch)) ###======================== SAVE MODEL ==========================### tl.files.save_npz(net.all_params, name=save_dir+'/u_net_{}.npz'.format(task), sess=sess)
#print(x_offset, y_offset) result = im[x_offset:im.shape[0] - x_offset, y_offset:im.shape[1] - y_offset] return (result) #input placeholders X = tf.placeholder(tf.float32, shape=[None, img_height, img_width, None], name='X') y = tf.placeholder(tf.float32, shape=[None, img_height, img_width, None], name='y') logits = m.u_net(X) print("logits shape is {}".format(logits.get_shape())) #create session session = tf.Session() #init global vars init = tf.global_variables_initializer() #preidction prediction = tf.argmax(tf.nn.softmax(logits), axis=1) probability_map = tf.nn.softmax(logits) #correct prediction correct_prediction = tf.equal( prediction, tf.cast(tf.reshape(y, prediction.get_shape(), name=None), tf.int64)) #accuracy accuracy_c = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
def main(task='all'): ## Create folder to save trained model and result images save_dir = "checkpoint" tl.files.exists_or_mkdir(save_dir) tl.files.exists_or_mkdir("samples/{}".format(task)) ###======================== LOAD DATA ===================================### ## by importing this, you can load a training set and a validation set. # you will get X_train_input, X_train_target, X_dev_input and X_dev_target # there are 4 labels in targets: # Label 0: background # Label 1: necrotic and non-enhancing tumor # Label 2: edema # Label 4: enhancing tumor import prepare_data_with_valid as dataset X_train = dataset.X_train_input y_train = dataset.X_train_target[:,:,:,np.newaxis] X_test = dataset.X_dev_input y_test = dataset.X_dev_target[:,:,:,np.newaxis] if task == 'all': y_train = (y_train > 0).astype(int) y_test = (y_test > 0).astype(int) elif task == 'necrotic': y_train = (y_train == 1).astype(int) y_test = (y_test == 1).astype(int) elif task == 'edema': y_train = (y_train == 2).astype(int) y_test = (y_test == 2).astype(int) elif task == 'enhance': y_train = (y_train == 4).astype(int) y_test = (y_test == 4).astype(int) else: exit("Unknow task %s" % task) ###======================== HYPER-PARAMETERS ============================### batch_size = 10 lr = 0.0001 # lr_decay = 0.5 # decay_every = 100 beta1 = 0.9 n_epoch = 100 print_freq_step = 100 ###======================== SHOW DATA ===================================### # show one slice X = np.asarray(X_train[80]) y = np.asarray(y_train[80]) # print(X.shape, X.min(), X.max()) # (240, 240, 4) -0.380588 2.62761 # print(y.shape, y.min(), y.max()) # (240, 240, 1) 0 1 nw, nh, nz = X.shape vis_imgs(X, y, 'samples/{}/_train_im.png'.format(task)) # show data augumentation results for i in range(10): x_flair, x_t1, x_t1ce, x_t2, label = distort_imgs([X[:,:,0,np.newaxis], X[:,:,1,np.newaxis], X[:,:,2,np.newaxis], X[:,:,3,np.newaxis], y])#[:,:,np.newaxis]]) # print(x_flair.shape, x_t1.shape, x_t1ce.shape, x_t2.shape, label.shape) # (240, 240, 1) (240, 240, 1) (240, 240, 1) (240, 240, 1) (240, 240, 1) X_dis = np.concatenate((x_flair, x_t1, x_t1ce, x_t2), axis=2) # print(X_dis.shape, X_dis.min(), X_dis.max()) # (240, 240, 4) -0.380588233471 2.62376139209 vis_imgs(X_dis, label, 'samples/{}/_train_im_aug{}.png'.format(task, i)) with tf.device('/cpu:0'): sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) with tf.device('/gpu:0'): #<- remove it if you train on CPU or other GPU ###======================== DEFIINE MODEL =======================### ## nz is 4 as we input all Flair, T1, T1c and T2. t_image = tf.placeholder('float32', [batch_size, nw, nh, nz], name='input_image') ## labels are either 0 or 1 t_seg = tf.placeholder('float32', [batch_size, nw, nh, 1], name='target_segment') ## train inference net = model.u_net(t_image, is_train=True, reuse=False, n_out=1) ## test inference net_test = model.u_net(t_image, is_train=False, reuse=True, n_out=1) ###======================== DEFINE LOSS =========================### ## train losses out_seg = net.outputs dice_loss = 1 - tl.cost.dice_coe(out_seg, t_seg, axis=[0,1,2,3])#, 'jaccard', epsilon=1e-5) iou_loss = tl.cost.iou_coe(out_seg, t_seg, axis=[0,1,2,3]) dice_hard = tl.cost.dice_hard_coe(out_seg, t_seg, axis=[0,1,2,3]) loss = dice_loss ## test losses test_out_seg = net_test.outputs test_dice_loss = 1 - tl.cost.dice_coe(test_out_seg, t_seg, axis=[0,1,2,3])#, 'jaccard', epsilon=1e-5) test_iou_loss = tl.cost.iou_coe(test_out_seg, t_seg, axis=[0,1,2,3]) test_dice_hard = tl.cost.dice_hard_coe(test_out_seg, t_seg, axis=[0,1,2,3]) ###======================== DEFINE TRAIN OPTS =======================### t_vars = tl.layers.get_variables_with_name('u_net', True, True) with tf.device('/gpu:0'): with tf.variable_scope('learning_rate'): lr_v = tf.Variable(lr, trainable=False) train_op = tf.train.AdamOptimizer(lr_v, beta1=beta1).minimize(loss, var_list=t_vars) ###======================== LOAD MODEL ==============================### tl.layers.initialize_global_variables(sess) ## load existing model if possible tl.files.load_and_assign_npz(sess=sess, name=save_dir+'/u_net_{}.npz'.format(task), network=net) ###======================== TRAINING ================================### for epoch in range(0, n_epoch+1): epoch_time = time.time() ## update decay learning rate at the beginning of a epoch # if epoch !=0 and (epoch % decay_every == 0): # new_lr_decay = lr_decay ** (epoch // decay_every) # sess.run(tf.assign(lr_v, lr * new_lr_decay)) # log = " ** new learning rate: %f" % (lr * new_lr_decay) # print(log) # elif epoch == 0: # sess.run(tf.assign(lr_v, lr)) # log = " ** init lr: %f decay_every_epoch: %d, lr_decay: %f" % (lr, decay_every, lr_decay) # print(log) total_dice, total_iou, total_dice_hard, n_batch = 0, 0, 0, 0 for batch in tl.iterate.minibatches(inputs=X_train, targets=y_train, batch_size=batch_size, shuffle=True): images, labels = batch step_time = time.time() ## data augumentation for a batch of Flair, T1, T1c, T2 images # and label maps synchronously. data = tl.prepro.threading_data([_ for _ in zip(images[:,:,:,0, np.newaxis], images[:,:,:,1, np.newaxis], images[:,:,:,2, np.newaxis], images[:,:,:,3, np.newaxis], labels)], fn=distort_imgs) # (10, 5, 240, 240, 1) b_images = data[:,0:4,:,:,:] # (10, 4, 240, 240, 1) b_labels = data[:,4,:,:,:] b_images = b_images.transpose((0,2,3,1,4)) b_images.shape = (batch_size, nw, nh, nz) ## update network _, _dice, _iou, _diceh, out = sess.run([train_op, dice_loss, iou_loss, dice_hard, net.outputs], {t_image: b_images, t_seg: b_labels}) total_dice += _dice; total_iou += _iou; total_dice_hard += _diceh n_batch += 1 ## you can show the predition here: # vis_imgs2(b_images[0], b_labels[0], out[0], "samples/{}/_tmp.png".format(task)) # exit() # if _dice == 1: # DEBUG # print("DEBUG") # vis_imgs2(b_images[0], b_labels[0], out[0], "samples/{}/_debug.png".format(task)) if n_batch % print_freq_step == 0: print("Epoch %d step %d 1-dice: %f hard-dice: %f iou: %f took %fs (2d with distortion)" % (epoch, n_batch, _dice, _diceh, _iou, time.time()-step_time)) ## check model fail if np.isnan(_dice): exit(" ** NaN loss found during training, stop training") if np.isnan(out).any(): exit(" ** NaN found in output images during training, stop training") print(" ** Epoch [%d/%d] train 1-dice: %f hard-dice: %f iou: %f took %fs (2d with distortion)" % (epoch, n_epoch, total_dice/n_batch, total_dice_hard/n_batch, total_iou/n_batch, time.time()-epoch_time)) ## save a predition of training set for i in range(batch_size): if np.max(b_images[i]) > 0: vis_imgs2(b_images[i], b_labels[i], out[i], "samples/{}/train_{}.png".format(task, epoch)) break elif i == batch_size-1: vis_imgs2(b_images[i], b_labels[i], out[i], "samples/{}/train_{}.png".format(task, epoch)) ###======================== EVALUATION ==========================### total_dice, total_iou, total_dice_hard, n_batch = 0, 0, 0, 0 for batch in tl.iterate.minibatches(inputs=X_test, targets=y_test, batch_size=batch_size, shuffle=True): b_images, b_labels = batch _dice, _iou, _diceh, out = sess.run([test_dice_loss, test_iou_loss, test_dice_hard, net_test.outputs], {t_image: b_images, t_seg: b_labels}) total_dice += _dice; total_iou += _iou; total_dice_hard += _diceh n_batch += 1 print(" **"+" "*17+"test 1-dice: %f hard-dice: %f iou: %f (2d no distortion)" % (total_dice/n_batch, total_dice_hard/n_batch, total_iou/n_batch)) print(" task: {}".format(task)) ## save a predition of test set for i in range(batch_size): if np.max(b_images[i]) > 0: vis_imgs2(b_images[i], b_labels[i], out[i], "samples/{}/test_{}.png".format(task, epoch)) break elif i == batch_size-1: vis_imgs2(b_images[i], b_labels[i], out[i], "samples/{}/test_{}.png".format(task, epoch)) ###======================== SAVE MODEL ==========================### tl.files.save_npz(net.all_params, name=save_dir+'/u_net_{}.npz'.format(task), sess=sess)
smooth=1 def dice_coef(y_true, y_pred): y_true_f = K.flatten(y_true) y_pred_f = K.flatten(y_pred) intersection = K.sum(y_true_f * y_pred_f) return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth) def dice_coef_loss(y_true, y_pred): return -dice_coef(y_true, y_pred) input_shape = [256,256,1] dropout_rate = 0.3 l2_lambda = 0.0002 model = u_net(input_shape, dropout_rate, l2_lambda) model.summary() def preprocess(img_path,mask_path): images=[] a=[] print('\nLoading Volumes...') print('-'*30) for i in range(len(img_path)): a=nib.load(img_path[i]).get_data() print("image--%d--loaded"%(i)) a=np.resize(a,(a.shape[0],256,256)) a=a[:,:,:] for j in range(a.shape[0]): images.append((a[j,:,:]))
data_path = './data/MICCAI_BraTS17_Data_Training/HGG/' # ubuntu data path #data_path = './data/MICCAI_BraTS17_Data_Training_IPP/MICCAI_BraTS17_Data_Training/HGG/' #windows batch_size = 10 nw = 240 nh = 240 nz = 4 #init the network sess = tf.Session() t_image = tf.placeholder('float32', [batch_size, nw, nh, nz], name='input_image') ## labels are either 0 or 1 t_seg = tf.placeholder('float32', [batch_size, nw, nh, 1], name='target_segment') ## train inference net = model.u_net(t_image, is_train=False, reuse=False, n_out=1) net_seg = net.outputs sess = tf.Session() tl.layers.initialize_global_variables(sess) ## load existing model if possible tl.files.load_and_assign_npz(sess=sess, name=save_dir + '/u_net_{}.npz'.format(task), network=net) # read the files if not os.path.exists(save_seg_path): os.makedirs(save_seg_path) with open('HGG_list.txt', 'r') as f: files = f.readlines()
if y_.ndim == 2: y_ = y_[:, :, np.newaxis] assert X.ndim == 3 tl.vis.save_images(np.asarray([ X[:, :, 0, np.newaxis], X[:, :, 1, np.newaxis], X[:, :, 2, np.newaxis], X[:, :, 3, np.newaxis], y_, y ]), size=(1, 6), image_path=path) path = "E:/unet/checkpoint/" params = tl.files.load_npz(path, name='u_net.npz') x = tf.placeholder('float32', [batch_size, nw, nh, nz], name='input_image') y = tf.placeholder('float32', [batch_size, nw, nh, 1], name='target_segment') network = model.u_net(x, is_train=False, reuse=True) with tf.device('/cpu:0'): sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) tl.files.assign_params(sess, params, network) test_data_path = "E:/unet/testdata" test_path_list = tl.files.load_folder_list(path=test_data_path) test_name_list = [os.path.basename(p) for p in test_path_list] index_test = list(range(0, len(test_name_list))) data_types = ['flair', 't1', 't1ce', 't2'] data_types_mean_std_dict = {i: {'mean': 0.0, 'std': 1.0} for i in data_types} for i in data_types: data_temp_list = [] for j in test_name_list: img_path = os.path.join(test_data_path, j, j + '_' + i + '.nii.gz') img = nib.load(img_path).get_data()