def _handler(ir_path, vis_path, model_path, model_pre_path, ssim_weight, index, output_path=None): ir_img = get_train_images(ir_path, flag=False) vis_img = get_train_images(vis_path, flag=False) # ir_img = get_train_images_rgb(ir_path, flag=False) # vis_img = get_train_images_rgb(vis_path, flag=False) dimension = ir_img.shape ir_img = ir_img.reshape([1, dimension[0], dimension[1], dimension[2]]) vis_img = vis_img.reshape([1, dimension[0], dimension[1], dimension[2]]) ir_img = np.transpose(ir_img, (0, 2, 1, 3)) vis_img = np.transpose(vis_img, (0, 2, 1, 3)) print('img shape final:', ir_img.shape) with tf.Graph().as_default(), tf.Session() as sess: infrared_field = tf.placeholder( tf.float32, shape=ir_img.shape, name='content') visible_field = tf.placeholder( tf.float32, shape=ir_img.shape, name='style') dfn = DenseFuseNet(model_pre_path) output_image = dfn.transform_addition(infrared_field, visible_field) # restore the trained model and run the style transferring saver = tf.train.Saver() saver.restore(sess, model_path) output = sess.run(output_image, feed_dict={infrared_field: ir_img, visible_field: vis_img}) save_images(ir_path, output, output_path, prefix='fused' + str(index), suffix='_densefuse_addition_'+str(ssim_weight))
def _handler_l1(ir_path, vis_path, model_path, model_pre_path, ssim_weight, index, output_path=None): ir_img = get_train_images(ir_path, flag=False) vis_img = get_train_images(vis_path, flag=False) dimension = ir_img.shape ir_img = ir_img.reshape([1, dimension[0], dimension[1], dimension[2]]) vis_img = vis_img.reshape([1, dimension[0], dimension[1], dimension[2]]) ir_img = np.transpose(ir_img, (0, 2, 1, 3)) vis_img = np.transpose(vis_img, (0, 2, 1, 3)) print('img shape final:', ir_img.shape) with tf.Graph().as_default(), tf.Session() as sess: # build the dataflow graph infrared_field = tf.placeholder(tf.float32, shape=ir_img.shape, name='content') visible_field = tf.placeholder(tf.float32, shape=ir_img.shape, name='style') dfn = DenseFuseNet(model_pre_path) enc_ir = dfn.transform_encoder(infrared_field) enc_vis = dfn.transform_encoder(visible_field) target = tf.placeholder(tf.float32, shape=enc_ir.shape, name='target') output_image = dfn.transform_decoder(target) # restore the trained model and run the style transferring saver = tf.train.Saver() saver.restore(sess, model_path) enc_ir_temp, enc_vis_temp = sess.run([enc_ir, enc_vis], feed_dict={ infrared_field: ir_img, visible_field: vis_img }) feature = L1_norm(enc_ir_temp, enc_vis_temp) output = sess.run(output_image, feed_dict={target: feature}) save_images(ir_path, output, output_path, prefix='fused' + str(index), suffix='_densefuse_l1norm_' + str(ssim_weight))
def _get_attention(ir_path,vis_path,model_path_a,model_pre_path_a): ir_img = get_train_images(ir_path, flag=False) vis_img = get_train_images(vis_path, flag=False) dimension = ir_img.shape ir_img = ir_img.reshape([1, dimension[0], dimension[1], dimension[2]]) vis_img = vis_img.reshape([1, dimension[0], dimension[1], dimension[2]]) ir_img = np.transpose(ir_img, (0, 2, 1, 3)) vis_img = np.transpose(vis_img, (0, 2, 1, 3)) g1 = tf.Graph() # 加载到Session 1的graph sess1 = tf.Session(graph=g1) # Session1 with sess1.as_default(): with g1.as_default(), tf.Session() as sess: infrared_field = tf.placeholder( tf.float32, shape=ir_img.shape, name='content') visible_field = tf.placeholder( tf.float32, shape=vis_img.shape, name='style') edge_ir = tf.placeholder(tf.float32, shape=ir_img.shape, name='attention') edge_vis = tf.placeholder(tf.float32, shape=ir_img.shape, name='attention') # ----------------------------------------------- image_ir = sess.run(infrared_field, feed_dict={infrared_field: ir_img}) image_vis = sess.run(visible_field, feed_dict={visible_field: vis_img}) p_vis = image_vis[0] p_ir = image_ir[0] p_vis = np.squeeze(p_vis) # 降维 p_ir = np.squeeze(p_ir) guideFilter_img_vis = Grad(p_vis) guideFilter_img_ir = Grad(p_ir) guideFilter_img_vis[guideFilter_img_vis < 0] = 0 guideFilter_img_ir[guideFilter_img_ir < 0] = 0 guideFilter_img_vis = np.expand_dims(guideFilter_img_vis, axis=-1) guideFilter_img_ir = np.expand_dims(guideFilter_img_ir, axis=-1) guideFilter_img_vis = np.expand_dims(guideFilter_img_vis, axis=0) guideFilter_img_ir = np.expand_dims(guideFilter_img_ir, axis=0) a = attention.Attention(model_pre_path_a) saver = tf.train.Saver() saver.restore(sess, model_path_a) feature_a=a.get_attention(edge_ir) feature_b=a.get_attention(edge_vis) edge_ir_temp = sess.run([feature_a], feed_dict={edge_ir: guideFilter_img_ir}) edge_vis_temp = sess.run([feature_b], feed_dict={edge_vis: guideFilter_img_vis}) '''feature_a = a.get_attention(edge_ir_temp) feature_b = a.get_attention(edge_vis_temp)''' return edge_ir_temp,edge_vis_temp
def _handler(content_path, style_path, encoder_path, model_path, model_pre_path, output_path=None): with tf.Graph().as_default(), tf.Session() as sess: index = 2 content_path = content_path + str(index) + '.jpg' style_path = style_path + str(index) + '.jpg' content_img = get_train_images(content_path) style_img = get_train_images(style_path) # build the dataflow graph content = tf.placeholder(tf.float32, shape=content_img.shape, name='content') style = tf.placeholder(tf.float32, shape=style_img.shape, name='style') stn = StyleTransferNet(encoder_path, model_pre_path) enc_c, enc_s = stn.encoder_process(content, style) target = tf.placeholder(tf.float32, shape=enc_c.shape, name='target') # output_image = stn.transform(content, style) output_image = stn.decoder_process(target) # restore the trained model and run the style transferring saver = tf.train.Saver() saver.restore(sess, model_path) # get the output enc_c, enc_s = sess.run([enc_c, enc_s], feed_dict={ content: content_img, style: style_img }) feature = L1_Max(enc_c, enc_s) # feature = enc_s output = sess.run(output_image, feed_dict={target: feature}) if output_path is not None: save_images(content_path, output, output_path, prefix='fused' + str(index) + '_', suffix='deep') return output
def _handler(content_name, style_name, model_path, model_pre_path, index, output_path=None): content_path = content_name style_path = style_name content_img = get_train_images(content_path, flag=False) style_img = get_train_images(style_path, flag=False) dimension = content_img.shape content_img = content_img.reshape( [1, dimension[0], dimension[1], dimension[2]]) style_img = style_img.reshape( [1, dimension[0], dimension[1], dimension[2]]) content_img = np.transpose(content_img, (0, 2, 1, 3)) style_img = np.transpose(style_img, (0, 2, 1, 3)) print('content_img shape final:', content_img.shape) with tf.Graph().as_default(), tf.Session() as sess: # build the dataflow graph content = tf.placeholder(tf.float32, shape=content_img.shape, name='content') style = tf.placeholder(tf.float32, shape=style_img.shape, name='style') dfn = DeepFuseNet(model_pre_path) output_image = dfn.transform_addition(content, style) # output_image = dfn.transform_recons(style) # output_image = dfn.transform_recons(content) # restore the trained model and run the style transferring saver = tf.train.Saver() saver.restore(sess, model_path) output = sess.run(output_image, feed_dict={ content: content_img, style: style_img }) save_images(content_path, output, output_path, prefix='fused' + str(index), suffix='_deepfuse_bs2_epoch2') return output
def _handler_video(ir_path, vis_path, model_path, model_pre_path, ssim_weight, output_path=None): infrared = ir_path[0] img = get_train_images(infrared, flag=False) img = img.reshape([1, img.shape[0], img.shape[1], img.shape[2]]) img = np.transpose(img, (0, 2, 1, 3)) print('img shape final:', img.shape) num_imgs = len(ir_path) with tf.Graph().as_default(), tf.Session() as sess: # build the dataflow graph infrared_field = tf.placeholder( tf.float32, shape=img.shape, name='content') visible_field = tf.placeholder( tf.float32, shape=img.shape, name='style') dfn = DenseFuseNet(model_pre_path) output_image = dfn.transform_addition(infrared_field, visible_field) # restore the trained model and run the style transferring saver = tf.train.Saver() saver.restore(sess, model_path) ##################GET IMAGES################################################################################### start_time = datetime.now() for i in range(num_imgs): print('image number:', i) infrared = ir_path[i] visible = vis_path[i] ir_img = get_train_images(infrared, flag=False) vis_img = get_train_images(visible, flag=False) dimension = ir_img.shape ir_img = ir_img.reshape([1, dimension[0], dimension[1], dimension[2]]) vis_img = vis_img.reshape([1, dimension[0], dimension[1], dimension[2]]) ir_img = np.transpose(ir_img, (0, 2, 1, 3)) vis_img = np.transpose(vis_img, (0, 2, 1, 3)) ################FEED######################################## output = sess.run(output_image, feed_dict={infrared_field: ir_img, visible_field: vis_img}) save_images(infrared, output, output_path, prefix='fused' + str(i), suffix='_addition_' + str(ssim_weight)) ###################################################################################################### elapsed_time = datetime.now() - start_time print('Dense block video==> elapsed time: %s' % (elapsed_time))
def _handler(ir_path, vis_path, model_path, model_pre_path, ssim_weight, index, output_path=None): ir_img = get_train_images(ir_path, flag=False) vis_img = get_train_images(vis_path, flag=False) # ir_img = get_train_images_rgb(ir_path, flag=False) # vis_img = get_train_images_rgb(vis_path, flag=False) dimension = ir_img.shape ir_img = ir_img.reshape([1, dimension[0], dimension[1], dimension[2]]) vis_img = vis_img.reshape([1, dimension[0], dimension[1], dimension[2]]) ir_img = np.transpose(ir_img, (0, 2, 1, 3)) vis_img = np.transpose(vis_img, (0, 2, 1, 3)) print('img shape final:', ir_img.shape) with tf.Graph().as_default(), tf.Session() as sess: infrared_field = tf.placeholder(tf.float32, shape=ir_img.shape, name='content') visible_field = tf.placeholder(tf.float32, shape=ir_img.shape, name='style') dfn = DenseFuseNet(model_pre_path) output_image = dfn.transform_addition(infrared_field, visible_field) # restore the trained model and run the style transferring saver = tf.train.Saver() saver.restore(sess, model_path) output = sess.run(output_image, feed_dict={ infrared_field: ir_img, visible_field: vis_img }) save_images(ir_path, output, output_path, prefix='3', suffix='') img333 = Image.open('/home/bingyang/wby/1/3.png') imgout = img333.transpose(Image.FLIP_LEFT_RIGHT) imgout = imgout.transpose(Image.ROTATE_90) imgout.save('/home/bingyang/wby/1/3.png')
def _handler_l1(content_name, style_name, model_path, model_pre_path, ssim_weight, index, output_path=None): infrared_path = content_name visible_path = style_name content_img = get_train_images(infrared_path, flag=False) style_img = get_train_images(visible_path, flag=False) dimension = content_img.shape content_img = content_img.reshape([1, dimension[0], dimension[1], dimension[2]]) style_img = style_img.reshape([1, dimension[0], dimension[1], dimension[2]]) content_img = np.transpose(content_img, (0, 2, 1, 3)) style_img = np.transpose(style_img, (0, 2, 1, 3)) print('content_img shape final:', content_img.shape) with tf.Graph().as_default(), tf.Session() as sess: # build the dataflow graph content = tf.placeholder( tf.float32, shape=content_img.shape, name='content') style = tf.placeholder( tf.float32, shape=style_img.shape, name='style') dfn = DenseFuseNet(model_pre_path) enc_c = dfn.transform_encoder(content) enc_s = dfn.transform_encoder(style) target = tf.placeholder( tf.float32, shape=enc_c.shape, name='target') output_image = dfn.transform_decoder(target) # restore the trained model and run the style transferring saver = tf.train.Saver() saver.restore(sess, model_path) enc_c_temp, enc_s_temp = sess.run([enc_c, enc_s], feed_dict={content: content_img, style: style_img}) feature = L1_norm(enc_c_temp, enc_s_temp) output = sess.run(output_image, feed_dict={target: feature}) save_images(infrared_path, output, output_path, prefix='fused' + str(index), suffix='_densefuse_l1norm_'+str(ssim_weight)) return output
def test_utils_imoperations(): from utils import imread, imresize_square, get_train_images, imsave path_read = '/tmp/panda.jpg' path_save = '/tmp/panda_resized.jpg' # image = imread(path_read, mode='RGB') # image = imresize_square(image, long_side=256, interp = 'nearest') # imsave(path_save, image) images = get_train_images([path_read]) imsave(path_save, images[0])
def run(gpuids, q): # scan all files under img_path names = get_train_images() # init scheduler x = Scheduler(gpuids, q) # start processing and wait for complete return x.start(names)
def train_recons(inputPath, validationPath, save_path, model_pre_path, EPOCHES_set, BATCH_SIZE, debug=False, logging_period=1): from datetime import datetime start_time = datetime.now() path = './models/performanceData/' fileName = 'TrainPerformanceData_'+str(start_time)+'.txt' fileName = fileName.replace(" ", "_") fileName = fileName.replace(":", "_") file = open(path+fileName, 'w') file.close() folders = list_folders(inputPath) valFolders = list_folders(validationPath) EPOCHS = EPOCHES_set print("EPOCHES : ", EPOCHS) print("BATCH_SIZE: ", BATCH_SIZE) # get the traing image shape HEIGHT, WIDTH, CHANNELS = TRAINING_IMAGE_SHAPE INPUT_SHAPE = (BATCH_SIZE, HEIGHT, WIDTH, CHANNELS) HEIGHT_OR, WIDTH_OR, CHANNELS_OR = TRAINING_IMAGE_SHAPE_OR INPUT_SHAPE_OR = (BATCH_SIZE, HEIGHT_OR, WIDTH_OR, CHANNELS_OR) GROUNDTRUTH_SHAPE_OR = (1, HEIGHT_OR, WIDTH_OR, CHANNELS_OR) # create the graph with tf.Graph().as_default(), tf.Session() as sess: original = tf.placeholder(tf.float32, shape=INPUT_SHAPE_OR, name='original') groundtruth = tf.placeholder(tf.float32, shape=GROUNDTRUTH_SHAPE_OR, name='groundtruth') source = original print('source :', source.shape) print('original:', original.shape) print('groundtruth:', groundtruth.shape) # create the deepfuse net (encoder and decoder) dfn = DenseFuseNet(model_pre_path) generated_img = dfn.transform_recons_train(source) print('generate:', generated_img.shape) pixel_loss = tf.reduce_sum(tf.square(groundtruth - generated_img)) pixel_loss = tf.math.sqrt(pixel_loss / (BATCH_SIZE * HEIGHT * WIDTH)) loss = pixel_loss train_op = tf.train.AdamOptimizer(LEARNING_RATE).minimize(loss) sess.run(tf.global_variables_initializer()) # saver = tf.train.Saver() saver = tf.train.Saver(keep_checkpoint_every_n_hours=1) # ** Start Training ** step = 0 count_loss = 0 numTrainSite = len(folders) numValSite = len(valFolders) for epoch in range(EPOCHS): save_path_epoc = './models_intermediate/'+str(epoch)+'.ckpt' start_time_epoc = datetime.now() for site in range(numTrainSite): start_time_site = datetime.now() file = open(path + fileName, 'a') groundtruth_imgs_path = list_images(inputPath+folders[site] + '/gt/') training_imgs_path = list_images(inputPath+folders[site]) np.random.shuffle(training_imgs_path) gt = get_train_images(groundtruth_imgs_path, crop_height=HEIGHT, crop_width=WIDTH, flag=False) gtImgTrain = np.zeros(GROUNDTRUTH_SHAPE_OR) gtImgTrain[0] = gt n_batches = int(len(training_imgs_path) // BATCH_SIZE) for batch in range(n_batches): original_path = training_imgs_path[batch * BATCH_SIZE:(batch * BATCH_SIZE + BATCH_SIZE)] original_batch = get_train_images(original_path, crop_height=HEIGHT, crop_width=WIDTH, flag=False) original_batch = original_batch.reshape([BATCH_SIZE, 256, 256, 1]) sess.run(train_op, feed_dict={original: original_batch, groundtruth: gtImgTrain}) if debug: for batch in range(n_batches): original_path = training_imgs_path[batch * BATCH_SIZE:(batch * BATCH_SIZE + BATCH_SIZE)] original_batch = get_train_images(original_path, crop_height=HEIGHT, crop_width=WIDTH, flag=False) original_batch = original_batch.reshape([BATCH_SIZE, 256, 256, 1]) # print('original_batch shape final:', original_batch.shape) # run the training step _p_loss = sess.run(pixel_loss, feed_dict={original: original_batch, groundtruth: gtImgTrain}) # add text file to add mode(validation/training), epoch#, site#, batch#, _p_loss) ------------------------------ file.write('Train[Epoch#: %d, Site#: %d, Batch#: %d, _p_loss: %d]\n' % (epoch, site, batch, _p_loss)) print('Train[Epoch#: %d, Site#: %d, Batch#: %d, _p_loss: %d]' % (epoch, site, batch, _p_loss)) print('Time taken per site: %s' %(datetime.now() - start_time_site)) file.close() for site in range(numValSite): file = open(path + fileName, 'a') start_time_validation = datetime.now() groundtruth_val_imgs_path = list_images(validationPath+valFolders[site] + '/gt/') validation_imgs_path = list_images(validationPath+valFolders[site]) np.random.shuffle(validation_imgs_path) gt = get_train_images(groundtruth_val_imgs_path, crop_height=HEIGHT, crop_width=WIDTH, flag=False) gtImgVal = np.zeros(GROUNDTRUTH_SHAPE_OR) gtImgVal[0] = gt val_batches = int(len(validation_imgs_path) // BATCH_SIZE) val_pixel_acc = 0 for batch in range(val_batches): val_original_path = validation_imgs_path[batch * BATCH_SIZE:(batch * BATCH_SIZE + BATCH_SIZE)] val_original_batch = get_train_images(val_original_path, crop_height=HEIGHT, crop_width=WIDTH, flag=False) val_original_batch = val_original_batch.reshape([BATCH_SIZE, 256, 256, 1]) val_pixel = sess.run(pixel_loss, feed_dict={original: val_original_batch, groundtruth: gtImgVal}) file.write('Validation[Epoch#: %d, Site#: %d, Batch#: %d, _p_loss: %d]\n' % (epoch, site, batch, val_pixel)) val_pixel_acc = val_pixel_acc + val_pixel print('Time taken per validation site: %s' % (datetime.now() - start_time_validation)) val_loss = val_pixel_acc/val_batches file.write('ValidationAcc[Epoch#: %d, Site#: %d, Batch#: %d, val_loss: %d]\n' % (epoch, site, batch, val_loss)) print('ValidationAcc[Epoch#: %d, Site#: %d, Batch#: %d, _p_loss: %d]' % (epoch, site, batch, val_loss)) file.close() print('------------------------------------------------------------------------------') print('Time taken per epoc: %s' % (datetime.now() - start_time_epoc)) saver.save(sess, save_path_epoc) saver.save(sess, save_path) print('Done training!') print('Total Time taken (training): %s' % (datetime.now() - start_time)) file.close()
def train(ssim_weight, original_imgs_path_name, source_a_imgs_path, source_b_imgs_path_name, encoder_path, save_path, model_pre_path, debug=False, logging_period=100): if debug: from datetime import datetime start_time = datetime.now() # num_imgs = len(source_a_imgs_path) num_imgs = 10000 source_a_imgs_path = source_a_imgs_path[:num_imgs] mod = num_imgs % BATCH_SIZE print('Train images number %d.\n' % num_imgs) print('Train images samples %s.\n' % str(num_imgs / BATCH_SIZE)) if mod > 0: print('Train set has been trimmed %d samples...\n' % mod) source_a_imgs_path = source_a_imgs_path[:-mod] # get the traing image shape HEIGHT, WIDTH, CHANNELS = TRAINING_IMAGE_SHAPE INPUT_SHAPE = (BATCH_SIZE, HEIGHT, WIDTH, CHANNELS) HEIGHT_OR, WIDTH_OR, CHANNELS_OR = TRAINING_IMAGE_SHAPE_OR INPUT_SHAPE_OR = (BATCH_SIZE, HEIGHT_OR, WIDTH_OR, CHANNELS_OR) # create the graph with tf.Graph().as_default(), tf.Session() as sess: original = tf.placeholder(tf.float32, shape=INPUT_SHAPE_OR, name='original') source_a = tf.placeholder(tf.float32, shape=INPUT_SHAPE, name='source_a') source_b = tf.placeholder(tf.float32, shape=INPUT_SHAPE, name='source_b') print('source:', source_a.shape) # create the style transfer net stn = StyleTransferNet(encoder_path, model_pre_path) # pass content and style to the stn, getting the generated_img, fused image generated_img = stn.transform(source_a, source_b) # # get the target feature maps which is the output of AdaIN # target_features = stn.target_features pixel_loss = tf.reduce_sum( tf.reduce_mean(tf.square(original - generated_img), axis=[1, 2])) pixel_loss = pixel_loss / (HEIGHT * WIDTH) # compute the SSIM loss ssim_loss = 1 - SSIM.tf_ssim(original, generated_img) # compute the total loss loss = pixel_loss + ssim_weight * ssim_loss # Training step train_op = tf.train.AdamOptimizer(LEARNING_RATE).minimize(loss) sess.run(tf.global_variables_initializer()) # saver = tf.train.Saver() saver = tf.train.Saver(keep_checkpoint_every_n_hours=1) # ** Start Training ** step = 0 count_loss = 0 n_batches = int(len(source_a_imgs_path) // BATCH_SIZE) if debug: elapsed_time = datetime.now() - start_time print( '\nElapsed time for preprocessing before actually train the model: %s' % elapsed_time) print('Now begin to train the model...\n') start_time = datetime.now() Loss_all = [i for i in range(EPOCHS * n_batches)] for epoch in range(EPOCHS): np.random.shuffle(source_a_imgs_path) for batch in range(n_batches): # retrive a batch of content and style images source_a_path = source_a_imgs_path[batch * BATCH_SIZE:( batch * BATCH_SIZE + BATCH_SIZE)] source_a_str = source_a_path[0] name_f = source_a_str.find('\\') source_image_name = source_a_str[name_f + 1:] source_image_name_comm = source_image_name[2:] source_b_path = [source_b_imgs_path_name + source_image_name] original_path = [ original_imgs_path_name + source_image_name_comm ] original_batch = get_train_images(original_path, crop_height=HEIGHT, crop_width=WIDTH, flag=False) source_a_batch = get_train_images(source_a_path, crop_height=HEIGHT, crop_width=WIDTH) source_b_batch = get_train_images(source_b_path, crop_height=HEIGHT, crop_width=WIDTH) original_batch = original_batch.reshape([1, 256, 256, 1]) # run the training step sess.run(train_op, feed_dict={ original: original_batch, source_a: source_a_batch, source_b: source_b_batch }) step += 1 # if step % 1000 == 0: # saver.save(sess, save_path, global_step=step) if debug: is_last_step = (epoch == EPOCHS - 1) and (batch == n_batches - 1) if is_last_step or step % logging_period == 0: elapsed_time = datetime.now() - start_time _pixel_loss, _ssim_loss, _loss = sess.run( [pixel_loss, ssim_loss, loss], feed_dict={ original: original_batch, source_a: source_a_batch, source_b: source_b_batch }) Loss_all[count_loss] = _loss count_loss += 1 print( 'step: %d, total loss: %.3f, elapsed time: %s' % (step, _loss, elapsed_time)) print('pixel loss: %.3f' % (_pixel_loss)) print('ssim loss : %.3f\n' % (_ssim_loss)) # print('pca or shape : ', _pca_or.shape) # print('pca gen shape : ', _pca_gen.shape) # ** Done Training & Save the model ** saver.save(sess, save_path) iter_index = [i for i in range(count_loss)] plt.plot(iter_index, Loss_all[:count_loss]) plt.show() if debug: elapsed_time = datetime.now() - start_time print('Done training! Elapsed time: %s' % elapsed_time) print('Model is saved to: %s' % save_path)
def train(style_weight, content_imgs_path, style_imgs_path, encoder_path, model_save_path, debug=False, logging_period=100): if debug: from datetime import datetime start_time = datetime.now() # guarantee the size of content and style images to be a multiple of BATCH_SIZE num_imgs = min(len(content_imgs_path), len(style_imgs_path)) content_imgs_path = content_imgs_path[:num_imgs] style_imgs_path = style_imgs_path[:num_imgs] mod = num_imgs % BATCH_SIZE if mod > 0: print('Train set has been trimmed %d samples...\n' % mod) content_imgs_path = content_imgs_path[:-mod] style_imgs_path = style_imgs_path[:-mod] # get the traing image shape HEIGHT, WIDTH, CHANNELS = TRAINING_IMAGE_SHAPE INPUT_SHAPE = (BATCH_SIZE, HEIGHT, WIDTH, CHANNELS) # create the graph with tf.Graph().as_default(), tf.Session() as sess: content = tf.placeholder(tf.float32, shape=INPUT_SHAPE, name='content') style = tf.placeholder(tf.float32, shape=INPUT_SHAPE, name='style') # create the style transfer net stn = StyleTransferNet(encoder_path) # pass content and style to the stn, getting the generated_img generated_img = stn.transform(content, style) # get the target feature maps which is the output of AdaIN target_features = stn.target_features # pass the generated_img to the encoder, and use the output compute loss generated_img = tf.reverse(generated_img, axis=[-1]) # switch RGB to BGR generated_img = stn.encoder.preprocess(generated_img) # preprocess image enc_gen, enc_gen_layers = stn.encoder.encode(generated_img) # compute the content loss # content_loss = fft_loss(enc_gen, target_features) content_loss = tf.reduce_sum(tf.reduce_mean(tf.square(enc_gen - target_features), axis=[1, 2])) # compute the style loss style_layer_loss = [] for layer in STYLE_LAYERS: enc_style_feat = stn.encoded_style_layers[layer] enc_gen_feat = enc_gen_layers[layer] meanS, varS = tf.nn.moments(enc_style_feat, [1, 2]) meanG, varG = tf.nn.moments(enc_gen_feat, [1, 2]) # fft_pred = tf.abs(tf.spectral.rfft2d(tf.transpose(enc_gen_feat, (0, 3, 1, 2)))) # fft_true = tf.abs(tf.spectral.rfft2d(tf.transpose(enc_style_feat, (0, 3, 1, 2)))) # meanS, varS = tf.nn.moments(fft_pred, [2, 3]) # meanG, varG = tf.nn.moments(fft_true, [2, 3]) sigmaS = tf.sqrt(varS + EPSILON) sigmaG = tf.sqrt(varG + EPSILON) l2_mean = tf.reduce_sum(tf.square(meanG - meanS)) l2_sigma = tf.reduce_sum(tf.square(sigmaG - sigmaS)) style_layer_loss.append(l2_mean + l2_sigma) style_loss = tf.reduce_sum(style_layer_loss) # compute the total loss loss = content_loss + style_weight * style_loss # Training step global_step = tf.Variable(0, trainable=False) learning_rate = tf.train.inverse_time_decay(LEARNING_RATE, global_step, DECAY_STEPS, LR_DECAY_RATE) train_op = tf.train.AdamOptimizer(learning_rate).minimize(loss, global_step=global_step) sess.run(tf.global_variables_initializer()) # saver saver = tf.train.Saver(max_to_keep=10) ###### Start Training ###### step = 0 n_batches = int(len(content_imgs_path) // BATCH_SIZE) if debug: elapsed_time = datetime.now() - start_time start_time = datetime.now() print('\nElapsed time for preprocessing before actually train the model: %s' % elapsed_time) print('Now begin to train the model...\n') try: for epoch in range(EPOCHS): np.random.shuffle(content_imgs_path) np.random.shuffle(style_imgs_path) for batch in range(n_batches): # retrive a batch of content and style images content_batch_path = content_imgs_path[batch*BATCH_SIZE:(batch*BATCH_SIZE + BATCH_SIZE)] style_batch_path = style_imgs_path[batch*BATCH_SIZE:(batch*BATCH_SIZE + BATCH_SIZE)] content_batch = get_train_images(content_batch_path, crop_height=HEIGHT, crop_width=WIDTH) style_batch = get_train_images(style_batch_path, crop_height=HEIGHT, crop_width=WIDTH) # run the training step sess.run(train_op, feed_dict={content: content_batch, style: style_batch}) step += 1 if step % 1000 == 0: saver.save(sess, model_save_path, global_step=step, write_meta_graph=False) if debug: is_last_step = (epoch == EPOCHS - 1) and (batch == n_batches - 1) if is_last_step or step == 1 or step % logging_period == 0: elapsed_time = datetime.now() - start_time _content_loss, _style_loss, _loss = sess.run([content_loss, style_loss, loss], feed_dict={content: content_batch, style: style_batch}) print('step: %d, total loss: %.3f, elapsed time: %s' % (step, _loss, elapsed_time)) print('content loss: %.3f' % (_content_loss)) print('style loss : %.3f, weighted style loss: %.3f\n' % (_style_loss, style_weight * _style_loss)) except Exception as ex: saver.save(sess, model_save_path, global_step=step) print('\nSomething wrong happens! Current model is saved to <%s>' % tmp_save_path) print('Error message: %s' % str(ex)) ###### Done Training & Save the model ###### saver.save(sess, model_save_path) if debug: elapsed_time = datetime.now() - start_time print('Done training! Elapsed time: %s' % elapsed_time) print('Model is saved to: %s' % model_save_path)
def _handler_l1(ir_path, vis_path, model_path, model_pre_path, ssim_weight, index, output_path=None): ir_img = get_train_images(ir_path, True, flag=False) vis_img, Cr, Cb = get_train_images(vis_path, False, flag=False) dimension = ir_img.shape ir_img = ir_img.reshape([1, dimension[0], dimension[1], dimension[2]]) vis_img = vis_img.reshape([1, dimension[0], dimension[1], dimension[2]]) #ir_img = np.transpose(ir_img, (0, 2, 1, 3)) #vis_img = np.transpose(vis_img, (0, 2, 1, 3)) print('img shape final:', ir_img.shape) with tf.Graph().as_default(), tf.Session() as sess: # build the dataflow graph infrared_field = tf.placeholder(tf.float32, shape=ir_img.shape, name='content') visible_field = tf.placeholder(tf.float32, shape=ir_img.shape, name='style') dfn = DenseFuseNet(model_pre_path) enc_ir1, enc_ir2, enc_ir3 = dfn.transform_encoder(infrared_field) enc_vis1, enc_vis2, enc_vis3 = dfn.transform_encoder(visible_field) temp_enc = tf.concat([enc_ir1, enc_ir2], 3) temp_out = tf.concat([temp_enc, enc_ir3], 3) target = tf.placeholder(tf.float32, shape=temp_out.shape, name='target') output_image = dfn.transform_decoder(target) # restore the trained model and run the style transferring saver = tf.train.Saver() saver.restore(sess, model_path) enc_ir1, enc_ir2, enc_ir3, enc_vis1, enc_vis2, enc_vis3 = sess.run( [enc_ir1, enc_ir2, enc_ir3, enc_vis1, enc_vis2, enc_vis3], feed_dict={ infrared_field: ir_img, visible_field: vis_img }) target_feature1_l1 = L1_norm(enc_ir1, enc_vis1) target_feature2_l1 = L1_norm(enc_ir2, enc_vis2) target_feature3_l1 = L1_norm(enc_ir3, enc_vis3) temp_l1 = tf.concat([target_feature1_l1, target_feature2_l1], 3) out_l1 = tf.concat([temp_l1, target_feature3_l1], 3) out_l1 = out_l1.eval() output = sess.run(output_image, feed_dict={target: out_l1}) output = output.squeeze() result = np.dstack([output, Cr, Cb]) result = cv2.cvtColor(result, cv2.COLOR_YCrCb2BGR) cv2.imwrite((output_path + str(index) + '_multiScale_densefuse_l1_MRSPECT_ssim' + str(ssim_weight) + '.jpg'), result)
proc.start() out_queue = run(gpuids, q) out_dict = {} while out_queue.qsize() > 0: item = out_queue.get() out_dict[item['image_name']] = item['embedding'] with open("data/train_embeddings.p", "wb") as file: pickle.dump(out_dict, file) q.put(None) proc.join() train_images = get_train_images() def calculate_distance_list(image_i): embedding_i = embeddings[image_i] distance_list = np.empty(shape=(num_train_samples,), dtype=np.float32) for j, image_j in enumerate(train_images): embedding_j = embeddings[image_j] dist = np.square(np.linalg.norm(embedding_i - embedding_j)) distance_list[j] = dist return distance_list if __name__ == '__main__': print('creating train embeddings') create_train_embeddings()
def train_recons(original_imgs_path, save_path, model_pre_path, EPOCHES_set, BATCH_SIZE, debug=False, logging_period=100): if debug: from datetime import datetime start_time = datetime.now() EPOCHS = EPOCHES_set print("EPOCHES : ", EPOCHS) print("BATCH_SIZE: ", BATCH_SIZE) num_imgs = len(original_imgs_path) # num_imgs = 100 original_imgs_path = original_imgs_path[:num_imgs] mod = num_imgs % BATCH_SIZE print('Train images number %d.\n' % num_imgs) print('Train images samples %s.\n' % str(num_imgs / BATCH_SIZE)) if mod > 0: print('Train set has been trimmed %d samples...\n' % mod) original_imgs_path = original_imgs_path[:-mod] # get the traing image shape HEIGHT, WIDTH, CHANNELS = TRAINING_IMAGE_SHAPE INPUT_SHAPE = (BATCH_SIZE, HEIGHT, WIDTH, CHANNELS) HEIGHT_OR, WIDTH_OR, CHANNELS_OR = TRAINING_IMAGE_SHAPE_OR INPUT_SHAPE_OR = (BATCH_SIZE, HEIGHT_OR, WIDTH_OR, CHANNELS_OR) # create the graph with tf.Graph().as_default(), tf.Session() as sess: original = tf.placeholder(tf.float32, shape=INPUT_SHAPE_OR, name='original') source = original print('source :', source.shape) print('original:', original.shape) # create the deepfuse net (encoder and decoder) dfn = DeepFuseNet(model_pre_path) generated_img = dfn.transform_recons(source) print('generate:', generated_img.shape) ssim_loss = SSIM_LOSS(original, generated_img) loss = 1 - ssim_loss train_op = tf.train.AdamOptimizer(LEARNING_RATE).minimize(loss) sess.run(tf.global_variables_initializer()) # saver = tf.train.Saver() saver = tf.train.Saver(keep_checkpoint_every_n_hours=1) # ** Start Training ** step = 0 count_loss = 0 n_batches = int(len(original_imgs_path) // BATCH_SIZE) if debug: elapsed_time = datetime.now() - start_time print( '\nElapsed time for preprocessing before actually train the model: %s' % elapsed_time) print('Now begin to train the model...\n') start_time = datetime.now() Loss_all = [i for i in range(EPOCHS * n_batches)] for epoch in range(EPOCHS): np.random.shuffle(original_imgs_path) for batch in range(n_batches): # retrive a batch of content and style images original_path = original_imgs_path[batch * BATCH_SIZE:( batch * BATCH_SIZE + BATCH_SIZE)] original_batch = get_train_images(original_path, crop_height=HEIGHT, crop_width=WIDTH, flag=False) original_batch = original_batch.reshape( [BATCH_SIZE, 256, 256, 1]) # print('original_batch shape final:', original_batch.shape) # run the training step sess.run(train_op, feed_dict={original: original_batch}) step += 1 # if step % 1000 == 0: # saver.save(sess, save_path, global_step=step) if debug: is_last_step = (epoch == EPOCHS - 1) and (batch == n_batches - 1) if is_last_step or step % logging_period == 0: elapsed_time = datetime.now() - start_time _ssim_loss, _loss = sess.run( [ssim_loss, loss], feed_dict={original: original_batch}) Loss_all[count_loss] = _loss count_loss += 1 print( 'Deep fuse==>>step: %d, total loss: %s, elapsed time: %s' % (step, _loss, elapsed_time)) print('ssim_loss: %s ' % (_ssim_loss)) # ** Done Training & Save the model ** saver.save(sess, save_path) loss_data = Loss_all[:count_loss] scio.savemat( 'D:/project/GitHub/ImageFusion/Imagefusion_deepfuse/DeepFuseLossData.mat', {'loss': loss_data}) # iter_index = [i for i in range(count_loss)] # plt.plot(iter_index, Loss_all[:count_loss]) # plt.show() if debug: elapsed_time = datetime.now() - start_time print('Done training! Elapsed time: %s' % elapsed_time) print('Model is saved to: %s' % save_path)
def train_recons(original_imgs_path, save_path, model_pre_path, EPOCHES_set, BATCH_SIZE_set, debug=False, logging_period=1): from datetime import datetime if debug: start_time = datetime.now() EPOCHS = EPOCHES_set BATCH_SIZE = BATCH_SIZE_set print("EPOCHES : ", EPOCHS) print("BATCH_SIZE: ", BATCH_SIZE) num_imgs = len(original_imgs_path) mod = num_imgs % BATCH_SIZE print('Train images number {}.'.format(num_imgs)) print('Train images samples {}.'.format(num_imgs // BATCH_SIZE)) if mod > 0: print('Train set has been trimmed {} samples...'.format(mod)) original_imgs_path = original_imgs_path[:-mod] # get the traing image shape INPUT_SHAPE = (BATCH_SIZE, HEIGHT, WIDTH, CHANNELS) # create the graph with tf.compat.v1.Graph().as_default(), tf.compat.v1.Session() as sess: with tf.compat.v1.name_scope('Input'): original = tf.placeholder(tf.float32, shape=INPUT_SHAPE, name='original') source = original print('source :', source.shape) print('original :', original.shape) # create the deepfuse net (encoder and decoder) dfn = DenseFuseNet(model_pre_path) generated_img = dfn.transform_recons(source) print('generate:', generated_img.shape) epsilon_1 = tf.reduce_mean(tf.square(generated_img - original)) epsilon_2 = 1 - tf.reduce_mean( tf.image.ssim(generated_img, original, max_val=1.0)) total_loss = epsilon_1 + 1000 * epsilon_2 tf.compat.v1.summary.scalar('epsilon_1', epsilon_1) tf.compat.v1.summary.scalar('epsilon_2', epsilon_2) tf.compat.v1.summary.scalar('total_loss', total_loss) train_op = tf.compat.v1.train.AdamOptimizer(LEARNING_RATE).minimize( total_loss) summary_op = tf.compat.v1.summary.merge_all() train_writer = tf.compat.v1.summary.FileWriter('./models/log', sess.graph, flush_secs=60) train_writer.add_graph(sess.graph) sess.run(tf.compat.v1.global_variables_initializer()) # saver = tf.train.Saver() saver = tf.compat.v1.train.Saver(max_to_keep=20) # ** Start Training ** step = 0 n_batches = int(len(original_imgs_path) // BATCH_SIZE) if debug: elapsed_time = datetime.now() - start_time print( 'Elapsed time for preprocessing before actually train the model: {}' .format(elapsed_time)) print('Now begin to train the model...') start_time = datetime.now() Loss_1 = [] Loss_2 = [] Loss_all = [] for epoch in range(EPOCHS): for batch in range(n_batches): # retrive a batch of infrared and visiable images original_path = original_imgs_path[batch * BATCH_SIZE:( batch * BATCH_SIZE + BATCH_SIZE)] original_batch = get_train_images(original_path) # print(original_batch.shape) original_batch = original_batch.transpose((3, 0, 1, 2)) # run the training step step += 1 _, summary_str, _epsilon_1, _epsilon_2, _total_loss = sess.run( [train_op, summary_op, epsilon_1, epsilon_2, total_loss], feed_dict={original: original_batch}) train_writer.add_summary(summary_str, step) Loss_1.append(_epsilon_1) Loss_2.append(_epsilon_2) Loss_all.append(_total_loss) if debug: is_last_step = (epoch == EPOCHS - 1) and (batch == n_batches - 1) if is_last_step or step % logging_period == 0: elapsed_time = datetime.now() - start_time print( 'epoch:{:>2}/{}, step:{:>4}, total loss: {:.4f}, elapsed time: {}' .format(epoch + 1, EPOCHS, step, _total_loss, elapsed_time)) print('epsilon_1: {}, epsilon_2: {}\n'.format( _epsilon_1, _epsilon_2)) # ** Done Training & Save the model ** saver.save(sess, save_path, global_step=epoch + 1) if not os.path.exists('./models/loss/'): os.mkdir('./models/loss/') scio.savemat('./models/loss/TotalLoss_' + str(epoch + 1) + '.mat', {'total_loss': Loss_all}) scio.savemat('./models/loss/Epsilon1_' + str(epoch + 1) + '.mat', {'epsilon_1': Loss_1}) scio.savemat('./models/loss/Epsilon2_' + str(epoch + 1) + '.mat', {'epsilon_2': Loss_2}) if debug: elapsed_time = datetime.now() - start_time print('Done training! Elapsed time: {}'.format(elapsed_time)) print('Model is saved to: {}'.format(save_path))
def train_recons(original_imgs_path, validatioin_imgs_path, save_path, model_pre_path, ssim_weight, EPOCHES_set, BATCH_SIZE, debug=False, logging_period=1): if debug: from datetime import datetime start_time = datetime.now() EPOCHS = EPOCHES_set print("EPOCHES : ", EPOCHS) print("BATCH_SIZE: ", BATCH_SIZE) num_val = len(validatioin_imgs_path) num_imgs = len(original_imgs_path) # num_imgs = 100 original_imgs_path = original_imgs_path[:num_imgs] mod = num_imgs % BATCH_SIZE print('Train images number %d.\n' % num_imgs) print('Train images samples %s.\n' % str(num_imgs / BATCH_SIZE)) if mod > 0: print('Train set has been trimmed %d samples...\n' % mod) original_imgs_path = original_imgs_path[:-mod] # get the traing image shape HEIGHT, WIDTH, CHANNELS = TRAINING_IMAGE_SHAPE INPUT_SHAPE = (BATCH_SIZE, HEIGHT, WIDTH, CHANNELS) HEIGHT_OR, WIDTH_OR, CHANNELS_OR = TRAINING_IMAGE_SHAPE_OR INPUT_SHAPE_OR = (BATCH_SIZE, HEIGHT_OR, WIDTH_OR, CHANNELS_OR) # create the graph with tf.Graph().as_default(), tf.Session() as sess: original = tf.placeholder(tf.float32, shape=INPUT_SHAPE_OR, name='original') source = original print('source :', source.shape) print('original:', original.shape) # create the deepfuse net (encoder and decoder) dfn = DenseFuseNet(model_pre_path) generated_img = dfn.transform_recons(source) print('generate:', generated_img.shape) ssim_loss_value = SSIM_LOSS(original, generated_img) pixel_loss = tf.reduce_sum(tf.square(original - generated_img)) pixel_loss = pixel_loss / (BATCH_SIZE * HEIGHT * WIDTH) ssim_loss = 1 - ssim_loss_value loss = ssim_weight * ssim_loss + pixel_loss train_op = tf.train.AdamOptimizer(LEARNING_RATE).minimize(loss) sess.run(tf.global_variables_initializer()) # saver = tf.train.Saver() saver = tf.train.Saver(keep_checkpoint_every_n_hours=1) # ** Start Training ** step = 0 count_loss = 0 n_batches = int(len(original_imgs_path) // BATCH_SIZE) val_batches = int(len(validatioin_imgs_path) // BATCH_SIZE) if debug: elapsed_time = datetime.now() - start_time print( '\nElapsed time for preprocessing before actually train the model: %s' % elapsed_time) print('Now begin to train the model...\n') start_time = datetime.now() Loss_all = [i for i in range(EPOCHS * n_batches)] Loss_ssim = [i for i in range(EPOCHS * n_batches)] Loss_pixel = [i for i in range(EPOCHS * n_batches)] Val_ssim_data = [i for i in range(EPOCHS * n_batches)] Val_pixel_data = [i for i in range(EPOCHS * n_batches)] for epoch in range(EPOCHS): np.random.shuffle(original_imgs_path) for batch in range(n_batches): # retrive a batch of content and style images original_path = original_imgs_path[batch * BATCH_SIZE:( batch * BATCH_SIZE + BATCH_SIZE)] original_batch = get_train_images(original_path, crop_height=HEIGHT, crop_width=WIDTH, flag=False) original_batch = original_batch.reshape( [BATCH_SIZE, 256, 256, 1]) # print('original_batch shape final:', original_batch.shape) # run the training step sess.run(train_op, feed_dict={original: original_batch}) step += 1 if debug: is_last_step = (epoch == EPOCHS - 1) and (batch == n_batches - 1) if is_last_step or step % logging_period == 0: elapsed_time = datetime.now() - start_time _ssim_loss, _loss, _p_loss = sess.run( [ssim_loss, loss, pixel_loss], feed_dict={original: original_batch}) Loss_all[count_loss] = _loss Loss_ssim[count_loss] = _ssim_loss Loss_pixel[count_loss] = _p_loss print( 'epoch: %d/%d, step: %d, total loss: %s, elapsed time: %s' % (epoch, EPOCHS, step, _loss, elapsed_time)) print('p_loss: %s, ssim_loss: %s ,w_ssim_loss: %s ' % (_p_loss, _ssim_loss, ssim_weight * _ssim_loss)) # calculate the accuracy rate for 1000 images, every 100 steps val_ssim_acc = 0 val_pixel_acc = 0 np.random.shuffle(validatioin_imgs_path) val_start_time = datetime.now() for v in range(val_batches): val_original_path = validatioin_imgs_path[ v * BATCH_SIZE:(v * BATCH_SIZE + BATCH_SIZE)] val_original_batch = get_train_images( val_original_path, crop_height=HEIGHT, crop_width=WIDTH, flag=False) val_original_batch = val_original_batch.reshape( [BATCH_SIZE, 256, 256, 1]) val_ssim, val_pixel = sess.run( [ssim_loss, pixel_loss], feed_dict={original: val_original_batch}) val_ssim_acc = val_ssim_acc + (1 - val_ssim) val_pixel_acc = val_pixel_acc + val_pixel Val_ssim_data[count_loss] = val_ssim_acc / val_batches Val_pixel_data[ count_loss] = val_pixel_acc / val_batches val_es_time = datetime.now() - val_start_time print( 'validation value, SSIM: %s, Pixel: %s, elapsed time: %s' % (val_ssim_acc / val_batches, val_pixel_acc / val_batches, val_es_time)) print( '------------------------------------------------------------------------------' ) count_loss += 1 # ** Done Training & Save the model ** saver.save(sess, save_path) loss_data = Loss_all[:count_loss] scio.savemat( './models/loss/DeepDenseLossData' + str(ssim_weight) + '.mat', {'loss': loss_data}) loss_ssim_data = Loss_ssim[:count_loss] scio.savemat( './models/loss/DeepDenseLossSSIMData' + str(ssim_weight) + '.mat', {'loss_ssim': loss_ssim_data}) loss_pixel_data = Loss_pixel[:count_loss] scio.savemat( './models/loss/DeepDenseLossPixelData.mat' + str(ssim_weight) + '', {'loss_pixel': loss_pixel_data}) validation_ssim_data = Val_ssim_data[:count_loss] scio.savemat( './models/val/Validation_ssim_Data.mat' + str(ssim_weight) + '', {'val_ssim': validation_ssim_data}) validation_pixel_data = Val_pixel_data[:count_loss] scio.savemat( './models/val/Validation_pixel_Data.mat' + str(ssim_weight) + '', {'val_pixel': validation_pixel_data}) if debug: elapsed_time = datetime.now() - start_time print('Done training! Elapsed time: %s' % elapsed_time) print('Model is saved to: %s' % save_path)
def _handler_mix(ir_path, vis_path, model_path, model_pre_path, ssim_weight, index, output_path=None): mix_block = [] ir_img = get_train_images(ir_path, flag=False) vis_img = get_train_images(vis_path, flag=False) dimension = ir_img.shape ir_img = ir_img.reshape([1, dimension[0], dimension[1], dimension[2]]) vis_img = vis_img.reshape([1, dimension[0], dimension[1], dimension[2]]) ir_img = np.transpose(ir_img, (0, 2, 1, 3)) vis_img = np.transpose(vis_img, (0, 2, 1, 3)) print('img shape final:', ir_img.shape) with tf.Graph().as_default(), tf.Session() as sess: infrared_field = tf.placeholder(tf.float32, shape=ir_img.shape, name='content') visible_field = tf.placeholder(tf.float32, shape=vis_img.shape, name='style') # ----------------------------------------------- dfn = DenseFuseNet(model_pre_path) #sess.run(tf.global_variables_initializer()) enc_ir, enc_ir_res_block, enc_ir_block, enc_ir_block2 = dfn.transform_encoder( infrared_field) enc_vis, enc_vis_res_block, enc_vis_block, enc_vis_block2 = dfn.transform_encoder( visible_field) result = tf.placeholder(tf.float32, shape=enc_ir.shape, name='target') saver = tf.train.Saver() saver.restore(sess, model_path) enc_ir_temp, enc_ir_res_block_temp, enc_ir_block_temp, enc_ir_block2_temp = sess.run( [enc_ir, enc_ir_res_block, enc_ir_block, enc_ir_block2], feed_dict={infrared_field: ir_img}) enc_vis_temp, enc_vis_res_block_temp, enc_vis_block_temp, enc_vis_block2_temp = sess.run( [enc_vis, enc_vis_res_block, enc_vis_block, enc_vis_block2], feed_dict={visible_field: vis_img}) block = L1_norm(enc_ir_block_temp, enc_vis_block_temp) block2 = L1_norm(enc_ir_block2_temp, enc_vis_block2_temp) first_first = L1_norm(enc_ir_res_block_temp[0], enc_vis_res_block_temp[0]) first_second = Strategy(enc_ir_res_block_temp[1], enc_vis_res_block_temp[1]) #first_third = L1_norm_attention(enc_ir_res_block_temp[2],feation_ir, enc_vis_res_block_temp[2],feation_vis) #first_four = L1_norm_attention(enc_ir_res_block_temp[3],feation_ir, enc_vis_res_block_temp[3],feation_vis) first_third = L1_norm(enc_ir_res_block_temp[2], enc_vis_res_block_temp[2]) first_four = Strategy(enc_ir_res_block_temp[3], enc_vis_res_block_temp[3]) first_first = tf.concat( [first_first, tf.to_int32(first_second, name='ToInt')], 3) first_first = tf.concat( [first_first, tf.to_int32(first_third, name='ToInt')], 3) first_first = tf.concat([first_first, first_four], 3) first = first_first second = L1_norm(enc_ir_res_block_temp[6], enc_vis_res_block_temp[6]) third = L1_norm(enc_ir_res_block_temp[9], enc_vis_res_block_temp[9]) feature = 1 * first + 0.1 * second + 0.1 * third #--------------------------------------------------------- # block=Strategy(enc_ir_block_temp,enc_vis_block_temp) # block2=L1_norm(enc_ir_block2_temp,enc_vis_block2_temp) #--------------------------------------------------------- feature = feature.eval() output_image = dfn.transform_decoder(result, block, block2) # output = dfn.transform_decoder(feature) # print(type(feature)) # output = sess.run(output_image, feed_dict={result: feature,enc_res_block:block,enc_res_block2:block2}) output = sess.run(output_image, feed_dict={result: feature}) save_images(ir_path, output, output_path, prefix='fused' + str(index), suffix='_mix_' + str(ssim_weight))
def _handler_mix_a(ir_path, vis_path, model_path, model_pre_path, model_path_a, model_pre_path_a, ssim_weight, index, output_path=None): ir_img = get_train_images(ir_path, flag=False) vis_img = get_train_images(vis_path, flag=False) dimension = ir_img.shape ir_img = ir_img.reshape([1, dimension[0], dimension[1], dimension[2]]) vis_img = vis_img.reshape([1, dimension[0], dimension[1], dimension[2]]) ir_img = np.transpose(ir_img, (0, 2, 1, 3)) vis_img = np.transpose(vis_img, (0, 2, 1, 3)) g2 = tf.Graph() # 加载到Session 2的graph sess2 = tf.Session(graph=g2) # Session2 with sess2.as_default(): # 1 with g2.as_default(), tf.Session() as sess: infrared_field = tf.placeholder(tf.float32, shape=ir_img.shape, name='content') visible_field = tf.placeholder(tf.float32, shape=vis_img.shape, name='style') dfn = DenseFuseNet(model_pre_path) # sess.run(tf.global_variables_initializer()) enc_ir, enc_ir_res_block, enc_ir_block, enc_ir_block2 = dfn.transform_encoder( infrared_field) enc_vis, enc_vis_res_block, enc_vis_block, enc_vis_block2 = dfn.transform_encoder( visible_field) result = tf.placeholder(tf.float32, shape=enc_ir.shape, name='target') saver = tf.train.Saver() saver.restore(sess, model_path) # ------------------------attention------------------------------------------------------ #feature_a,feature_b=_get_attention(ir_path,vis_path,model_path_a,model_pre_path_a) #print("______+++________") #print(feature_a[0].shape) # ------------------------attention------------------------------------------------------ enc_ir_temp, enc_ir_res_block_temp, enc_ir_block_temp, enc_ir_block2_temp = sess.run( [enc_ir, enc_ir_res_block, enc_ir_block, enc_ir_block2], feed_dict={infrared_field: ir_img}) enc_vis_temp, enc_vis_res_block_temp, enc_vis_block_temp, enc_vis_block2_temp = sess.run( [enc_vis, enc_vis_res_block, enc_vis_block, enc_vis_block2], feed_dict={visible_field: vis_img}) # ------------------------------------------------------------------------------------------------------------ #------------------------------------------------------------------------------------------------------------ block = 0.8 * enc_vis_block_temp + 0.2 * enc_ir_block_temp block2 = 0.4 * enc_ir_block2_temp + 0.6 * enc_vis_block2_temp #first_first = Strategy(enc_ir_res_block_temp[0], enc_vis_res_block_temp[0]) #first_first = L1_norm(enc_ir_res_block_temp[0], enc_vis_res_block_temp[0]) #first_second = Strategy(enc_ir_res_block_temp[1], enc_vis_res_block_temp[1]) #first_second = L1_norm(enc_ir_res_block_temp[1], enc_vis_res_block_temp[1]) #first_third = Strategy(enc_ir_res_block_temp[2], enc_vis_res_block_temp[2]) #first_third = L1_norm_attention(enc_ir_res_block_temp[2],feature_a, enc_vis_res_block_temp[2],feature_b) #first_four = Strategy(enc_ir_res_block_temp[3], enc_vis_res_block_temp[3]) #first_four = L1_norm_attention(enc_ir_res_block_temp[3],feature_a, enc_vis_res_block_temp[3],feature_b) #first_first = tf.concat([first_first, tf.to_int32(first_second, name='ToInt')], 3) #first_first = tf.concat([first_first, tf.to_int32(first_third, name='ToInt')], 3) #first_first = tf.concat([first_first, first_four], 3) #first = first_first first = Strategy(enc_ir_res_block_temp[3], enc_vis_res_block_temp[3]) second = Strategy(enc_ir_res_block_temp[6], enc_vis_res_block_temp[6]) third = Strategy(enc_ir_res_block_temp[9], enc_vis_res_block_temp[9]) # ------------------------------------------------------------------------------------------------------------ # ------------------------------------------------------------------------------------------------------------ feature = 1 * first + 1 * second + 1 * third # --------------------------------------------------------- # block=Strategy(enc_ir_block_temp,enc_vis_block_temp) # block2=L1_norm(enc_ir_block2_temp,enc_vis_block2_temp) # --------------------------------------------------------- #feature = feature.eval() # --------------将特征图压成单通道---------------------------------- #feature_map_vis_out = sess.run(tf.reduce_sum(feature_a[0], 3, keep_dims=True)) #feature_map_ir_out = sess.run(tf.reduce_sum(feature_b[0],3, keep_dims=True)) # ------------------------------------------------------------------ output_image = dfn.transform_decoder(result, block, block2) # output = dfn.transform_decoder(feature) # print(type(feature)) # output = sess.run(output_image, feed_dict={result: feature,enc_res_block:block,enc_res_block2:block2}) output = sess.run(output_image, feed_dict={result: feature}) save_images(ir_path, output, output_path, prefix='fused' + str(index), suffix='_mix_' + str(ssim_weight))
print('\nElapsed time for preprocessing before actually train the model: %s' % elapsed_time) print('Now begin to train the model...\n') start_time = datetime.now() for epoch in range(EPOCHS): np.random.shuffle(content_images) np.random.shuffle(style_images) for batch in range(n_batches): # retrive a batch of content and style images content_batch_path = content_images[batch * BATCH_SIZE:(batch * BATCH_SIZE + BATCH_SIZE)] style_batch_path = style_images[batch * BATCH_SIZE:(batch * BATCH_SIZE + BATCH_SIZE)] content_batch = utils.get_train_images(content_batch_path, crop_height=HEIGHT, crop_width=WIDTH) style_batch = utils.get_train_images(style_batch_path, crop_height=HEIGHT, crop_width=WIDTH) # run the training step sess.run(train_op, feed_dict={content_input: content_batch, style_input: style_batch}) if step % 100 == 0: _content_loss, _style_loss, _loss = sess.run([content_loss, style_loss, loss], feed_dict={content_input: content_batch, style_input: style_batch}) elapsed_time = datetime.now() - start_time print('step: %d, total loss: %.3f, elapsed time: %s' % (step, _loss,elapsed_time)) print('content loss: %.3f' % (_content_loss)) print('style loss : %.3f, weighted style loss: %.3f\n' % (
def train_recons(original_imgs_path, validatioin_imgs_path, save_path, model_pre_path, ssim_weight, EPOCHES_set, BATCH_SIZE, debug=False, logging_period=1): if debug: from datetime import datetime start_time = datetime.now() EPOCHS = EPOCHES_set print("EPOCHES : ", EPOCHS) #EPOCHS = 4 遍历整个数据集的次数,训练网络一共要执行n*4次 print("BATCH_SIZE: ", BATCH_SIZE) #BATCH_SIZE = 2 每个Batch有2个样本,共n/2个Batch,每处理两个样本模型权重就更新 num_val = len(validatioin_imgs_path) #测试集样本个数 num_imgs = len(original_imgs_path) #训练集样本个数 # num_imgs = 100 original_imgs_path = original_imgs_path[:num_imgs] #迷惑行为,自己赋给自己 mod = num_imgs % BATCH_SIZE #Batch个数 print('Train images number %d.\n' % num_imgs) print('Train images samples %s.\n' % str(num_imgs / BATCH_SIZE)) if mod > 0: print('Train set has been trimmed %d samples...\n' % mod) original_imgs_path = original_imgs_path[:-mod] #original_imags_path 数组移除最后两个 # get the traing image shape #训练图像的长宽及通道数 255,255,1 HEIGHT, WIDTH, CHANNELS = TRAINING_IMAGE_SHAPE INPUT_SHAPE = (BATCH_SIZE, HEIGHT, WIDTH, CHANNELS) #定义元组,意义不明 HEIGHT_OR, WIDTH_OR, CHANNELS_OR = TRAINING_IMAGE_SHAPE_OR INPUT_SHAPE_OR = (BATCH_SIZE, HEIGHT_OR, WIDTH_OR, CHANNELS_OR) #OR是什么意思,意义不明 # create the graph with tf.Graph().as_default(), tf.Session() as sess: original = tf.placeholder(tf.float32, shape=INPUT_SHAPE_OR, name='original') #神经网络构建graph的时候在模型中的占位,只分配必要的内存,运行模型时通过feed_dict()向占位符喂入数据 #第一个参数,数据类型,常用tf.float32,tf.float64 #第二个参数,数据形状,矩阵形状,图像的长宽及通道数 #第三个参数,名称 #返回Tensor类型 source = original #迷惑行为,意义不明 print('source :', source.shape) print('original:', original.shape) # create the deepfuse net (encoder and decoder) #创建深度学习网络 dfn = DenseFuseNet(model_pre_path) #这里的model_pre_path是自己设置的模型参数,默认是None,若不为None则起始训练的参数为设置的文件 generated_img = dfn.transform_recons(source) #输出图像 print('generate:', generated_img.shape) ######################################################################################### # COST FUNCTION 部分 ssim_loss_value = SSIM_LOSS(original, generated_img) #计算SSIM pixel_loss = tf.reduce_sum(tf.square(original - generated_img)) pixel_loss = pixel_loss/(BATCH_SIZE*HEIGHT*WIDTH) #计算pixel loss ssim_loss = 1 - ssim_loss_value #SSIM loss数值 loss = ssim_weight*ssim_loss + pixel_loss #整体loss #train_op = tf.train.AdamOptimizer(LEARNING_RATE).minimize(loss) #自适应矩估计(梯度下降的一种方法) train_op = tf.train.AdamOptimizer(LEARNING_RATE_2).minimize(loss) # 自适应矩估计(梯度下降的一种方法) ########################################################################################## sess.run(tf.global_variables_initializer()) # saver = tf.train.Saver() saver = tf.train.Saver(keep_checkpoint_every_n_hours=1) # ** Start Training ** step = 0 count_loss = 0 n_batches = int(len(original_imgs_path) // BATCH_SIZE) val_batches = int(len(validatioin_imgs_path) // BATCH_SIZE) if debug: elapsed_time = datetime.now() - start_time print('\nElapsed time for preprocessing before actually train the model: %s' % elapsed_time) print('Now begin to train the model...\n') start_time = datetime.now() Loss_all = [i for i in range(EPOCHS * n_batches)] Loss_ssim = [i for i in range(EPOCHS * n_batches)] Loss_pixel = [i for i in range(EPOCHS * n_batches)] Val_ssim_data = [i for i in range(EPOCHS * n_batches)] Val_pixel_data = [i for i in range(EPOCHS * n_batches)] for epoch in range(EPOCHS): np.random.shuffle(original_imgs_path) for batch in range(n_batches): # retrive a batch of content and style images original_path = original_imgs_path[batch*BATCH_SIZE:(batch*BATCH_SIZE + BATCH_SIZE)] original_batch = get_train_images(original_path, crop_height=HEIGHT, crop_width=WIDTH, flag=False) original_batch = original_batch.reshape([BATCH_SIZE, 256, 256, 1]) # print('original_batch shape final:', original_batch.shape) # run the training step sess.run(train_op, feed_dict={original: original_batch}) step += 1 if debug: is_last_step = (epoch == EPOCHS - 1) and (batch == n_batches - 1) if is_last_step or step % logging_period == 0: elapsed_time = datetime.now() - start_time _ssim_loss, _loss, _p_loss = sess.run([ssim_loss, loss, pixel_loss], feed_dict={original: original_batch}) Loss_all[count_loss] = _loss Loss_ssim[count_loss] = _ssim_loss Loss_pixel[count_loss] = _p_loss print('epoch: %d/%d, step: %d, total loss: %s, elapsed time: %s' % (epoch, EPOCHS, step, _loss, elapsed_time)) print('p_loss: %s, ssim_loss: %s ,w_ssim_loss: %s ' % (_p_loss, _ssim_loss, ssim_weight * _ssim_loss)) # calculate the accuracy rate for 1000 images, every 100 steps val_ssim_acc = 0 val_pixel_acc = 0 np.random.shuffle(validatioin_imgs_path) val_start_time = datetime.now() for v in range(val_batches): val_original_path = validatioin_imgs_path[v * BATCH_SIZE:(v * BATCH_SIZE + BATCH_SIZE)] val_original_batch = get_train_images(val_original_path, crop_height=HEIGHT, crop_width=WIDTH,flag=False) val_original_batch = val_original_batch.reshape([BATCH_SIZE, 256, 256, 1]) val_ssim, val_pixel = sess.run([ssim_loss, pixel_loss], feed_dict={original: val_original_batch}) val_ssim_acc = val_ssim_acc + (1 - val_ssim) val_pixel_acc = val_pixel_acc + val_pixel Val_ssim_data[count_loss] = val_ssim_acc/val_batches Val_pixel_data[count_loss] = val_pixel_acc / val_batches val_es_time = datetime.now() - val_start_time print('validation value, SSIM: %s, Pixel: %s, elapsed time: %s' % (val_ssim_acc/val_batches, val_pixel_acc / val_batches, val_es_time)) print('------------------------------------------------------------------------------') count_loss += 1 # ** Done Training & Save the model ** saver.save(sess, save_path) #---------------------------------------------------------------------------------------------------------------- loss_data = Loss_all[:count_loss] scio.savemat('/data/ljy/1-Project-Go/01-06-upsampling/models/loss/DeepDenseLossData' + str(ssim_weight) + '.mat', {'loss': loss_data}) loss_ssim_data = Loss_ssim[:count_loss] scio.savemat('/data/ljy/1-Project-Go/01-06-upsampling/models/loss/DeepDenseLossSSIMData' + str( ssim_weight) + '.mat', {'loss_ssim': loss_ssim_data}) loss_pixel_data = Loss_pixel[:count_loss] scio.savemat('/data/ljy/1-Project-Go/01-06-upsampling/models/loss/DeepDenseLossPixelData.mat' + str( ssim_weight) + '', {'loss_pixel': loss_pixel_data}) validation_ssim_data = Val_ssim_data[:count_loss] scio.savemat('/data/ljy/1-Project-Go/01-06-upsampling/models/val/Validation_ssim_Data.mat' + str( ssim_weight) + '', {'val_ssim': validation_ssim_data}) validation_pixel_data = Val_pixel_data[:count_loss] scio.savemat('/data/ljy/1-Project-Go/01-06-upsampling/models/val/Validation_pixel_Data.mat' + str( ssim_weight) + '', {'val_pixel': validation_pixel_data}) #---------------------------------------------------------------------------------------------------- if debug: elapsed_time = datetime.now() - start_time print('Done training! Elapsed time: %s' % elapsed_time) print('Model is saved to: %s' % save_path)
def train(opt): ''' Input: opt : optins for training model ''' content_img_list = utils.list_images(opt.content_img_dir) style_img_list = utils.list_images(opt.style_img_dir) #import pdb; pdb.set_trace() assert(content_img_list) # ensure not empty assert(style_img_list) with tf.Graph().as_default(), tf.Session() as sess: content_img = tf.placeholder(tf.float32, shape=(opt.batch_size, opt.img_size, opt.img_size, 3), name='content_img') style_img = tf.placeholder(tf.float32, shape=(opt.batch_size, opt.img_size, opt.img_size, 3), name='style_img') model = Model(opt.checkpoint_encoder) # Encode Image generated_img = model.transform(content_img, style_img) generated_img = tf.reverse(generated_img, axis=[-1]) generated_img = model.encoder.preprocess(generated_img) enc_gen, enc_gen_layers = model.encoder.encode(generated_img) target_features = model.target_features # Content Loss content_loss = tf.reduce_sum(tf.reduce_mean(tf.square(enc_gen - target_features), axis=[1, 2])) # Style Loss style_layer_loss = [] style_layers = ['relu1_1', 'relu2_1', 'relu3_1', 'relu4_1'] for layer in style_layers: enc_style_feat = model.encoded_style_layers[layer] enc_gen_feat = enc_gen_layers[layer] meanS, varS = tf.nn.moments(enc_style_feat, [1, 2]) meanG, varG = tf.nn.moments(enc_gen_feat, [1, 2]) sigmaS = tf.sqrt(varS + opt.epsilon) sigmaG = tf.sqrt(varG + opt.epsilon) l2_mean = tf.reduce_sum(tf.square(meanG - meanS)) l2_sigma = tf.reduce_sum(tf.square(sigmaG - sigmaS)) style_layer_loss.append(l2_mean + l2_sigma) style_loss = tf.reduce_sum(style_layer_loss) # Total loss loss = opt.content_weight * content_loss + opt.style_weight * style_loss # Train global_step = tf.Variable(0, trainable=False) learning_rate = tf.train.inverse_time_decay(opt.lr, global_step, opt.lr_decay_step, opt.lr_decay) train_op = tf.train.AdamOptimizer(learning_rate).minimize(loss, global_step=global_step) sess.run(tf.global_variables_initializer()) # saver saver = tf.train.Saver(max_to_keep=10) step = 0 n_batches = int(len(content_img_list) // opt.batch_size) try: for epoch in range(opt.epoch): np.random.shuffle(content_img_list) np.random.shuffle(style_img_list) for batch in range(n_batches): print('iteration {}'.format(step)) # retrive a batch of content and style images content_batch_path = content_img_list[batch*opt.batch_size:(batch*opt.batch_size + opt.batch_size)] style_batch_path = style_img_list[batch*opt.batch_size:(batch*opt.batch_size + opt.batch_size)] content_batch = utils.get_train_images(content_batch_path, crop_height=opt.img_size, crop_width=opt.img_size) style_batch = utils.get_train_images(style_batch_path, crop_height=opt.img_size, crop_width=opt.img_size) # run the training step sess.run(train_op, feed_dict={ content_img: content_batch, style_img: style_batch }) step += 1 if step % 1000 == 0: saver.save(sess, opt.checkpoint_save_dir, global_step=step, write_meta_graph=False) except Exception as ex: saver.save(sess, opt.checkpoint_save_dir, global_step=step) print('Error message: %s' % str(ex)) # Finish saver.save(sess, opt.checkpoint_save_dir)