import os
import glob
import time
import shutil
import numpy as np
from collections import OrderedDict
import __future__
import logging
import matplotlib
import tensorflow as tf
import csv
from tensorflow.python import debug as tf_debug
from layers import *
from ops import *
from lib import _dice_eval, _save, _save_nii_prediction, _jaccard, _dice, _label_decomp, _indicator_eval, read_nii_image
np.random.seed(0)
contour_map = { # a map used for mapping label value to its name, used for output
"bg": 0,
"la_myo": 1,
"la_blood": 2,
"lv_blood": 3,
"aa": 4
}
verbose = True
logging.basicConfig(filename = "curr_log", level=logging.DEBUG, format='%(asctime)s %(message)s')
if verbose == True:
logging.getLogger().addHandler(logging.StreamHandler())
raw_size = [256, 256, 3] # original raw input size
volume_size = [256, 256, 3] # volume size after processing, for the tfrecord file
label_size = [256, 256, 1] # size of label
decomp_feature = { # configuration for decoding tf_record file
'dsize_dim0': tf.FixedLenFeature([], tf.int64),
'dsize_dim1': tf.FixedLenFeature([], tf.int64),
'dsize_dim2': tf.FixedLenFeature([], tf.int64),
'lsize_dim0': tf.FixedLenFeature([], tf.int64),
'lsize_dim1': tf.FixedLenFeature([], tf.int64),
'lsize_dim2': tf.FixedLenFeature([], tf.int64),
'data_vol': tf.FixedLenFeature([], tf.string),
'label_vol': tf.FixedLenFeature([], tf.string)}
class Full_DRN(object):
def __init__(self, channels, n_class, batch_size, cost_kwargs={}, network_config = {}):
##### Done this function
tf.reset_default_graph()
self.n_class = n_class # please note background is another class
self.batch_size = batch_size
self.mr_front_weights = [] # conv weights of MR path
self.ct_front_weights = [] # conv weights of CT path
self.cls_weights = [] # weights of feature discriminator
self.m_cls_weights = [] # weights for segmentation mask discriminator
self.joint_weights = [] # weights of joint part between CT and MRI. The final segmentor in our case
self.mr = tf.placeholder("float", shape=[None, volume_size[0], volume_size[1], channels], name = "mr_ph")
self.ct = tf.placeholder("float", shape=[None, volume_size[0], volume_size[1], channels])
self.ct_y = tf.placeholder("float", shape=[None, label_size[0], label_size[1], self.n_class])
self.mr_y = tf.placeholder("float", shape=[None, label_size[0], label_size[1], self.n_class])
self.mr_front_bn = tf.placeholder_with_default(False, shape = None, name = "main_batchnorm_training_switch")
self.joint_bn = tf.placeholder_with_default(False, shape = None, name = "joint_batchnorm_training_switch")
self.ct_front_bn = tf.placeholder_with_default(True, shape = None, name = "adapt_batchnorm_training_switch")
# these two are useless. They are not passed into the program
self.cls_bn = tf.placeholder_with_default(True, shape = None, name = "cls_batchnorm_training_switch")
self.m_cls_bn = tf.placeholder_with_default(True, shape = None, name = "mask_cls_batchnorm_training_switch")
self.network_config = network_config
self.mr_front_trainable = self.network_config["mr_front_trainable"]
self.ct_front_trainable = self.network_config["ct_front_trainable"]
self.joint_trainable = self.network_config["joint_trainable"]
self.cls_trainable = self.network_config["cls_trainable"]
self.m_cls_trainable = self.network_config["m_cls_trainable"]
self.keep_prob = tf.placeholder(tf.float32) # dropout keep probability
# Get features from MRI and CT path, for early layers
_mr_c4_2, _ct_c4_2, _mr_c6_2, _ct_c6_2 = self.create_zip_network(input_channel = channels,\
feature_base = 16, num_cls = n_class, keep_prob = self.keep_prob,\
main_bn = self.mr_front_bn, main_trainable = self.mr_front_trainable,\
adapt_bn = self.ct_front_bn, adapt_trainable = self.ct_front_trainable)
# Get features from MRI and CT, fromt the shared higher layers
with tf.variable_scope("", reuse = tf.AUTO_REUSE) as scope:
_ct_c9_2, _ct_b8, _ct_b7, _ct_logits = self.create_second_half( _ct_c6_2, feature_base = 16, input_channel = 3, num_cls = n_class, keep_prob = self.keep_prob, joint_bn = self.joint_bn, joint_trainable = self.joint_trainable)
_mr_c9_2, _mr_b8, _mr_b7, _mr_logits = self.create_second_half( _mr_c6_2, feature_base = 16, input_channel = 3, num_cls = n_class, keep_prob = self.keep_prob, joint_bn = self.joint_bn, joint_trainable = self.joint_trainable)
self.ct_conv9_2 = _ct_c9_2
self.mr_conv9_2 = _mr_c9_2
with tf.variable_scope("cls_scope", reuse = tf.AUTO_REUSE) as scope:
self._ct_class_logits = self.create_classifier(_ct_c4_2, _ct_c6_2, _ct_b7, _ct_c9_2, _ct_logits)
self._mr_class_logits = self.create_classifier(_mr_c4_2, _mr_c6_2, _mr_b7, _mr_c9_2, _mr_logits)
self.predictor = pixel_wise_softmax_2(_ct_logits) # segmentation logits of CT
self.compact_pred = tf.argmax(self.predicter, 3) # predictions
self.compact_y = tf.argmax(self.ct_y, 3) # ground truth
self.ct_dice_eval, self.ct_dice_eval_arr = _dice_eval(self.compact_pred, self.ct_y, self.n_class) # used for monitoring training process
self.ct_dice_eval_c1 = self.ct_dice_eval_arr[1]
self.ct_dice_eval_c2 = self.ct_dice_eval_arr[2]
self.ct_dice_eval_c3 = self.ct_dice_eval_arr[3]
self.ct_dice_eval_c4 = self.ct_dice_eval_arr[4]
self.mr_seg_valid = pixel_wise_softmax_2(_mr_logits) # segmentation logits of MRI
self.compact_mr_valid = tf.argmax(self.mr_seg_valid, 3)
self.compact_mr_y = tf.argmax(self.mr_y, 3)
self.mr_dice_eval, self.mr_dice_eval_arr = _dice_eval(self.compact_mr_valid, self.mr_y, self.n_class)
with tf.variable_scope("mask_cls_scope", reuse = tf.AUTO_REUSE) as scope:
self._ct_mask_logits = self.create_mask_critic(_ct_logits, num_cls = n_class) # auxilary D loss for masks
self._mr_mask_logits = self.create_mask_critic(_mr_logits, num_cls = n_class)
self.cost_kwargs = cost_kwargs
self.dis_loss, self.ct_gen_loss, self.fixed_coeff_reg, self.dis_reg, self.gen_reg = self._get_cost(_ct_logits, _mr_logits, self._ct_class_logits, self._mr_class_logits,\
self._ct_mask_logits, self._mr_mask_logits, self.cost_kwargs) # get cost
self.confusion_matrix = tf.confusion_matrix( tf.reshape(self.compact_y,[-1]), tf.reshape(self.compact_pred, [-1]), num_classes = self.n_class )
def create_zip_network(self, main_bn, main_trainable, adapt_bn, adapt_trainable, num_cls, feature_base = 16, input_channel = 3, keep_prob = 0.75):
# MR path starts from here
with tf.variable_scope('group_1') as scope:
w1_1 = weight_variable(shape = [3, 3, input_channel, feature_base], trainable = main_trainable)
conv1_1 = conv2d(self.mr, w1_1, keep_prob )
wr1_1 = weight_variable(shape = [ 3, 3, feature_base,feature_base], trainable = main_trainable)
wr1_2 = weight_variable(shape = [3, 3, feature_base, feature_base], trainable = main_trainable)
block1_1 = residual_block(conv1_1, wr1_1, wr1_2, keep_prob , is_train = main_bn, leak = True, bn_trainable = main_trainable , scope = 'pred_1_1' ) # here the scope is for bn
out1 = max_pool2d(block1_1, n = 2)
self.mr_front_weights.append(w1_1)
self.mr_front_weights.append(wr1_1)
self.mr_front_weights.append(wr1_2)
with tf.variable_scope('group_2') as scope:
wr2_1 = weight_variable(shape = [3, 3, feature_base, feature_base * 2], trainable = main_trainable)
wr2_2 = weight_variable(shape = [3, 3, feature_base * 2, feature_base * 2], trainable = main_trainable)
block2_1 = residual_block(out1, wr2_1, wr2_2, inc_dim = True,keep_prob = keep_prob, leak = True, is_train = main_bn, bn_trainable = main_trainable, scope = 'pred_2_1' )
out2 = max_pool2d(block2_1, n = 2)
self.mr_front_weights.append(wr2_1)
self.mr_front_weights.append(wr2_2)
with tf.variable_scope('group_3') as scope:
wr3_1 = weight_variable( shape = [3, 3, feature_base * 2, feature_base * 4], trainable = main_trainable )
wr3_2 = weight_variable( shape = [3, 3, feature_base * 4, feature_base * 4], trainable = main_trainable )
block3_1 = residual_block( out2, wr3_1, wr3_2, keep_prob, inc_dim = True, is_train = main_bn, leak = True, bn_trainable = main_trainable , scope = 'pred_3_1' )
wr3_3 = weight_variable( shape = [3, 3, feature_base * 4, feature_base * 4], trainable = main_trainable )
wr3_4 = weight_variable( shape = [3, 3, feature_base * 4, feature_base * 4], trainable = main_trainable )
block3_2 = residual_block( block3_1, wr3_3, wr3_4,keep_prob = keep_prob, is_train = main_bn, leak = True, bn_trainable = main_trainable , scope = 'pred_3_2' )
out3 = max_pool2d(block3_2, n = 2)
self.mr_front_weights.append(wr3_1)
self.mr_front_weights.append(wr3_2)
self.mr_front_weights.append(wr3_3)
self.mr_front_weights.append(wr3_4)
with tf.variable_scope('group_4') as scope:
wr4_1 = weight_variable( shape = [3, 3, feature_base * 4, feature_base * 8], trainable = main_trainable )
wr4_2 = weight_variable( shape = [3, 3, feature_base * 8, feature_base * 8], trainable = main_trainable )
block4_1 = residual_block( out3, wr4_1, wr4_2, keep_prob, inc_dim = True, is_train = main_bn, leak = True, bn_trainable = main_trainable , scope = 'pred_4_1' )
wr4_3 = weight_variable( shape = [3, 3, feature_base * 8, feature_base * 8], trainable = main_trainable )
wr4_4 = weight_variable( shape = [3, 3, feature_base * 8, feature_base * 8], trainable = main_trainable )
block4_2 = residual_block( block4_1, wr4_3, wr4_4, keep_prob, is_train = main_bn, leak = True, bn_trainable = main_trainable , scope = 'pred_4_2' )
self.mr_front_weights.append(wr4_1)
self.mr_front_weights.append(wr4_2)
self.mr_front_weights.append(wr4_3)
self.mr_front_weights.append(wr4_4)
with tf.variable_scope('group_5') as scope:
wr5_1 = sharable_weight_variable( shape = [3, 3, feature_base * 8, feature_base * 16], trainable = main_trainable, name = "Variable" )
wr5_2 = sharable_weight_variable( shape = [3, 3, feature_base * 16, feature_base * 16], trainable = main_trainable , name = "Variable_1" )
block5_1 = residual_block( block4_2, wr5_1, wr5_2, keep_prob = keep_prob, inc_dim = True, leak = True, is_train = main_bn, bn_trainable = main_trainable, scope = 'pred_5_1' )
wr5_3 = sharable_weight_variable( shape = [3, 3, feature_base * 16, feature_base * 16], trainable = main_trainable , name = "Variable_2" )
wr5_4 = sharable_weight_variable( shape = [3, 3, feature_base * 16, feature_base * 16], trainable = main_trainable , name = "Variable_3" )
block5_2 = residual_block( block5_1, wr5_3, wr5_4, keep_prob = keep_prob, is_train = main_bn, leak = True, bn_trainable = main_trainable , scope = 'pred_5_2' )
self.mr_front_weights.append( wr5_1 )
self.mr_front_weights.append( wr5_2 )
self.mr_front_weights.append( wr5_3 )
self.mr_front_weights.append( wr5_4 )
with tf.variable_scope('group_6') as scope:
wr6_1 = sharable_weight_variable( shape = [3, 3, feature_base * 16, feature_base * 16], trainable = main_trainable , name = "Variable" )
wr6_2 = sharable_weight_variable( shape = [3, 3, feature_base * 16, feature_base * 16], trainable = main_trainable , name = "Variable_1" )
block6_1 = residual_block( block5_2, wr6_1, wr6_2, keep_prob = keep_prob, is_train = main_bn, leak = True, bn_trainable = main_trainable , scope = 'pred_6_1' )
wr6_3 = sharable_weight_variable( shape = [3, 3, feature_base * 16, feature_base * 16], trainable = main_trainable , name = "Variable_2" )
wr6_4 = sharable_weight_variable( shape = [3, 3, feature_base * 16, feature_base * 16], trainable = main_trainable, name = "Variable_3" )
block6_2 = residual_block( block6_1, wr6_3, wr6_4, keep_prob = keep_prob, is_train = main_bn, leak = True, bn_trainable = main_trainable , scope = 'pred_6_2' )
self.mr_front_weights.append( wr6_1 )
self.mr_front_weights.append( wr6_2 )
self.mr_front_weights.append( wr6_3 )
self.mr_front_weights.append( wr6_4 )
# DAM for CT path starts from here
with tf.variable_scope('adapt_1') as scope:
w1_1a = sharable_weight_variable(shape = [3, 3, input_channel, feature_base ], trainable = adapt_trainable, name = "Variable")
conv1_1a = conv2d(self.ct, w1_1a, keep_prob )
wr1_1a = sharable_weight_variable(shape = [ 3, 3, feature_base ,feature_base ], trainable = adapt_trainable, name = "Variable_1")
wr1_2a = sharable_weight_variable(shape = [3, 3, feature_base , feature_base ], trainable = adapt_trainable, name = "Variable_2")
block1_1a = residual_block(conv1_1a, wr1_1a, wr1_2a, keep_prob , is_train = adapt_bn, leak = True, bn_trainable = adapt_trainable, scope = 'adapt_1' )
out1a = max_pool2d(block1_1a, n = 2)
self.ct_front_weights.append(w1_1a)
self.ct_front_weights.append(wr1_1a)
self.ct_front_weights.append(wr1_2a)
with tf.variable_scope('adapt_2') as scope:
wr2_1a = sharable_weight_variable(shape = [3, 3, feature_base , feature_base * 2], trainable = adapt_trainable, name = "Variable")
wr2_2a = sharable_weight_variable(shape = [3, 3, feature_base * 2, feature_base * 2], trainable = adapt_trainable, name = "Variable_1")
block2_1a = residual_block(out1a, wr2_1a, wr2_2a, inc_dim = True,keep_prob = keep_prob, leak = True, is_train = adapt_bn, bn_trainable = adapt_trainable, scope = 'adapt_2' )
out2a = max_pool2d(block2_1a, n = 2)
self.ct_front_weights.append(wr2_1a)
self.ct_front_weights.append(wr2_2a)
with tf.variable_scope('adapt_3') as scope:
wr3_1a = sharable_weight_variable( shape = [3, 3, feature_base * 2, feature_base * 4], trainable = adapt_trainable, name = "Variable" )
wr3_2a = sharable_weight_variable( shape = [3, 3, feature_base * 4, feature_base * 4], trainable = adapt_trainable, name = "Variable_1" )
block3_1a = residual_block( out2a, wr3_1a, wr3_2a, keep_prob, inc_dim = True, leak = True, is_train = adapt_bn, bn_trainable = adapt_trainable , scope = 'adapt_3_1' )
wr3_3a = sharable_weight_variable( shape = [3, 3, feature_base * 4, feature_base * 4], trainable = adapt_trainable, name = "Variable_2" )
wr3_4a = sharable_weight_variable( shape = [3, 3, feature_base * 4, feature_base * 4], trainable = adapt_trainable , name = "Variable_3" )
block3_2a = residual_block( block3_1a, wr3_3a, wr3_4a,keep_prob = keep_prob, leak = True, is_train = adapt_bn, bn_trainable = adapt_trainable, scope = 'adapt_3_2' )
out3a = max_pool2d(block3_2a, n = 2)
self.ct_front_weights.append(wr3_1a)
self.ct_front_weights.append(wr3_2a)
self.ct_front_weights.append(wr3_3a)
self.ct_front_weights.append(wr3_4a)
with tf.variable_scope('adapt_4') as scope:
wr4_1a = sharable_weight_variable( shape = [3, 3, feature_base * 4, feature_base * 8], trainable = adapt_trainable, name = "Variable" )
wr4_2a = sharable_weight_variable( shape = [3, 3, feature_base * 8, feature_base * 8], trainable = adapt_trainable , name = "Variable_1" )
block4_1a = residual_block( out3a, wr4_1a, wr4_2a, keep_prob, inc_dim = True, leak = True, is_train = adapt_bn, bn_trainable = adapt_trainable, scope = 'adapt_4_1' )
wr4_3a = sharable_weight_variable( shape = [3, 3, feature_base * 8, feature_base * 8], trainable = adapt_trainable , name = "Variable_2" )
wr4_4a = sharable_weight_variable( shape = [3, 3, feature_base * 8, feature_base * 8], trainable = adapt_trainable , name = "Variable_3" )
block4_2a = residual_block( block4_1a, wr4_3a, wr4_4a, keep_prob, is_train = adapt_bn, leak = True, bn_trainable = adapt_trainable, scope = 'adapt_4_2' )
self.ct_front_weights.append(wr4_1a)
self.ct_front_weights.append(wr4_2a)
self.ct_front_weights.append(wr4_3a)
self.ct_front_weights.append(wr4_4a)
with tf.variable_scope('adapt_5') as scope:
wr5_1a = sharable_weight_variable( shape = [3, 3, feature_base * 8, feature_base * 16], trainable = adapt_trainable, name = "Variable" )
wr5_2a = sharable_weight_variable( shape = [3, 3, feature_base * 16, feature_base * 16], trainable = adapt_trainable , name = "Variable_1" )
block5_1a = residual_block( block4_2a, wr5_1a, wr5_2a, keep_prob = keep_prob, leak = True, inc_dim = True, is_train = adapt_bn, bn_trainable = adapt_trainable, scope = 'adapt_5_1' )
wr5_3a = sharable_weight_variable( shape = [3, 3, feature_base * 16, feature_base * 16], trainable = adapt_trainable , name = "Variable_2" )
wr5_4a = sharable_weight_variable( shape = [3, 3, feature_base * 16, feature_base * 16], trainable = adapt_trainable , name = "Variable_3" )
block5_2a = residual_block( block5_1a, wr5_3a, wr5_4a, keep_prob = keep_prob, leak = True, is_train = adapt_bn, bn_trainable = adapt_trainable , scope = 'adapt_5_2' )
self.ct_front_weights.append( wr5_1a )
self.ct_front_weights.append( wr5_2a )
self.ct_front_weights.append( wr5_3a )
self.ct_front_weights.append( wr5_4a )
with tf.variable_scope('adapt_6') as scope:
wr6_1a = sharable_weight_variable( shape = [3, 3, feature_base * 16, feature_base * 16], trainable = adapt_trainable , name = "Variable" )
wr6_2a = sharable_weight_variable( shape = [3, 3, feature_base * 16, feature_base * 16], trainable = adapt_trainable , name = "Variable_1" )
block6_1a = residual_block( block5_2a, wr6_1a, wr6_2a, keep_prob = keep_prob, leak = True, is_train = adapt_bn, bn_trainable = adapt_trainable , scope = 'adapt_6_1' )
wr6_3a = sharable_weight_variable( shape = [3, 3, feature_base * 16, feature_base * 16], trainable = adapt_trainable , name = "Variable_2" )
wr6_4a = sharable_weight_variable( shape = [3, 3, feature_base * 16, feature_base * 16], trainable = adapt_trainable, name = "Variable_3" )
block6_2a = residual_block( block6_1a, wr6_3a, wr6_4a, keep_prob = keep_prob, leak = True, is_train = adapt_bn, bn_trainable = adapt_trainable , scope = 'adapt_6_2' )
self.ct_front_weights.append( wr6_1a )
self.ct_front_weights.append( wr6_2a )
self.ct_front_weights.append( wr6_3a )
self.ct_front_weights.append( wr6_4a )
return block4_2, block4_2a, block6_2, block6_2a
def create_second_half(self, input_feature, joint_bn, joint_trainable, num_cls, feature_base = 16, input_channel = 3, keep_prob = 0.75):
with tf.variable_scope('group_7', reuse = tf.AUTO_REUSE) as scope:
wr7_1 = sharable_weight_variable( shape = [3, 3, feature_base * 16, feature_base * 32], trainable = joint_trainable , name = "Variable" )
wr7_2 = sharable_weight_variable( shape = [3, 3, feature_base * 32, feature_base * 32], trainable = joint_trainable , name = "Variable_1" )
block7_1 = residual_block( input_feature, wr7_1, wr7_2, keep_prob = keep_prob, leak = True, inc_dim = True, is_train = joint_bn, bn_trainable = joint_trainable , scope = 'pred_7_1' )
wr7_3 = sharable_weight_variable( shape = [3, 3, feature_base * 32, feature_base * 32], trainable = joint_trainable , name = "Variable_2" )
wr7_4 = sharable_weight_variable( shape = [3, 3, feature_base * 32, feature_base * 32], trainable = joint_trainable , name = "Variable_3" )
block7_2 = residual_block( block7_1, wr7_3, wr7_4, keep_prob = keep_prob, leak = True, is_train = joint_bn, bn_trainable = joint_trainable , scope = 'pred_7_2' )
self.mr_front_weights.append( wr7_1 )
self.mr_front_weights.append( wr7_2 )
self.mr_front_weights.append( wr7_3 )
self.mr_front_weights.append( wr7_4 )
with tf.variable_scope('group_8', reuse = tf.AUTO_REUSE) as scope:
wr8_1 = sharable_weight_variable( shape = [3, 3, feature_base * 32, feature_base * 32], trainable = joint_trainable , name = "Variable" )
wr8_2 = sharable_weight_variable( shape = [3, 3, feature_base * 32, feature_base * 32], trainable = joint_trainable , name = "Variable_1" )
block8_1 = DR_block( block7_2, wr8_1, wr8_2, keep_prob = keep_prob, leak = True, is_train = joint_bn, rate = 2, bn_trainable = joint_trainable , scope = 'pred_8_1' )
wr8_3 = sharable_weight_variable( shape = [3, 3, feature_base * 32, feature_base * 32], trainable = joint_trainable , name = "Variable_2" )
wr8_4 = sharable_weight_variable( shape = [3, 3, feature_base * 32, feature_base * 32], trainable = joint_trainable , name = "Variable_3" )
block8_2 = DR_block( block8_1, wr8_3, wr8_4, keep_prob = keep_prob, leak = True, is_train = joint_bn, rate = 2, bn_trainable = joint_trainable , scope = 'pred_8_2' )
self.mr_front_weights.append( wr8_1 )
self.mr_front_weights.append( wr8_2 )
self.mr_front_weights.append( wr8_3 )
self.mr_front_weights.append( wr8_4 )
with tf.variable_scope('group_9', reuse = tf.AUTO_REUSE) as scope:
w9_1 = sharable_weight_variable( shape = [3, 3, feature_base * 32, feature_base * 32], trainable = joint_trainable , name = "Variable" )
conv9_1 = conv_bn_relu2d( block8_2, w9_1, keep_prob, leak = True, is_train = joint_bn, bn_trainable = joint_trainable , scope = 'pred_9_1' )
w9_2 = sharable_weight_variable( shape = [3, 3, feature_base * 32, feature_base * 32], trainable = joint_trainable , name = "Variable_1" )
conv9_2 = conv_bn_relu2d( conv9_1, w9_2, keep_prob, leak = True, is_train = joint_bn, bn_trainable = joint_trainable , scope = 'pred_9_2' )
self.mr_front_weights.append( w9_1 )
self.mr_front_weights.append( w9_2 )
with tf.variable_scope('group_10', reuse = tf.AUTO_REUSE) as scope:
local_size = 8 * 8
w10_1 = sharable_weight_variable( shape = [3, 3, feature_base * 32, local_size * num_cls * 8], trainable = joint_trainable , name = "Variable" )
conv10_1 = conv2d( conv9_2, w10_1, keep_prob_ = keep_prob, padding = 'SYMMETRIC')
self.mr_front_weights.append(w10_1)
flat_conv10_1 = PS(conv10_1, r = 8, n_channel = num_cls * 8, batch_size = self.batch_size) # phase shift
with tf.variable_scope('output', reuse = tf.AUTO_REUSE) as scope:
w11_1 = sharable_weight_variable( shape = [5, 5, num_cls * 8, num_cls], trainable = joint_trainable , name = "Variable" )
logits = conv2d( flat_conv10_1, w11_1, keep_prob_ = 1., padding = 'SYMMETRIC' )
return conv9_2, block8_2, block7_2, logits
def create_classifier(self, input_conv4, input_conv6, input_b7, input_conv9, seg_logits, feature_base = 16, keep_prob = 0.75, cls_bn = True, cls_trainable = True):
"""
domain discriminator for MRI features and CT features
"""
with tf.variable_scope('cls_0') as scope:
flat_input_conv4 = PS(input_conv4, r=8, n_channel=2, batch_size=self.batch_size) # 2
flat_input_conv4 = tf.tile(flat_input_conv4, [1, 1, 1, 3]) # 6 in total
flat_input_conv6 = PS(input_conv6, r=8, n_channel=4, batch_size=self.batch_size) # 10 in total
flat_input_b7 = PS(input_b7, r=8, n_channel=8, batch_size=self.batch_size) # 18 in total
flat_input_conv9 = PS(input_conv9, r = 8, n_channel = 8, batch_size = self.batch_size) # 26 in total
input_comp = simple_concat2d(flat_input_conv4, flat_input_conv6) # 10
input_comp = simple_concat2d(input_comp, flat_input_b7) # 18
input_comp = simple_concat2d(input_comp, flat_input_conv9) # 26
input_comp = simple_concat2d(input_comp, seg_logits) # 31 in total
input_comp = simple_concat2d(input_comp, tf.expand_dims(tf.cast(tf.argmax(seg_logits, 3), tf.float32), 3)) # 1
with tf.variable_scope('cls_1') as scope:
wr1_1c = sharable_weight_variable( shape = [3, 3, feature_base * 2, feature_base * 4], trainable = cls_trainable , name = "Variable" )
wr1_2c = sharable_weight_variable( shape = [3, 3, feature_base * 4, feature_base * 4], trainable = cls_trainable , name = "Variable_1" )
block1_1c = residual_block( input_comp, wr1_1c, wr1_2c, keep_prob = keep_prob, inc_dim = True, is_train = cls_bn, bn_trainable = cls_trainable, scope = 'cls_1' , leak = True )
wr1_3d = sharable_weight_variable( shape = [3,3, feature_base * 4, feature_base * 4], trainable = cls_trainable, name = "Variable_2" )
out1c = conv_bn_relu2d( block1_1c, wr1_3d, keep_prob, strides = [1,2,2,1], is_train = cls_bn, bn_trainable = cls_trainable, scope = 'cls_1_3', leak = True )
self.cls_weights.append( wr1_1c )
self.cls_weights.append( wr1_2c )
self.cls_weights.append( wr1_3d )
with tf.variable_scope('cls_2') as scope:
wr2_1c = sharable_weight_variable( shape = [3, 3, feature_base * 4, feature_base *8], trainable = cls_trainable , name = "Variable" )
wr2_2c = sharable_weight_variable( shape = [3, 3, feature_base * 8, feature_base *8], trainable = cls_trainable , name = "Variable_1" )
block2_1c = residual_block( out1c, wr2_1c, wr2_2c, keep_prob = keep_prob, inc_dim = True, is_train = cls_bn, bn_trainable = cls_trainable, scope = 'cls_2' , leak = True )
wr2_3d = sharable_weight_variable( shape = [5,5, feature_base * 8, feature_base * 8], trainable = cls_trainable, name = "Variable_2" )
out2c = conv_bn_relu2d( block2_1c, wr2_3d, keep_prob, strides = [1,2,2,1], is_train = cls_bn, bn_trainable = cls_trainable, scope = 'cls_2_3', leak = True )
self.cls_weights.append( wr2_1c )
self.cls_weights.append( wr2_2c )
self.cls_weights.append( wr2_3d )
self.debug_out2c = out2c
self.debug_wr2_2c = wr2_2c
with tf.variable_scope('cls_3') as scope:
wr3_1c = sharable_weight_variable( shape = [3, 3, feature_base * 8, feature_base *16], trainable = cls_trainable , name = "Variable" )
wr3_2c = sharable_weight_variable( shape = [3, 3, feature_base * 16, feature_base *16], trainable = cls_trainable , name = "Variable_1" )
block3_1c = residual_block( out2c, wr3_1c, wr3_2c, keep_prob = keep_prob, inc_dim = True, is_train = cls_bn, bn_trainable = cls_trainable, scope = 'cls_3' , leak = True )
wr3_3d = sharable_weight_variable( shape = [3,3, feature_base * 16, feature_base * 16], trainable = cls_trainable, name = "Variable_2" )
out3c = conv_bn_relu2d( block3_1c, wr3_3d, keep_prob, strides = [1,2,2,1], is_train = cls_bn, bn_trainable = cls_trainable, scope = 'cls_3_3', leak = True )
self.cls_weights.append( wr3_1c )
self.cls_weights.append( wr3_2c )
self.cls_weights.append( wr3_3d )
with tf.variable_scope('cls_4') as scope:
wr4_1c = sharable_weight_variable( shape = [3, 3, feature_base * 16, feature_base *32], trainable = cls_trainable , name = "Variable" )
wr4_2c = sharable_weight_variable( shape = [3, 3, feature_base * 32, feature_base *32], trainable = cls_trainable , name = "Variable_1" )
block4_1c = residual_block( out3c, wr4_1c, wr4_2c, keep_prob = keep_prob, inc_dim = True, is_train = cls_bn, bn_trainable = cls_trainable, scope = 'cls_4' , leak = True )
wr4_3d = sharable_weight_variable( shape = [3,3, feature_base * 32, feature_base * 32], trainable = cls_trainable, name = "Variable_2" )
out4c = conv_bn_relu2d( block4_1c, wr4_3d, keep_prob, strides = [1,2,2,1], is_train = cls_bn, bn_trainable = cls_trainable, scope = 'cls_4_3', leak = True )
self.cls_weights.append( wr4_1c )
self.cls_weights.append( wr4_2c )
self.cls_weights.append( wr4_3d )
with tf.variable_scope('cls_5') as scope:
wr5_1c = sharable_weight_variable( shape = [3, 3, feature_base * 32, feature_base *32], trainable = cls_trainable , name = "Variable" )
wr5_2c = sharable_weight_variable( shape = [3, 3, feature_base * 32, feature_base *32], trainable = cls_trainable , name = "Variable_1" )
block5_1c = residual_block( out4c, wr5_1c, wr5_2c, keep_prob = keep_prob, is_train = cls_bn, bn_trainable = cls_trainable, scope = 'cls_5' , leak = True )
wr5_3d = sharable_weight_variable( shape = [5,5, feature_base * 32, feature_base * 32], trainable = cls_trainable, name = "Variable_2" )
out5c = conv_bn_relu2d( block5_1c, wr5_3d, keep_prob, strides = [1,4,4,1], is_train = cls_bn, bn_trainable = cls_trainable, scope = 'cls_5_3', leak = True )
self.cls_weights.append( wr5_1c )
self.cls_weights.append( wr5_2c )
self.cls_weights.append( wr5_3d )
with tf.variable_scope('cls_6') as scope:
wr6_1c = sharable_weight_variable( shape = [3, 3, feature_base * 32, feature_base *32], trainable = cls_trainable , name = "Variable" )
conv_6c = conv_bn_relu2d(out5c, wr6_1c, strides = [1,2,2,1], keep_prob = keep_prob, padding = "SYMMETRIC", scope = 'cls_6', is_train = cls_bn, bn_trainable = cls_trainable, leak = True)
self.cls_weights.append( wr6_1c )
with tf.variable_scope('cls_out') as scope:
wc_out = sharable_weight_variable( shape = [ feature_base* 32 * 4,1 ], trainable = cls_trainable , name = "Variable" )
out6c_flat = tf.reshape(conv_6c, [-1, feature_base * 32 * 4])
cls_logits = tf.matmul(out6c_flat, wc_out)
self.cls_weights.append(wc_out)
return cls_logits
def create_mask_critic(self, input_mask, feature_base = 16, keep_prob = 0.75, num_cls = 5, m_cls_bn = True, m_cls_trainable = True):
"""
domain discriminator for MRI and CT segmentation maskS
"""
with tf.variable_scope('mask_cls_1') as scope:
wr1_1m = sharable_weight_variable( shape = [3, 3, num_cls, feature_base], trainable = m_cls_trainable , name = "Variable" )
out1m = conv_bn_relu2d( input_mask, wr1_1m, keep_prob, strides = [1,2,2,1], is_train = m_cls_bn, bn_trainable = m_cls_trainable, scope = 'mask_cls_1', leak = True ) # use strided conv instead of maxpool to
self.m_cls_weights.append( wr1_1m )
with tf.variable_scope('mask_cls_2') as scope:
wr2_1m = sharable_weight_variable( shape = [3, 3, feature_base, feature_base ], trainable = m_cls_trainable , name = "Variable" )
wr2_2m = sharable_weight_variable( shape = [3, 3, feature_base, feature_base ], trainable = m_cls_trainable , name = "Variable_1" )
block2_1m = residual_block( out1m, wr2_1m, wr2_2m, keep_prob = keep_prob, inc_dim = False, is_train = m_cls_bn, bn_trainable = m_cls_trainable, scope = 'm_cls_2' , leak = True )
wr2_3d = sharable_weight_variable( shape = [5,5, feature_base, feature_base * 2], trainable = m_cls_trainable, name = "Variable_2" )
out2m = conv_bn_relu2d( block2_1m, wr2_3d, keep_prob, strides = [1,4,4,1], is_train = m_cls_bn, bn_trainable = m_cls_trainable, scope = 'm_cls_2_3', leak = True )
self.m_cls_weights.append( wr2_1m )
self.m_cls_weights.append( wr2_2m )
self.m_cls_weights.append( wr2_3d )
with tf.variable_scope('mask_cls_3') as scope:
wr3_1m = sharable_weight_variable( shape = [3, 3, feature_base * 2, feature_base * 4], trainable = m_cls_trainable , name = "Variable" )
wr3_2m = sharable_weight_variable( shape = [3, 3, feature_base * 4, feature_base * 4 ], trainable = m_cls_trainable , name = "Variable_1" )
block3_1m = residual_block( out2m, wr3_1m, wr3_2m, keep_prob = keep_prob, inc_dim = True, is_train = m_cls_bn, bn_trainable = m_cls_trainable, scope = 'm_cls_3' , leak = True )
wr3_3d = sharable_weight_variable( shape = [5,5, feature_base * 4, feature_base * 8], trainable = m_cls_trainable, name = "Variable_2" )
out3m = conv_bn_relu2d( block3_1m, wr3_3d, keep_prob, strides = [1,4,4,1], is_train = m_cls_bn, bn_trainable = m_cls_trainable, scope = 'm_cls_3_3', leak = True )
self.m_cls_weights.append( wr3_1m )
self.m_cls_weights.append( wr3_2m )
self.m_cls_weights.append( wr3_3d )
with tf.variable_scope('mask_cls_4') as scope:
wr4_1m = sharable_weight_variable( shape = [5, 5, feature_base * 8, feature_base * 16], trainable = m_cls_trainable , name = "Variable" )
conv_4m = conv_bn_relu2d(out3m, wr4_1m, strides = [1,4,4,1], keep_prob = keep_prob, padding = "SYMMETRIC", scope = 'm_cls_4', is_train = m_cls_bn, bn_trainable = m_cls_trainable, leak = True)
self.m_cls_weights.append( wr4_1m )
with tf.variable_scope('m_cls_out') as scope:
wm_out = sharable_weight_variable( shape = [ feature_base* 16 * 4,1 ], trainable = m_cls_trainable , name = "Variable" )
out5m_flat = tf.reshape(conv_4m, [-1, feature_base * 16 * 4])
m_cls_logits = tf.matmul(out5m_flat, wm_out)
self.m_cls_weights.append(wm_out)
return m_cls_logits
def _get_cost(self, ct_logits, mr_logits, ct_cls_logits, mr_cls_logits, ct_mask_logits, mr_mask_logits, cost_kwargs):
miu_dis = cost_kwargs["miu_dis"] # coefficient for discriminator loss
miu_gen = cost_kwargs["miu_gen"] # used to be 0.5 0.5 1
lambda_mask_loss = cost_kwargs.pop("lambda_mask_loss", 1.0) # weighting of mask critic score
self.miu_dis = tf.Variable(miu_dis, name = "miu_dis") # coefficient for discrminator
self.miu_gen = tf.Variable(miu_gen, name = "miu_gen")
# loss for main critic and mask critic
dis_loss = -1 * self.miu_dis * tf.reduce_mean( mr_cls_logits - ct_cls_logits ) # loss functions of WGAN
gen_loss = -1 * self.miu_gen * tf.reduce_mean( ct_cls_logits )
m_dis_loss = -1 * self.miu_dis * tf.reduce_mean( mr_mask_logits - ct_mask_logits )
m_gen_loss = -1 * self.miu_gen * tf.reduce_mean( ct_mask_logits )
############ L2 norm regularizer ######################
reg_coeff = cost_kwargs.pop("regularizer", 1.0e-4) # regularizer coefficients for non-GAN parts
mr_front_reg = sum([tf.nn.l2_loss(variable) for variable in self.mr_front_weights]) # regulizer for MRI varibles, fixed for the unsupervised setting
joint_reg = sum([tf.nn.l2_loss(variable) for variable in self.joint_weights]) # regularizer for joint part, fixed for the unsupervised setting
fixed_coeff_reg = reg_coeff * (mr_front_reg + joint_reg) # for training observation to confirm the source segmenter is not updated
gan_reg_coeff = cost_kwargs.pop("gan_regularizer", 1.0e-4) # regularizer coefficients for GAN parts, note, seems that it works well when it is larger
gen_reg = gan_reg_coeff * self.miu_gen * sum([tf.nn.l2_loss(variable) for variable in self.ct_front_weights]) # regulizers for WGAN
dis_reg = gan_reg_coeff * self.miu_dis * sum([tf.nn.l2_loss(variable) for variable in self.cls_weights])
m_dis_reg = gan_reg_coeff * self.miu_dis * sum([tf.nn.l2_loss(variable) for variable in self.m_cls_weights])
dis_loss += lambda_mask_loss * m_dis_loss
gen_loss += lambda_mask_loss * m_gen_loss
dis_reg += lambda_mask_loss * m_dis_reg
return dis_loss, gen_loss, fixed_coeff_reg, dis_reg, gen_reg
def _get_variables_by_scope(self):
"""
Group different variables (MR, CT, GAN, etc)to different groups
"""
logging.info("extent of joint part and segmenter need to be manually set, including variables and bns")
self.adapt_vars = [] # variables for adaptation (CT)
self.cls_vars = [] # variables for domain-classifier (i.e. discriminator) for WGAN
self.seg_vars = [] # variables for segmentation, fixed higher layers in source segmenter
self.mri_seg_vars = [] # variables for segmentation, MRI early players, fixed as well
var_list = tf.contrib.framework.get_variables()
for var in var_list:
if "cls" in var.name:
self.cls_vars.append(var)
elif "adapt" in var.name:
self.adapt_vars.append(var)
elif "output" in var.name:
self.seg_vars.append(var)
self.mri_seg_vars.append(var)
elif "group" in var.name:
_group_name = var.name.split("/")[0]
_group_no = float(_group_name.split("_")[-1] )
self.mri_seg_vars.append(var)
def restore(self, sess, model_path, no_gan=False, clear_rms=False):
"""
Restores a session from a checkpoint
:param sess: current session instance
:param model_path: path to file system checkpoint location
:param no_gan: only restore mr variables
:param clear_rms: does not restore RMSprop internal variables, please set is true
"""
saver = tf.train.Saver(tf.contrib.framework.get_variables() + tf.get_collection_ref("internal_batchnorm_variables") )
logging.info("Model restored from file: %s" % model_path)
if no_gan is True:
logging.info("I only load the main variables! without batchnorm!!!")
variables = tf.global_variables()
reader = tf.pywrap_tensorflow.NewCheckpointReader(model_path)
var_keep_dic = reader.get_variable_to_shape_map()
variables_to_restore = []
for v in variables:
if v.name.split(':')[0] in var_keep_dic:
if ("adapt" in v.name) or ("cls" in v.name) or("Adam" in v.name):
continue
if ("group" in v.name) or ("output" in v.name):
logging.info("restoring "+str(v.name))
variables_to_restore.append(v)
restorer = tf.train.Saver(variables_to_restore)
restorer.restore(sess, model_path)
logging.info("Model restored from file: %s, the pre-trained MRI model (without bn params)" % model_path)
return 0
if clear_rms is True:
logging.info("Calculating RMS parameters from beginning")
variables = tf.global_variables()
reader = tf.pywrap_tensorflow.NewCheckpointReader(model_path)
var_keep_dic = reader.get_variable_to_shape_map()
variables_to_restore = []
for v in variables:
if v.name.split(':')[0] in var_keep_dic:
if ("RMS" in v.name) :
continue
else:
logging.info("restoring "+str(v.name))
variables_to_restore.append(v)
restorer = tf.train.Saver(variables_to_restore)
restorer.restore(sess, model_path)
logging.info("Model restored from file: %s and RMS variables are ignored" % model_path)
return 0
try: # else, just restore as much as possible
saver.restore(sess, model_path)
logging.info("Model restored from file: %s" % model_path)
except:
variables = tf.global_variables()
reader = tf.pywrap_tensorflow.NewCheckpointReader(model_path)
var_keep_dic = reader.get_variable_to_shape_map()
variables_to_restore = []
for v in variables:
if v.name.split(':')[0] in var_keep_dic:
skip_flg = False
for kwd in self.network_config["restore_skip_kwd"]: # if it is manully specified to be skipped, don't restore it
if kwd in v.name:
skip_flg = True
break
if skip_flg is False:
variables_to_restore.append(v)
logging.info("cannot fully restore the model, restoring "+str(v.name))
restorer = tf.train.Saver(variables_to_restore)
restorer.restore(sess, model_path)
logging.info("Model restored from file: %s with relaxation" % model_path)
class Trainer(object):
"""
Train a unet instance
"""
def __init__(self, net, mr_train_list, mr_val_list, ct_train_list, ct_val_list,\
adapt_var_list, mr_var_list, old_bn_list, new_bn_list,\
test_label_list = None, test_nii_list = None,\
num_cls=None, batch_size = 6,\
opt_kwargs={}, train_config = {}):
self.net = net
self.batch_size = batch_size
self.num_cls = num_cls # including background
self.opt_kwargs = opt_kwargs
self.ct_train_list = ct_train_list # a list of training files
self.ct_val_list = ct_val_list # a list of validation files
self.mr_train_list = mr_train_list # a list of training files for MRI
self.mr_val_list = mr_val_list
self.test_label_list = test_label_list # test files (npz format)
self.test_nii_list = test_nii_list # test files (npz format)
self.adapt_var_list = adapt_var_list # a list of variables in CT path
self.mr_var_list = mr_var_list # a list of variables in MRI path in correspondance with variables in adapt_var_list, this is used for manually initialize variables in CT path with those of MRI path
self.old_bn_list = old_bn_list # a list of batch_norm internal variables in baseline model
self.new_bn_list = new_bn_list # a list of batch_norm internal variables for the MRI path in current model
self.ct_train_queue = tf.train.string_input_producer(ct_train_list, num_epochs = None, shuffle = True) # tensorflow input queue for CT supervision (disabled), CT and MRI
self.ct_val_queue = tf.train.string_input_producer(ct_val_list, num_epochs = None, shuffle = True)
self.mr_train_queue = tf.train.string_input_producer(mr_train_list, num_epochs = None, shuffle = True)
self.mr_val_queue = tf.train.string_input_producer(mr_val_list, num_epochs = None, shuffle = True)
self.train_config = train_config # configuations for training
self.lr_update_flag = train_config["lr_update"]
def next_batch(self, input_queue, capacity = 120, num_threads = 2, min_after_dequeue = 30, label_type = 'float'):
reader = tf.TFRecordReader()
fid, serialized_example = reader.read(input_queue)
parser = tf.parse_single_example(serialized_example, features = decomp_feature)
dsize_dim0 = tf.cast(parser['dsize_dim0'], tf.int32)
dsize_dim1 = tf.cast(parser['dsize_dim1'], tf.int32)
dsize_dim2 = tf.cast(parser['dsize_dim2'], tf.int32)
lsize_dim0 = tf.cast(parser['lsize_dim0'], tf.int32)
lsize_dim1 = tf.cast(parser['lsize_dim1'], tf.int32)
lsize_dim2 = tf.cast(parser['dsize_dim2'], tf.int32)
data_vol = tf.decode_raw(parser['data_vol'], tf.float32)
label_vol = tf.decode_raw(parser['label_vol'], tf.float32)
data_vol = tf.reshape(data_vol, raw_size)
label_vol = tf.reshape(label_vol, raw_size)
data_vol = tf.slice(data_vol, [0,0,0],volume_size)
label_vol = tf.slice(label_vol, [0,0,1], label_size)
data_feed, label_feed, fid_feed = tf.train.shuffle_batch([data_vol, label_vol, fid], batch_size =self.batch_size , capacity = capacity, \
num_threads = num_threads, min_after_dequeue = min_after_dequeue)
pair_feed = tf.concat([data_feed, label_feed], axis = 3) # concatenate them
return pair_feed, fid_feed
def _get_optimizer(self, training_iters, global_step):
"""
Use RMSprop instead of Adam for training WGAN
"""
learning_rate = self.opt_kwargs.pop("learning_rate", None) # default set to 0.0002
self.LR_refresh = learning_rate
self.learning_rate_node = tf.Variable(learning_rate)
# optimizer for discriminator/ domain classifier
dis_optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate_node,
**self.opt_kwargs).minimize(self.net.dis_loss + 1.0 / self.train_config['dis_sub_iter'] * self.net.dis_reg,
global_step=global_step,\
var_list = self.net.cls_vars)
# optimizer for training generator
gen_optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate_node,
**self.opt_kwargs).minimize(self.net.ct_gen_loss + 1.0 / self.train_config['gen_sub_iter'] * self.net.gen_reg,
global_step=global_step,\
var_list = self.net.adapt_vars)
# clip operation for WGAN for Lipschitz constrain
self.clip_op = [tf.assign(var, tf.clip_by_value(var, -0.03, 0.03)) for var in self.net.cls_vars if "Variable" in var.name]
return dis_optimizer, gen_optimizer
def _initialize(self, training_iters, output_path):
"""
initialization and tensorboard setting
"""
self.global_step = tf.Variable(0)
scalar_summaries = [] # tensorboard summaries
scalar_summaries.append(tf.summary.scalar('fixed_coeff_reg', self.net.fixed_coeff_reg)) # regulizer of MRI segemter weights, monitor MRI weights unchanged
scalar_summaries.append(tf.summary.scalar('discriminator_loss', self.net.dis_loss))
scalar_summaries.append(tf.summary.scalar('generator_loss', self.net.ct_gen_loss))
scalar_summaries.append(tf.summary.scalar('ct_dice_eval_c1_lv_myo', self.net.ct_dice_eval_c1))
scalar_summaries.append(tf.summary.scalar('ct_dice_eval_c2_la_blood', self.net.ct_dice_eval_c2))
scalar_summaries.append(tf.summary.scalar('ct_dice_eval_c3_lv_blood', self.net.ct_dice_eval_c3))
scalar_summaries.append(tf.summary.scalar('ct_dice_eval_c4_aa', self.net.ct_dice_eval_c4))
scalar_summaries.append(tf.summary.scalar('mri_dice', self.net.mr_dice_eval)) # set to show absolute value for mr segmentation
train_images = []
train_images.append(tf.summary.image('ct_pred', tf.expand_dims(tf.cast(self.net.compact_pred, tf.float32), 3 )) ) # ct prediction
train_images.append(tf.summary.image('ct_image', tf.expand_dims(tf.cast(self.net.ct[:,:,:,1], tf.float32), 3 )) )
train_images.append(tf.summary.image('ct_gt', tf.expand_dims(tf.cast(self.net.compact_y, tf.float32), 3))) # ground truth for CT segmentation
train_images.append(tf.summary.image('mri_validation_pred', tf.expand_dims(tf.cast(self.net.compact_mr_valid, tf.float32), 3 )) ) # mri segmentation for debugging
train_images.append(tf.summary.image('mri_image', tf.expand_dims(tf.cast(self.net.mr[:,:,:,1], tf.float32), 3 )) )
train_images.append(tf.summary.image('mri_gt', tf.expand_dims(tf.cast(self.net.compact_mr_y, tf.float32), 3))) # ground truth for CT segmentation
val_images = []
val_images.append(tf.summary.image('ct_val_pred', tf.expand_dims(tf.cast(self.net.compact_pred, tf.float32), 3))) # prediction for validation
val_images.append(tf.summary.image('ct_image', tf.expand_dims(tf.cast(self.net.ct[:,:,:,1], tf.float32), 3)))
val_images.append(tf.summary.image('ct_val_gt', tf.expand_dims(tf.cast(self.net.compact_y, tf.float32), 3)))
self.net._get_variables_by_scope() # get variable groups
self.dis_optimizer, self.gen_optimizer = self._get_optimizer(training_iters, self.global_step) # get optimizers
scalar_summaries.append(tf.summary.scalar('learning_rate', self.learning_rate_node))
# get summary writers
self.scalar_summary_op = tf.summary.merge(scalar_summaries)
self.train_image_summary_op = tf.summary.merge(train_images)
self.val_image_summary_op = tf.summary.merge(val_images)
# variable initializers
init_glb = tf.global_variables_initializer()
init_loc = tf.variables_initializer(tf.local_variables())
return init_glb, init_loc
def _adapt_copy_weights(self, internal = False):
if internal is False:
if len(self.mr_var_list) != len(self.adapt_var_list):
raise ValueError("cannot copy weight to adaptation because of incorrect varaible lists")
with tf.variable_scope("", reuse = True):
for idx in range(len(self.mr_var_list)):
logging.info("Now initializing adaptation variable %s with mainstream variable %s"%( self.adapt_var_list[idx], self.mr_var_list[idx] ))
_curr_mr_var = tf.get_default_graph().get_tensor_by_name(self.mr_var_list[idx])
_curr_adapt_var = tf.get_default_graph().get_tensor_by_name(self.adapt_var_list[idx])
upd_op = tf.assign(_curr_adapt_var,_curr_mr_var)
upd_op.eval()
else:
logging.info("automatically seeks for variable correspondance")
all_var_list = tf.contrib.framework.get_variables()
self.mr_var_list = []
self.adapt_var_list = []
for v in all_var_list:
if ("RMS" in v.name) or ("Adam" in v.name):
continue
else:
if "group" in v.name:
self.mr_var_list.append(v)
elif "adapt" in v.name:
self.adapt_var_list.append(v)
else:
continue
if len(self.mr_var_list) != len(self.adapt_var_list):
raise ValueError("cannot copy weight to adaptation because of incorrect varaible list")
for _curr_adapt_var, _curr_mr_var in zip(self.adapt_var_list, self.mr_var_list):
upd_op = tf.assign(_curr_adapt_var,_curr_mr_var)
upd_op.eval()
logging.info("adaptation module has been initialized! Please remember that it is a one-time operation")
def _load_batch_norm_weights(self, output_path):
"""
convenience function for loading weights from eariler version of baseline model for the CT/MR segmentation network
old_bn_list: a list of bn variable names in baseline model
new_bn_list: a list of bn Variable names in current model
"""
if len(self.old_bn_list) != len(self.new_bn_list):
raise ValueError("two mappings mismatch")
checkpoint = tf.train.get_checkpoint_state(output_path)
self.copy_bn_dict = {}
for old_var, new_var in zip(self.old_bn_list, self.new_bn_list):
n_group = new_var.split("_")[1]
new_var = "group_" + n_group + "/" + new_var
self.copy_bn_dict[old_var] = new_var
old_variable = tf.contrib.framework.load_variable( output_path, old_var )
new_variable = tf.get_default_graph().get_tensor_by_name(new_var)
upd_op = tf.assign(new_variable, old_variable)
upd_op.eval()
logging.info("%s has send value to %s"%(old_var, new_var))
return 0
def train(self, output_path, restore=True, restored_path=None, training_iters=200, epochs=1000, dropout=0.75, display_step=5):
self.output_path = output_path
if not os.path.exists(output_path):
logging.info("Allocating '{:}'".format(output_path))
os.makedirs(output_path)
self._initialize_logs()
save_path = os.path.join(output_path, "model.cpkt")
if epochs == 0:
return save_path
init_glb, init_loc = self._initialize(training_iters, output_path)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True # False
with tf.Session(config=config) as sess:
sess.run([ init_glb, init_loc] )
coord = tf.train.Coordinator()
# For restore models, there are three situations:
# 1. warming up discriminator, init from MRI segmenter: "restore_from_baseline=True, clear_rms=True"
# if restore_from_baseline set True, clear_rms whatever,
# this would restore the pre-trained MRI segmenter (without BN), this works together with following lines 1076-1079 to manually load BN
# 2. after warming up discriminator, start training GAN: "restpre_from_baseline=False, clear_rms=True"
# this would restore the entire GAN system with warmed up discriminator (excluding RMS from optimizer)
# 3. fine-tune GAN from a breakpoint: "restore_from_baseline=False, clear_rms=False"
if restore:
if restored_path is None:
raise Exception("No restore path is provided")
ckpt = tf.train.get_checkpoint_state(restored_path)
if ckpt and ckpt.model_checkpoint_path:
self.net.restore(sess, ckpt.model_checkpoint_path, no_gan = self.train_config["restore_from_baseline"], clear_rms = self.train_config["clear_rms"])
if self.train_config["restore_from_baseline"] is True: # here initialize the MRI and CT part with the pre-trained MRI segmenter, only call once beginning train
self._load_batch_norm_weights(restored_path) # load batchnorm variables of MRI-specific and joint part
print("initializing from baseline model!")
self._adapt_copy_weights() # copy MRI weights to CT adaptation layers for initialization
if self.lr_update_flag is True: # manually reset learning rate when needed
sess.run( tf.assign(self.learning_rate_node, self.LR_refresh) )
logging.info("New learning rate %s has been loaded"%str(self.LR_refresh))
train_summary_writer = tf.summary.FileWriter(output_path + "/train_log" + self.train_config['tag'], graph=sess.graph)
val_summary_writer = tf.summary.FileWriter(output_path + "/val_log" + self.train_config['tag'], graph=sess.graph)
ct_feed_all, ct_feed_fid = self.next_batch(self.ct_train_queue)
ct_feed_val, ct_feed_val_fid = self.next_batch(self.ct_val_queue)
mr_feed_all, mr_feed_fid = self.next_batch(self.mr_train_queue)
mr_feed_val, mr_feed_val_fid = self.next_batch(self.mr_val_queue)
threads = tf.train.start_queue_runners(sess = sess, coord = coord, start = True)
# read iteration configurations
dis_interval = self.train_config['dis_interval'] # frequency of discriminator updates, default 1. if set 2, update discriminator every 2 iterations
gen_interval = self.train_config['gen_interval'] # frequency of generator updates, default 1. if set 2, update generator every 2 iterations
dis_sub_iter = self.train_config['dis_sub_iter'] # number of sub-iteration in one updates, recommended to be larger than gen_sub_iter
gen_sub_iter = self.train_config['gen_sub_iter']
# set if we what to increase *_sub_iter every <sub_iter_upd_interval>.
# for example, if this is set 1, and sub_iter_upd_interval is 100, then increase dis_sub_iter by 1 every 100 iterations
dis_sub_iter_inc = self.train_config.pop('dis_sub_iter_inc', 0)
gen_sub_iter_inc = self.train_config.pop('gen_sub_iter_inc', 0)
sub_iter_upd_interval = self.train_config.pop('iter_upd_interval', 999999999999)
for epoch in range(epochs):
for step in range((epoch*training_iters), ((epoch+1)*training_iters)):
logging.info("Running step %s epoch %s ..."%(str(step), str(epoch)))
start = time.time()
# according to DCGAN paper, first update discriminator
if dis_interval == 0:
pass
elif (step % dis_interval == 0) and (step != 0):
for itr_dummy in range(dis_sub_iter):
# read samples from the pipeline, decomp them and feed them into the discriminator
ct_batch, ct_fid = sess.run([ct_feed_all, ct_feed_fid])
ct_raw_y = ct_batch[:,:,:,3]
ct_batch = ct_batch[:,:,:,0:3]
ct_batch_y = _label_decomp(self.num_cls, ct_raw_y)
mr_batch, mr_fid = sess.run([mr_feed_all, mr_feed_fid])
mr_raw_y = mr_batch[:,:,:,3]
mr_batch = mr_batch[:,:,:,0:3]
mr_batch_y = _label_decomp(self.num_cls, mr_raw_y)
_, _ = sess.run((self.dis_optimizer, self.learning_rate_node),
feed_dict={ self.net.mr: mr_batch,
self.net.ct: ct_batch,
self.net.mr_front_bn: False,
self.net.joint_bn: False,
self.net.ct_front_bn: False,
self.net.cls_bn: True,
self.net.keep_prob: dropout})
# clip operation
sess.run(self.clip_op)
logging.info("discriminator updated %s of %s"%(str(itr_dummy),str(dis_sub_iter)))
# update generator
if gen_interval == 0:
pass
elif (step % gen_interval == 0) and (step != 0):
for _ in range(gen_sub_iter):
ct_batch, ct_fid = sess.run([ct_feed_all, ct_feed_fid])
ct_raw_y = ct_batch[:,:,:,3]
ct_batch = ct_batch[:,:,:,0:3]
ct_batch_y = _label_decomp(self.num_cls, ct_raw_y)
_, _ = sess.run((self.gen_optimizer, self.learning_rate_node),
feed_dict={ self.net.ct: ct_batch,
self.net.mr_front_bn: False,
self.net.joint_bn: False,
self.net.ct_front_bn: True,
self.net.cls_bn: False,
self.net.keep_prob: dropout})
logging.info("generator updated")
# if we need to update iteration configurations, do it here
if (step % sub_iter_upd_interval == 0) and (step != 0):
dis_sub_iter += dis_sub_iter_inc
gen_sub_iter += gen_sub_iter_inc
logging.info("sub iterations updated!")
logging.info("Training step %s epoch %s has been finished!"%(str(step), str(epoch)))
logging.info("Time elapsed %s seconds"%(str(time.time() - start)))
# evaluation and write them to tensorboard
if step % display_step == 0:
# training batch
train_ct_batch = sess.run(ct_feed_all)
train_ct_raw_y = train_ct_batch[:,:,:,3]
train_ct_batch = train_ct_batch[:,:,:,0:3]
train_ct_batch_y = _label_decomp(self.num_cls, train_ct_raw_y)
mr_batch, mr_fid = sess.run([mr_feed_all, mr_feed_fid])
mr_raw_y = mr_batch[:,:,:,3]
mr_batch = mr_batch[:,:,:,0:3]
mr_batch_y = _label_decomp(self.num_cls, mr_raw_y)
self.output_minibatch_stats(sess, train_summary_writer, step, train_ct_batch, train_ct_batch_y, mr_batch, mr_batch_y)
if step % (display_step * 1) == 0:
# validation batch
ct_batch = sess.run(ct_feed_val)
ct_raw_y = ct_batch[:,:,:,3]
ct_batch = ct_batch[:,:,:,0:3]
ct_batch_y = _label_decomp(self.num_cls, ct_raw_y)
mr_batch = sess.run(mr_feed_val)
mr_raw_y = mr_batch[:,:,:,3]
mr_batch = mr_batch[:,:,:,0:3]
mr_batch_y = _label_decomp(self.num_cls, mr_raw_y)
self.output_minibatch_stats(sess, val_summary_writer, step, ct_batch, ct_batch_y, mr_batch, mr_batch_y, detail = True)
# save and restore the model periodically
if step % (self.train_config["checkpoint_space"]) == 0:
if step == 0:
continue
else:
save_path = _save(sess, save_path, global_step = self.global_step.eval())
print('*********************** save path ******************: ', save_path)
logging.info("Model has been saved ...")
last_ckpt = tf.train.get_checkpoint_state(output_path)
if last_ckpt and last_ckpt.model_checkpoint_path:
self.net.restore(sess, last_ckpt.model_checkpoint_path)
logging.info("Model has been restored for re-allocation")
# learning rate decay
_pre_lr = sess.run(self.learning_rate_node)
sess.run( tf.assign(self.learning_rate_node, _pre_lr *\
self.train_config['lr_decay_factor']) )
logging.info("Global step %s"%str(self.global_step.eval()))
logging.info("Optimization Finished!")
coord.request_stop()
coord.join(threads)
return save_path
def output_minibatch_stats(self, sess, summary_writer, step, ct_batch, ct_batch_y, mr_batch, mr_batch_y, detail = False):
"""
minibatch stats for tensorboard observation
"""
if detail is not True:
summary_str, summary_img = sess.run([\
self.scalar_summary_op,
self.train_image_summary_op],
feed_dict={\
self.net.ct_front_bn : False,
self.net.mr_front_bn : False,
self.net.joint_bn : False,
self.net.cls_bn : False,
self.net.mr: mr_batch,
self.net.mr_y: mr_batch_y,
self.net.ct: ct_batch,
self.net.ct_y: ct_batch_y,
self.net.keep_prob: 1.\
})
else:
_, curr_conf_mat, summary_str, summary_img = sess.run([\
self.net.compact_pred,
self.net.confusion_matrix,
self.scalar_summary_op,
self.train_image_summary_op],
feed_dict={\
self.net.ct_front_bn : False,
self.net.mr_front_bn : False,
self.net.joint_bn : False,
self.net.cls_bn : False,
self.net.mr: mr_batch,
self.net.mr_y: mr_batch_y,
self.net.ct: ct_batch,
self.net.ct_y: ct_batch_y,
self.net.keep_prob: 1.\
})
_indicator_eval(curr_conf_mat)
summary_writer.add_summary(summary_str, step)
summary_writer.add_summary(summary_img, step)
summary_writer.flush()
def test_eval(self, sess, output_path, flip_correction = True):
all_cm = np.zeros([self.num_cls, self.num_cls])
pred_folder = os.path.join(output_path, "dense_pred")
try:
os.makedirs(pred_folder)
except:
logging.info("prediction folder exists")
self.test_pair_list = list(zip(self.test_label_list, self.test_nii_list))
sample_eval_list = [] # evaluation of each sample
for idx_file, pair in enumerate(self.test_pair_list):
sample_cm = np.zeros([self.num_cls, self.num_cls]) # confusion matrix for each sample
label_fid = pair[0]
nii_fid = pair[1]
if not os.path.isfile(nii_fid):
raise Exception("cannot find sample %s"%str(nii_fid))
raw = read_nii_image(nii_fid)
raw_y = read_nii_image(label_fid)
if flip_correction is True:
raw = np.flip(raw, axis = 0)
raw = np.flip(raw, axis = 1)
raw_y = np.flip(raw_y, axis = 0)
raw_y = np.flip(raw_y, axis = 1)
tmp_y = np.zeros(raw_y.shape)
frame_list = [kk for kk in range(1, raw.shape[2] - 1)]
np.random.shuffle(frame_list)
for ii in range( int( floor( raw.shape[2] // self.net.batch_size ) ) ):
vol = np.zeros( [self.net.batch_size, raw_size[0], raw_size[1], raw_size[2]] )
slice_y = np.zeros( [self.net.batch_size, label_size[0], label_size[1]] )
for idx, jj in enumerate(frame_list[ ii * self.net.batch_size : (ii + 1) * self.net.batch_size ]):
vol[idx, ...] = raw[ ..., jj -1: jj+2 ].copy()
slice_y[idx,...] = raw_y[..., jj ].copy()
vol_y = _label_decomp(self.num_cls, slice_y)
pred, curr_conf_mat= sess.run([self.net.compact_pred, self.net.confusion_matrix], feed_dict =\
{self.net.ct: vol, self.net.ct_y: vol_y, self.net.keep_prob: 1.0, self.net.mr_front_bn : False,\
self.net.ct_front_bn: False})
for idx, jj in enumerate(frame_list[ii * self.net.batch_size: (ii + 1) * self.net.batch_size]):
tmp_y[..., jj] = pred[idx, ...].copy()
sample_cm += curr_conf_mat
all_cm += sample_cm
sample_dice = _dice(sample_cm)
sample_jaccard = _jaccard(sample_cm)
sample_eval_list.append((sample_dice, sample_jaccard))
subject_dice_list, subject_jaccard_list = self.sample_metric_stddev(sample_eval_list)
np.savetxt(os.path.join(output_path, "cm.csv"), all_cm)
return subject_dice_list, subject_jaccard_list
def sample_metric_stddev(self, sample_eval_list):
"""
calculate stddev of each organ across samples
"""
metric_mat = np.zeros( [len(sample_eval_list), self.num_cls, 2] )
for organ, ind in list(contour_map.items()):
for ii in range(len(sample_eval_list)):
metric_mat[ii, int(ind), 0] = sample_eval_list[ii][0][int(ind)] # dice
metric_mat[ii, int(ind), 1] = sample_eval_list[ii][1][int(ind)] # jaccard
print("------- inside the sample_metric_stddev file ---- ")
for organ, ind in list(contour_map.items()):
print(( "organ: %s"%organ ))
print(( "dice_stddev: %s"%( np.std(metric_mat[:, int(ind), 0] ) ) ))
print(( "jaccard_stddev: %s"%( np.std(metric_mat[:, int(ind), 1] ) ) ))
print("------- inside the sample_metric_stddev file ---- ")
for organ, ind in list(contour_map.items()):
print(( "organ: %s"%organ ))
print(( "dice_mean: %s"%( np.mean(metric_mat[:, int(ind), 0] ) ) ))
print(( "jaccard_mean %s"%( np.mean(metric_mat[:, int(ind), 1] ) ) ))
print("-------")
print(( "all_dice_mean: %s"%( np.mean(metric_mat[:, 1:, 0] ) ) ))
print(("all_jaccard_mean: %s" % (np.mean(metric_mat[:, 1:, 1] ) )))
subject_level_list = np.mean(metric_mat, axis=0)
subject_level_list_dice = subject_level_list[:,0]
subject_level_list_jaccard = subject_level_list[:1]
return subject_level_list_dice, subject_level_list_jaccard
def _initialize_logs(self):
"""
This log is actually useless so ignore it
"""
self.acc_dice_dict = {}
self.acc_jaccard_dict = {}
self.log_eval_fid = os.path.join(self.output_path, "acc_eval.csv")
for organ, ind in list(contour_map.items()):
self.acc_dice_dict[organ] = [organ]
self.acc_jaccard_dict[organ] = [organ]
def test_model(self, this_model, output_path):
init_glb, init_loc = self._initialize(1, output_path)
with tf.Session() as sess:
sess.run([init_glb, init_loc])
self.net.restore(sess, this_model)
logging.info("model has been loaded!")
dice, jac = self.test_eval(sess, output_path)
logging.info("testing finished")
return dice, jac
