def main(_): if FLAGS.batch_size is not 1: print("Error: to test images, batch size must be 1") exit() model_dir = os.path.join(FLAGS.log_root, FLAGS.train_dir) if not os.path.exists(FLAGS.log_root): os.makedirs(FLAGS.log_root) if not os.path.exists(model_dir): os.makedirs(model_dir) bart_dir = os.path.join(model_dir, "bart_recon") if not os.path.exists(bart_dir): os.makedirs(bart_dir) run_config = tf.ConfigProto() run_config.gpu_options.allow_growth = True with tf.Session(config=run_config) as sess: """Execute main function.""" os.environ["CUDA_VISIBLE_DEVICES"] = FLAGS.device if not FLAGS.dataset_dir: raise ValueError("You must supply the dataset directory with " + "--dataset_dir") if FLAGS.random_seed >= 0: random.seed(FLAGS.random_seed) np.random.seed(FLAGS.random_seed) tf.logging.set_verbosity(tf.logging.INFO) print("Preparing dataset...") out_shape = [FLAGS.shape_z, FLAGS.shape_y] test_dataset, num_files = mri_data.create_dataset( os.path.join(FLAGS.dataset_dir, "test"), FLAGS.mask_path, num_channels=FLAGS.num_channels, num_emaps=FLAGS.num_emaps, batch_size=FLAGS.batch_size, out_shape=out_shape, ) # channels first: (batch, channels, z, y) # placeholders ks_shape = [None, FLAGS.shape_z, FLAGS.shape_y, FLAGS.num_channels] ks_place = tf.placeholder(tf.complex64, ks_shape) sense_shape = [ None, FLAGS.shape_z, FLAGS.shape_y, 1, FLAGS.num_channels ] sense_place = tf.placeholder(tf.complex64, sense_shape) im_shape = [None, FLAGS.shape_z, FLAGS.shape_y, 1] im_truth_place = tf.placeholder(tf.complex64, im_shape) # run through unrolled im_out_place = mri_model.unroll_fista( ks_place, sense_place, is_training=True, verbose=True, do_hardproj=FLAGS.do_hard_proj, num_summary_image=FLAGS.num_summary_image, resblock_num_features=FLAGS.feat_map, num_grad_steps=FLAGS.num_grad_steps, conv=FLAGS.conv, do_conjugate=FLAGS.do_conjugate, ) saver = tf.train.Saver() summary_writer = tf.summary.FileWriter(model_dir, sess.graph) # initialize model print("[*] initializing network...") if not load(model_dir, saver, sess): sess.run(tf.global_variables_initializer()) coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(sess, coord) # See how many parameters are in model total_parameters = 0 for variable in tf.trainable_variables(): variable_parameters = 1 for dim in variable.get_shape(): variable_parameters *= dim.value total_parameters += variable_parameters print("Total number of trainable parameters: %d" % total_parameters) test_iterator = test_dataset.make_one_shot_iterator() features, labels = test_iterator.get_next() ks_truth = labels ks_in = features["ks_input"] sense_in = features["sensemap"] mask_recon = features["mask_recon"] im_truth = tf_util.model_transpose(ks_truth * mask_recon, sense_in) total_summary = tf.summary.merge_all() output_psnr = [] output_nrmse = [] output_ssim = [] cs_psnr = [] cs_nrmse = [] cs_ssim = [] for test_file in range(num_files): ks_in_run, sense_in_run, im_truth_run = sess.run( [ks_in, sense_in, im_truth]) im_out, total_summary_run = sess.run( [im_out_place, total_summary], feed_dict={ ks_place: ks_in_run, sense_place: sense_in_run, im_truth_place: im_truth_run, }, ) # CS recon bart_test = bart_cs(bart_dir, ks_in_run, sense_in_run, l1=0.007) # bart_test = None # handle batch dimension for b in range(FLAGS.batch_size): truth = im_truth_run[b, :, :, :] out = im_out[b, :, :, :] psnr, nrmse, ssim = metrics.compute_all(truth, out, sos_axis=-1) output_psnr.append(psnr) output_nrmse.append(nrmse) output_ssim.append(ssim) print("output mean +/ standard deviation psnr, nrmse, ssim") print( np.mean(output_psnr), np.std(output_psnr), np.mean(output_nrmse), np.std(output_nrmse), np.mean(output_ssim), np.std(output_ssim), ) psnr, nrmse, ssim = metrics.compute_all(im_truth_run, bart_test, sos_axis=-1) cs_psnr.append(psnr) cs_nrmse.append(nrmse) cs_ssim.append(ssim) print("cs mean +/ standard deviation psnr, nrmse, ssim") print( np.mean(cs_psnr), np.std(cs_psnr), np.mean(cs_nrmse), np.std(cs_nrmse), np.mean(cs_ssim), np.std(cs_ssim), ) print("End of testing loop") txt_path = os.path.join(model_dir, "metrics.txt") f = open(txt_path, "w") f.write("parameters = " + str(total_parameters) + "\n" + "output psnr = " + str(np.mean(output_psnr)) + " +\- " + str(np.std(output_psnr)) + "\n" + "output nrmse = " + str(np.mean(output_nrmse)) + " +\- " + str(np.std(output_nrmse)) + "\n" + "output ssim = " + str(np.mean(output_ssim)) + " +\- " + str(np.std(output_ssim)) + "\n" "cs psnr = " + str(np.mean(cs_psnr)) + " +\- " + str(np.std(cs_psnr)) + "\n" + "output nrmse = " + str(np.mean(cs_nrmse)) + " +\- " + str(np.std(cs_nrmse)) + "\n" + "output ssim = " + str(np.mean(cs_ssim)) + " +\- " + str(np.std(cs_ssim))) f.close()
def main(_): # path where model checkpoints and summaries will be saved model_dir = os.path.join(FLAGS.log_root, FLAGS.train_dir) if not os.path.exists(FLAGS.log_root): os.makedirs(FLAGS.log_root) if not os.path.exists(model_dir): os.makedirs(model_dir) run_config = tf.ConfigProto() run_config.gpu_options.allow_growth = True with tf.Session( config=tf.ConfigProto(allow_soft_placement=True, log_device_placement=True) ) as sess: """Execute main function.""" os.environ["CUDA_VISIBLE_DEVICES"] = FLAGS.device if not FLAGS.dataset_dir: raise ValueError( "You must supply the dataset directory with " + "--dataset_dir" ) if FLAGS.random_seed >= 0: random.seed(FLAGS.random_seed) np.random.seed(FLAGS.random_seed) tf.logging.set_verbosity(tf.logging.INFO) print("Preparing dataset...") out_shape = [FLAGS.shape_z, FLAGS.shape_y] train_dataset, num_files = mri_data.create_dataset( os.path.join(FLAGS.dataset_dir, "train"), FLAGS.mask_path, num_channels=FLAGS.num_channels, num_emaps=FLAGS.num_emaps, batch_size=FLAGS.batch_size, out_shape=out_shape, ) # channels last format: batch, z, y, channels # placeholders ks_shape = [None, FLAGS.shape_z, FLAGS.shape_y, FLAGS.num_channels] ks_place = tf.placeholder(tf.complex64, ks_shape) sense_shape = [None, FLAGS.shape_z, FLAGS.shape_y, 1, FLAGS.num_channels] sense_place = tf.placeholder(tf.complex64, sense_shape) im_shape = [None, FLAGS.shape_z, FLAGS.shape_y, 1] im_truth_place = tf.placeholder(tf.complex64, im_shape) # run through unrolled model im_out_place = mri_model.unroll_fista( ks_place, sense_place, is_training=True, verbose=True, do_hardproj=FLAGS.do_hard_proj, num_summary_image=FLAGS.num_summary_image, resblock_num_features=FLAGS.feat_map, num_grad_steps=FLAGS.num_grad_steps, conv=FLAGS.conv, do_conjugate=FLAGS.do_conjugate, activation=FLAGS.activation ) # tensorboard summary function _create_summary(sense_place, ks_place, im_out_place, im_truth_place) # define L1 loss between output and ground truth loss = tf.reduce_mean(tf.abs(im_out_place - im_truth_place), name="l1") loss_sum = tf.summary.scalar("loss/l1", loss) # optimize using Adam optimizer = tf.train.AdamOptimizer( learning_rate=FLAGS.learning_rate, name="opt", beta1=FLAGS.adam_beta1, beta2=FLAGS.adam_beta2, ).minimize(loss) # counter for saving checkpoints with tf.variable_scope("counter"): counter = tf.get_variable( "counter", shape=[1], initializer=tf.constant_initializer([0]), dtype=tf.int32, ) update_counter = tf.assign(counter, tf.add(counter, 1)) saver = tf.train.Saver() summary_writer = tf.summary.FileWriter(model_dir, sess.graph) # initialize model print("[*] initializing network...") if not load(model_dir, saver, sess): sess.run(tf.global_variables_initializer()) coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(sess, coord) # calculate number of parameters in model total_parameters = 0 for variable in tf.trainable_variables(): variable_parameters = 1 for dim in variable.get_shape(): variable_parameters *= dim.value total_parameters += variable_parameters print("Total number of trainable parameters: %d" % total_parameters) tf.summary.scalar("parameters/parameters", total_parameters) # use iterator to go through TFrecord dataset train_iterator = train_dataset.make_one_shot_iterator() features, labels = train_iterator.get_next() ks_truth = labels # ground truth kspace ks_in = features["ks_input"] # input kspace sense_in = features["sensemap"] # sensitivity maps mask_recon = features["mask_recon"] # reconstruction mask # ground truth kspace to image domain im_truth = tf_util.model_transpose(ks_truth * mask_recon, sense_in) # gather summaries for tensorboard total_summary = tf.summary.merge_all() print("Start from step %d." % (sess.run(counter))) for step in range(int(sess.run(counter)), FLAGS.max_steps): # evaluate input kspace, sensitivity maps, ground truth image ks_in_run, sense_in_run, im_truth_run = sess.run( [ks_in, sense_in, im_truth] ) # run optimizer and collect output image from model and tensorboard summary im_out, total_summary_run, _ = sess.run( [im_out_place, total_summary, optimizer], feed_dict={ ks_place: ks_in_run, sense_place: sense_in_run, im_truth_place: im_truth_run, }, ) print("step", step) # add summary to tensorboard summary_writer.add_summary(total_summary_run, step) # save checkpoint every 500 steps if step % 500 == 0: print("saving checkpoint") saver.save(sess, model_dir + "/model.ckpt") # update recorded step training is at sess.run(update_counter) print("End of training loop")
def create_model(sess, features, labels, masks, MY, s, architecture='resnet'): # sess: TF sesson # features: input, for SR/CS it is the input image # labels: output, for SR/CS it is the groundtruth image # architecture: aec for encode-decoder, resnet for upside down # Generator rows = int(features.get_shape()[1]) cols = int(features.get_shape()[2]) channels = int(features.get_shape()[3]) gene_minput = tf.placeholder( tf.float32, shape=[FLAGS.batch_size, rows, cols, channels]) gene_mMY = tf.placeholder(tf.complex64, shape=[FLAGS.batch_size, 8, rows, cols]) gene_ms = tf.placeholder(tf.complex64, shape=[FLAGS.batch_size, 8, rows, cols]) if (FLAGS.sampling_pattern != "nomask"): function_generator = lambda x, y, z, w: _generator_model_with_scale( sess, x, y, z, w, num_dc_layers=0, layer_output_skip=7) else: # with unmasked input, remove dc function_generator = lambda x, y, z, w: _generator_model_with_scale( sess, x, y, z, w, num_dc_layers=-1, layer_output_skip=7) rbs = 2 with tf.variable_scope('gene') as scope: if FLAGS.unrolled > 0: gene_output_1 = mri_model.unroll_fista( MY, s, num_grad_steps=FLAGS.unrolled, resblock_num_features=64, resblock_num_blocks=rbs, is_training=True, scope="MRI", mask_output=1, window=None, do_hardproj=True, num_summary_image=0, mask=masks, verbose=False) gene_var_list, gene_layers_1 = None, ['empty'] else: gene_output_1, gene_var_list, gene_layers_1 = function_generator( features, masks, MY, s) scope.reuse_variables() gene_output_real = gene_output_1 gene_output = tf.reshape(gene_output_real, [FLAGS.batch_size, rows, cols, 2]) gene_layers = gene_layers_1 #tf.summary.image('gene_train_last',abs(gene_layers),2) #print('gene_output_train', gene_output.get_shape()) # for testing input if FLAGS.unrolled > 0: gene_moutput_1 = mri_model.unroll_fista( gene_mMY, gene_ms, num_grad_steps=FLAGS.unrolled, resblock_num_features=64, resblock_num_blocks=rbs, is_training=False, scope="MRI", mask_output=1, window=None, do_hardproj=True, num_summary_image=0, mask=masks, verbose=False) gene_mlayers_1 = None else: gene_moutput_1, _, gene_mlayers_1 = function_generator( gene_minput, masks, gene_mMY, gene_ms) scope.reuse_variables() gene_moutput_real = gene_moutput_1 #gene_moutput_complex = tf.complex(gene_moutput_real[:,:,:,0], gene_moutput_real[:,:,:,1]) #gene_moutput = tf.abs(gene_moutput_complex) #print('gene_moutput_test', gene_moutput.get_shape())i gene_moutput = tf.reshape(gene_moutput_real, [FLAGS.batch_size, rows, cols, 2]) gene_mlayers = gene_mlayers_1 # Discriminator with real data if FLAGS.use_phase == True: disc_real_input = tf.identity(labels, name='disc_real_input') else: disc_real_input = tf.sqrt(labels[:, :, :, 0]**2 + labels[:, :, :, 1]**2) disc_real_input = tf.expand_dims(disc_real_input, -1) # TBD: Is there a better way to instance the discriminator? with tf.variable_scope('disc', reuse=tf.AUTO_REUSE) as scope: disc_real_output, disc_var_list, disc_layers_X = \ _discriminator_model(sess, features, disc_real_input, hybrid_disc=FLAGS.hybrid_disc) scope.reuse_variables() if FLAGS.use_phase == True: gene_output_abs = gene_output else: gene_output_abs = tf.sqrt(gene_output[:, :, :, 0]**2 + gene_output[:, :, :, 1]**2) gene_output_abs = tf.expand_dims(gene_output_abs, -1) disc_fake_output, _, disc_layers_Z = _discriminator_model( sess, features, gene_output_abs, hybrid_disc=FLAGS.hybrid_disc) return [ gene_minput, gene_mMY, gene_ms, gene_moutput, gene_output_abs, gene_var_list, gene_layers, gene_mlayers, disc_real_output, disc_fake_output, disc_var_list, disc_real_input ]