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