else:
os.environ["CUDA_VISIBLE_DEVICES"] = ''
config=tf.ConfigProto()
config.gpu_options.allow_growth=True
sess=tf.Session(config=config)
input_i = tf.placeholder(tf.float32, shape=[None,None,None,3])
input_target = tf.placeholder(tf.float32, shape=[None,None,None,3])
lossDict = {}
objDict={}
with tf.variable_scope(tf.get_variable_scope()):
frame_input = input_i[:, :, :, 0:3]
frame_processed = input_target[:, :, :, 0:3]
with tf.variable_scope('individual'):
if not with_IRT:
frame_out = net.VCN(frame_input, reuse=False)
lossDict["total"] = Lp_loss(frame_out, frame_processed)
else:
net_pred=net.VCN(frame_input, output_channel=6, reuse=False)
frame_out = net_pred[...,0:3]
frame_out_minor = net_pred[...,3:6]
diff_map = tf.reduce_max(tf.abs(frame_out - frame_processed) / (frame_input + 1e-1), axis=3, keep_dims=True)
diff_map1 = tf.reduce_max(tf.abs(frame_out_minor - frame_processed) / (frame_input + 1e-1), axis=3, keep_dims=True)
confidence_map = tf.tile(tf.cast(tf.less(diff_map, diff_map1), tf.float32), [1,1,1,3])
lossDict["total"] = tf.reduce_mean(tf.abs(frame_out - frame_processed)) + tf.reduce_mean(tf.abs(frame_out_minor - frame_processed))
lossDict["confidence_map"] = tf.reduce_mean(tf.abs(frame_out - frame_processed)*confidence_map) + tf.reduce_mean(
tf.abs(frame_out_minor - frame_processed) * (-confidence_map + 1.0))
opt=tf.train.AdamOptimizer(learning_rate=0.0001).minimize(lossDict["total"], var_list=[var for var in tf.trainable_variables()])
opt_IRT=tf.train.AdamOptimizer(learning_rate=0.0001).minimize(lossDict["confidence_map"], var_list=[var for var in tf.trainable_variables()])
print([var for var in tf.trainable_variables() if var.name.startswith('VCRN')])
def get_model(input):
if ARGS.network == "unet":
return net.VCN(input, reuse=False)