def _handler_wt(ir_path, vi_path, model_path, model_pre_path, index, output_path=None): ir_img = read_test_image(ir_path) vi_img = read_test_image(vi_path) dimension = ir_img.shape ir_img = ir_img.reshape([1, dimension[0], dimension[1], dimension[2]]) vi_img = vi_img.reshape([1, dimension[0], dimension[1], dimension[2]]) print('img shape final:', ir_img.shape) with tf.compat.v1.Graph().as_default(), tf.compat.v1.Session() as sess: infrared_field = tf.placeholder(tf.float32, shape=ir_img.shape, name='image_ir') visible_field = tf.placeholder(tf.float32, shape=ir_img.shape, name='image_vi') dfn = DenseFuseNet(model_pre_path) # restore the trained model and run the style transferring saver = tf.train.Saver() saver.restore(sess, model_path) output_image = dfn.transform_wt(infrared_field, visible_field) output = sess.run(output_image, feed_dict={infrared_field: ir_img, visible_field: vi_img}) save_images(ir_path, output, output_path, prefix='Dense_fuse_' + str(index).zfill(2))
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 _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 _handler_rgb_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) 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)) ir_img1 = ir_img[:, :, :, 0] ir_img1 = ir_img1.reshape([1, dimension[0], dimension[1], 1]) ir_img2 = ir_img[:, :, :, 1] ir_img2 = ir_img2.reshape([1, dimension[0], dimension[1], 1]) ir_img3 = ir_img[:, :, :, 2] ir_img3 = ir_img3.reshape([1, dimension[0], dimension[1], 1]) vis_img1 = vis_img[:, :, :, 0] vis_img1 = vis_img1.reshape([1, dimension[0], dimension[1], 1]) vis_img2 = vis_img[:, :, :, 1] vis_img2 = vis_img2.reshape([1, dimension[0], dimension[1], 1]) vis_img3 = vis_img[:, :, :, 2] vis_img3 = vis_img3.reshape([1, dimension[0], dimension[1], 1]) print('img shape final:', ir_img1.shape) with tf.Graph().as_default(), tf.Session() as sess: infrared_field = tf.placeholder( tf.float32, shape=ir_img1.shape, name='content') visible_field = tf.placeholder( tf.float32, shape=ir_img1.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_img1, visible_field: vis_img1}) feature = L1_norm(enc_ir_temp, enc_vis_temp) output1 = sess.run(output_image, feed_dict={target: feature}) enc_ir_temp, enc_vis_temp = sess.run([enc_ir, enc_vis], feed_dict={infrared_field: ir_img2, visible_field: vis_img2}) feature = L1_norm(enc_ir_temp, enc_vis_temp) output2 = sess.run(output_image, feed_dict={target: feature}) enc_ir_temp, enc_vis_temp = sess.run([enc_ir, enc_vis], feed_dict={infrared_field: ir_img3, visible_field: vis_img3}) feature = L1_norm(enc_ir_temp, enc_vis_temp) output3 = sess.run(output_image, feed_dict={target: feature}) output1 = output1.reshape([1, dimension[0], dimension[1]]) output2 = output2.reshape([1, dimension[0], dimension[1]]) output3 = output3.reshape([1, dimension[0], dimension[1]]) output = np.stack((output1, output2, output3), axis=-1) output = np.transpose(output, (0, 2, 1, 3)) save_images(ir_path, output, output_path, prefix='fused' + str(index), suffix='_densefuse_l1norm_'+str(ssim_weight))
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 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()
# -*- coding: utf-8 -*- """ Created on Sat Jun 15 17:32:38 2019 @author: win10 """ import torch from densefuse_net import DenseFuseNet from utils import test device = 'cuda' model = DenseFuseNet().to(device) model.load_state_dict(torch.load('./train_result/model_weight.pkl')['weight']) test_path = './images/IV_images/' test(test_path, model, mode='add')
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 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 _handler_rgb_patch_based(images_path, model_path, model_pre_path, index, output_path=None): size = len(images_path) images = ["" for x in range(size)] ir_img1 = ["" for x in range(size)] ir_img2 = ["" for x in range(size)] ir_img3 = ["" for x in range(size)] for x in range(0, size): images[x] = get_train_images_rgb(images_path[x], flag=False) dimension = images[x].shape images[x] = images[x].reshape( [1, dimension[0], dimension[1], dimension[2]]) images[x] = np.transpose(images[x], (0, 2, 1, 3)) ir_img1[x] = images[x][:, :, :, 0] ir_img1[x] = ir_img1[x].reshape([1, dimension[0], dimension[1], 1]) ir_img2[x] = images[x][:, :, :, 1] ir_img2[x] = ir_img2[x].reshape([1, dimension[0], dimension[1], 1]) ir_img3[x] = images[x][:, :, :, 2] ir_img3[x] = ir_img3[x].reshape([1, dimension[0], dimension[1], 1]) print('img shape final:', ir_img1[0].shape) dimension = images[0].shape TRAIN_TAIL_SIZE_X = 32 TRAIN_TAIL_SIZE_Y = 32 TILES_X = (int)(dimension[1] / TRAIN_TAIL_SIZE_X) TILES_Y = (int)(dimension[2] / TRAIN_TAIL_SIZE_Y) INPUT_SHAPE_TILE = (1, TRAIN_TAIL_SIZE_X, TRAIN_TAIL_SIZE_Y, 1) with tf.Graph().as_default(), tf.Session() as sess: images_field = ["" for x in range(size)] for x in range(0, size): images_field[x] = tf.placeholder(tf.float32, shape=INPUT_SHAPE_TILE) dfn = DenseFuseNet(model_pre_path) output_image = dfn.transform_addition(images_field) # restore the trained model and run the style transferring saver = tf.train.Saver() saver.restore(sess, model_path) dimension = images[0].shape output1 = np.zeros([1, dimension[1], dimension[2]]) output2 = np.zeros([1, dimension[1], dimension[2]]) output3 = np.zeros([1, dimension[1], dimension[2]]) for tile_x in range(TILES_X): for tile_y in range(TILES_Y): x1 = tile_x * TRAIN_TAIL_SIZE_X y1 = tile_y * TRAIN_TAIL_SIZE_Y x2 = x1 + TRAIN_TAIL_SIZE_X y2 = y1 + TRAIN_TAIL_SIZE_Y tile_img1 = ["" for x in range(size)] tile_img2 = ["" for x in range(size)] tile_img3 = ["" for x in range(size)] for x in range(0, size): tile_img1[x] = ir_img1[x][:, x1:x2, y1:y2, :] tile_img2[x] = ir_img2[x][:, x1:x2, y1:y2, :] tile_img3[x] = ir_img3[x][:, x1:x2, y1:y2, :] output1_t = sess.run( output_image, feed_dict={i: d for i, d in zip(images_field, tile_img1)}) output2_t = sess.run( output_image, feed_dict={i: d for i, d in zip(images_field, tile_img2)}) output3_t = sess.run( output_image, feed_dict={i: d for i, d in zip(images_field, tile_img3)}) output1_t = output1_t.reshape( [1, TRAIN_TAIL_SIZE_X, TRAIN_TAIL_SIZE_Y]) output2_t = output2_t.reshape( [1, TRAIN_TAIL_SIZE_X, TRAIN_TAIL_SIZE_Y]) output3_t = output3_t.reshape( [1, TRAIN_TAIL_SIZE_X, TRAIN_TAIL_SIZE_Y]) output1[0, x1:x2, y1:y2] = output1_t output2[0, x1:x2, y1:y2] = output2_t output3[0, x1:x2, y1:y2] = output3_t output = np.stack((output1, output2, output3), axis=-1) output = np.transpose(output, (0, 2, 1, 3)) save_images(images_path, output, output_path, prefix='fused' + str(index), suffix='_densefuse_addition_patch_based')
def _handler_rgb_l1(images_path, model_path, model_pre_path, index, output_path=None): size = len(images_path) images = ["" for x in range(size)] ir_img1 = ["" for x in range(size)] ir_img2 = ["" for x in range(size)] ir_img3 = ["" for x in range(size)] for x in range(0, size): images[x] = get_train_images_rgb(images_path[x], flag=False) dimension = images[x].shape images[x] = images[x].reshape( [1, dimension[0], dimension[1], dimension[2]]) images[x] = np.transpose(images[x], (0, 2, 1, 3)) ir_img1[x] = images[x][:, :, :, 0] ir_img1[x] = ir_img1[x].reshape([1, dimension[0], dimension[1], 1]) ir_img2[x] = images[x][:, :, :, 1] ir_img2[x] = ir_img2[x].reshape([1, dimension[0], dimension[1], 1]) ir_img3[x] = images[x][:, :, :, 2] ir_img3[x] = ir_img3[x].reshape([1, dimension[0], dimension[1], 1]) print('img shape final:', ir_img1[0].shape) with tf.Graph().as_default(), tf.Session() as sess: images_field = ["" for x in range(size)] for x in range(0, size): images_field[x] = tf.placeholder(tf.float32, shape=ir_img1[0].shape) dfn = DenseFuseNet(model_pre_path) enc_irs = ["" for x in range(size)] enc_irs = dfn.transform_encoder(images_field) target = tf.placeholder(tf.float32, shape=enc_irs[0].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_temps = sess.run( enc_irs, feed_dict={i: d for i, d in zip(images_field, ir_img1)}) feature = L1_norm(enc_ir_temps) output1 = sess.run(output_image, feed_dict={target: feature}) enc_ir_temps = sess.run( enc_irs, feed_dict={i: d for i, d in zip(images_field, ir_img2)}) feature = L1_norm(enc_ir_temps) output2 = sess.run(output_image, feed_dict={target: feature}) enc_ir_temps = sess.run( enc_irs, feed_dict={i: d for i, d in zip(images_field, ir_img3)}) feature = L1_norm(enc_ir_temps) output3 = sess.run(output_image, feed_dict={target: feature}) output1 = output1.reshape([1, dimension[0], dimension[1]]) output2 = output2.reshape([1, dimension[0], dimension[1]]) output3 = output3.reshape([1, dimension[0], dimension[1]]) output = np.stack((output1, output2, output3), axis=-1) output = np.transpose(output, (0, 2, 1, 3)) save_images(images_path, output, output_path, prefix='fused' + str(index), suffix='_densefuse_l1norm')
# data_vi = AEDataset(root_vi, resize= 128, gray = True) # loader_vi = DataLoader(data_vi, batch_size = batch_size, shuffle=True) data_TNO = AEDataset(root_ir_TNO, root_vi_TNO, resize=128, gray=True) loader_TNO = DataLoader(data_TNO, batch_size=batch_size, shuffle=True) img_num = (len(data_TNO) + len(data_k)) / 2 print("Load ---- {} pairs of images: KAIST:[{}] + TNO:[{}]".format( img_num, len(data_k) / 2, len(data_TNO) / 2)) # print("Load ---- {} pairs of images:TNO:[{}]".format(img_num,len(data_TNO))) # data_val = AEDataset(root_val, resize= [256,256], transform = None, gray = True) # loader_val = DataLoader(data_val, batch_size = 100, shuffle=True) # Model model = DenseFuseNet().to(device) # checkpoint = torch.load('./train_result/H_model_weight_new.pkl') # model.load_state_dict(checkpoint['weight']) print(model) optimizer = optim.Adam(model.parameters(), lr=lr) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=10, verbose=False, threshold=0.0001, threshold_mode='rel', cooldown=0, min_lr=0, eps=1e-10) # scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1, gamma=0.9, last_epoch=-1)
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))
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)
# -*- coding: utf-8 -*- from glob import glob from densefuse_net import DenseFuseNet import torch from PIL import Image import torchvision.transforms as transforms import time import os from channel_fusion import channel_f as channel_fusion from utils import mkdir, Strategy _tensor = transforms.ToTensor() _pil_gray = transforms.ToPILImage() os.environ["CUDA_VISIBLE_DEVICES"] = "2" device = 'cuda' model = DenseFuseNet().to(device) checkpoint = torch.load('./train_result/H_best.pkl') # checkpoint = torch.load('./train_result/model_weight_new.pkl') model.load_state_dict(checkpoint['weight']) mkdir("outputs/fea/") mkdir("outputs/fea/vi/") mkdir("outputs/fea/ir/") mkdir("result") test_ir = './Test_ir/' test_vi = './Test_vi/' def load_img(img_path, img_type='gray'): img = Image.open(img_path) if img_type == 'gray': img = img.convert('L')