def train():
parser = config_parser()
args = parser.parse_args()
setup_runtime(args)
log_config(args)
mlp = MLPRunner(args)
mlp.train()
def train():
parser = config_parser()
args = parser.parse_args()
setup_runtime(args)
log_config(args)
fcn = FCNRunner(args)
fcn.train()
def main():
component_map = {"pup": maps.PUP, "hbi": maps.HBI}
logger.info("Starting Component Tester")
config.log_config()
if config.ROLE == "consume":
data = component_map[config.COMPONENT]
consumer.consume(data["consume_topic"])
elif config.ROLE == "produce":
data = component_map[config.COMPONENT]
producer.produce(data["produce_topic"], data["msg"])
else:
logger.error("Role not recognized: %s", config.ROLE)
def main():
config.init_config()
print('go to model')
print '*' * 80
_log_file = open(
config.LOG_FILE_PRE +
time.strftime("_%Y_%m_%d_%H%M%S", time.localtime()), 'w')
# log configuration.
config.log_config(_log_file)
# initialize model
weights_path = None
if config.MODE == 2:
if config.DATASET == 'WSJ0':
weights_path = './_tmp_weights/ASAM_WSJ0_weight_00031.h5'
elif config.DATASET == 'THCHS-30':
weights_path = './_tmp_weights/ASAM_THCHS30_weight_00034.h5'
dl4ss_model = nnet.NNet(_log_file, weights_path)
if config.MODE == 1:
print 'Start to train model ...'
_log_file.write('Start to train model ...\n')
dl4ss_model.train()
print 'valid spk number: 2'
_log_file.write('valid spk number: 2\n')
dl4ss_model.predict(config.VALID_LIST, spk_num=2)
print 'test spk number: 2'
_log_file.write('test spk number: 2\n')
dl4ss_model.predict(config.TEST_LIST, spk_num=2)
print 'test spk number: 3'
_log_file.write('test spk number: 3\n')
dl4ss_model.predict(config.TEST_LIST, spk_num=3)
print 'test spk number: 2 with bg noise'
_log_file.write('test spk number: 2 with bg noise\n')
dl4ss_model.predict(config.TEST_LIST, spk_num=2, add_bgd_noise=True)
for supp_time in [0.25, 0.5, 1, 2, 4, 8, 16, 32]:
print 'unk spk and supplemental wav span: %02d' % supp_time
_log_file.write('unk spk and supplemental wav span: %02d\n' %
supp_time)
dl4ss_model.predict(config.UNK_LIST,
spk_num=2,
unk_spk=True,
supp_time=supp_time)
else:
print 'Wrong mode: %s' % config.MODE
_log_file.write('Wrong mode: %s\n' % config.MODE)
_log_file.close()
def main():
# ############## Test,测试区 ############
# ############## Test,测试区 ############
config.init_config()
print('go to model')
print '*' * 80
_log_file = open(
config.LOG_FILE_PRE +
time.strftime("_%Y_%m_%d_%H%M%S", time.localtime()), 'w')
# 记录参数
config.log_config(_log_file)
# 初始化网络 initialize model
weights_path = './_tmp_weights/WSJ0_weight_00021.h5'
weights_path = './_tmp_weights/WSJ0_weight_00150_forImages.h5'
# weights_path = None
dl4ss_model = nnet.NNet(_log_file, weights_path)
if config.MODE == 1:
print 'still train'
dl4ss_model.train()
elif (config.MODE == 2) and (weights_path is not None):
# 在训练好的模型上测试TEST
# print 'valid spk number: 2'
# _log_file.write('valid spk number: 2\n')
# dl4ss_model.predict(config.VALID_LIST, spk_num=2)
print 'test spk number: 2'
_log_file.write('test spk number: 2\n')
dl4ss_model.predict(config.TEST_LIST, spk_num=2)
# print 'tes spk number: 3'
# _log_file.write('test spk number: 3\n')
# dl4ss_model.predict(config.TEST_LIST, spk_num=3)
print 'test spk number: 2 with bg noise'
_log_file.write('test spk number: 2 with bg noise\n')
dl4ss_model.predict(config.TEST_LIST, spk_num=2, add_bgd_noise=True)
# 在训练好的模型上测试UNK
# for supp_time in [0.25, 0.5, 1, 2, 4, 8, 16, 32]:
# for supp_time in [0.25, 0.5, 1, 2, 4, 8, 16, 32]:
# print 'unk spk and supplemental wav span: %02d' % supp_time
# _log_file.write('unk spk and supplemental wav span: %02d\n' % supp_time)
# dl4ss_model.predict(config.UNK_LIST, spk_num=2, unk_spk=True, supp_time=supp_time)
# else:
# print 'Wrong mode: %s' % config.MODE
# _log_file.write('Wrong mode: %s\n' % config.MODE)
_log_file.close()
def main():
# load the external configuration
config.init_config()
print("Go to model")
print '*' * 80
_log_file = open(
config.log_file_pre +
time.strftime("_%Y_%m_%d_%H%M%S", time.localtime()), 'w')
# record current configuration
config.log_config(_log_file)
# initialize model
hmn4qa_model = model.Model(_log_file)
# train
hmn4qa_model.train(n_epochs=config.max_epoch,
shuffle_batch=config.shuffle_batch)
_log_file.close()
def main():
config.init_config()
print('go to model')
print '*' * 80
_log_file = open(config.LOG_FILE_PRE + time.strftime("_%Y_%m_%d_%H%M%S", time.localtime()), 'w')
# log configuration.
config.log_config(_log_file)
# initialize model
weights_path = None
if config.MODE == 2:
if config.DATASET == 'WSJ0':
weights_path = './_tmp_weights/ASAM_WSJ0_weight_00031.h5'
elif config.DATASET == 'THCHS-30':
weights_path = './_tmp_weights/ASAM_THCHS30_weight_00034.h5'
dl4ss_model = nnet.NNet(_log_file, weights_path)
if config.MODE == 1:
print 'Start to train model ...'
_log_file.write('Start to train model ...\n')
dl4ss_model.train()
print 'valid spk number: 2'
_log_file.write('valid spk number: 2\n')
dl4ss_model.predict(config.VALID_LIST, spk_num=2)
print 'test spk number: 2'
_log_file.write('test spk number: 2\n')
dl4ss_model.predict(config.TEST_LIST, spk_num=2)
print 'test spk number: 3'
_log_file.write('test spk number: 3\n')
dl4ss_model.predict(config.TEST_LIST, spk_num=3)
print 'test spk number: 2 with bg noise'
_log_file.write('test spk number: 2 with bg noise\n')
dl4ss_model.predict(config.TEST_LIST, spk_num=2, add_bgd_noise=True)
for supp_time in [0.25, 0.5, 1, 2, 4, 8, 16, 32]:
print 'unk spk and supplemental wav span: %02d' % supp_time
_log_file.write('unk spk and supplemental wav span: %02d\n' % supp_time)
dl4ss_model.predict(config.UNK_LIST, spk_num=2, unk_spk=True, supp_time=supp_time)
else:
print 'Wrong mode: %s' % config.MODE
_log_file.write('Wrong mode: %s\n' % config.MODE)
_log_file.close()
def config():
args = parse_args()
log_config(args, 'test')
return args
def train_with_synthetic():
checkpoint_dir ="test_checkpoint/{}".format(tl.global_flag['mode']) # checkpoint_resize_conv
tl.files.exists_or_mkdir(checkpoint_dir)
log_config(checkpoint_dir + '/config', config)
save_dir_sample = "samples/{}".format(tl.global_flag['mode'])
tl.files.exists_or_mkdir(save_dir_sample)
input_path = config.TRAIN.synthetic_blur_path
train_blur_img_list = sorted(tl.files.load_file_list(path=input_path, regx='(out_of_focus|motion).*.(jpg|JPG)', printable=False))
train_mask_img_list=[]
for str in train_blur_img_list:
if ".jpg" in str:
train_mask_img_list.append(str.replace(".jpg",".png"))
else:
train_mask_img_list.append(str.replace(".JPG", ".png"))
#augmented dataset read
gt_path =config.TRAIN.synthetic_gt_path
print train_mask_img_list
train_blur_imgs = read_all_imgs(train_blur_img_list, path=input_path, n_threads=100 ,mode='RGB')
train_mask_imgs = read_all_imgs(train_mask_img_list, path=gt_path, n_threads=100,mode='GRAY_cv')
index= 0
train_classification_mask= []
#print train_mask_imgs
#img_n = 0
for img in train_mask_imgs:
tmp_class = img
tmp_classification = np.concatenate((img,img,img),axis = 2)
tmp_class[np.where(tmp_classification[:,:,0]==0)] =0 #sharp
tmp_class[np.where(tmp_classification[:,:,0]==100)] =1 #motion blur
tmp_class[np.where(tmp_classification[:,:,0]==200)] =2 #defocus blur
train_classification_mask.append(tmp_class)
index =index +1
input_path2 = config.TRAIN.CUHK_blur_path
ori_train_blur_img_list = sorted(tl.files.load_file_list(path=input_path2, regx='(out_of_focus|motion).*.(jpg|JPG)', printable=False))
ori_train_mask_img_list=[]
for str in ori_train_blur_img_list:
if ".jpg" in str:
ori_train_mask_img_list.append(str.replace(".jpg",".png"))
else:
ori_train_mask_img_list.append(str.replace(".JPG", ".png"))
#augmented dataset read
gt_path2 = config.TRAIN.CUHK_gt_path
print train_blur_img_list
ori_train_blur_imgs = read_all_imgs(ori_train_blur_img_list, path=input_path2, n_threads=batch_size ,mode='RGB')
ori_train_mask_imgs = read_all_imgs(ori_train_mask_img_list, path=gt_path2, n_threads=batch_size,mode='GRAY')
train_edge_imgs = []
index= 0
ori_train_classification_mask= []
#img_n = 0
for img in ori_train_mask_imgs:
if(index<236):
tmp_class = img
tmp_classification = np.concatenate((img,img,img),axis = 2)
tmp_class[np.where(tmp_classification[:,:,0]==0)] =0 #sharp
tmp_class[np.where(tmp_classification[:,:,0]>0)] =1 #defocus blur
else:
tmp_class = img
tmp_classification = np.concatenate((img, img, img), axis=2)
tmp_class[np.where(tmp_classification[:,:,0]==0)] =0 #sharp
tmp_class[np.where(tmp_classification[:,:,0]>0)] =2 #defocus blur
ori_train_classification_mask.append(tmp_class)
index =index +1
train_mask_imgs= train_classification_mask
for i in range(10):
train_blur_imgs = train_blur_imgs + ori_train_blur_imgs;
train_mask_imgs = train_mask_imgs + ori_train_classification_mask;
print len(train_blur_imgs), len(train_mask_imgs)
### DEFINE MODEL ###
patches_blurred = tf.placeholder('float32', [batch_size, h, w, 3], name = 'input_patches')
labels_sigma = tf.placeholder('float32', [batch_size,h,w, 1], name = 'lables')
classification_map= tf.placeholder('int32', [batch_size, h, w,1], name='labels')
#class_map = tf.placeholder('int32', [batch_size, h, w], name='classes')
#attention_edge = tf.placeholder('float32', [batch_size, h, w, 1], name='attention')
with tf.variable_scope('Unified'):
with tf.variable_scope('VGG') as scope1:
n, f0, f0_1, f1_2, f2_3 ,hrg,wrg= VGG19_pretrained(patches_blurred,reuse=False, scope=scope1)
with tf.variable_scope('UNet') as scope2:
net_regression,m1,m2,m3= Decoder_Network_classification(n, f0, f0_1, f1_2, f2_3 ,hrg,wrg, reuse = False, scope = scope2)
### DEFINE LOSS ###
loss1 = tl.cost.cross_entropy((net_regression.outputs), tf.squeeze( classification_map), name='loss1')
loss2 = tl.cost.cross_entropy((m1.outputs), tf.squeeze( tf.image.resize_images(classification_map, [128,128],method = tf.image.ResizeMethod.NEAREST_NEIGHBOR )),name ='loss2')
loss3 = tl.cost.cross_entropy((m2.outputs), tf.squeeze( tf.image.resize_images(classification_map, [64,64],method = tf.image.ResizeMethod.NEAREST_NEIGHBOR) ),name='loss3')
loss4 = tl.cost.cross_entropy((m3.outputs), tf.squeeze( tf.image.resize_images(classification_map, [32,32],method = tf.image.ResizeMethod.NEAREST_NEIGHBOR )),name='loss4')
out =(net_regression.outputs)
loss = loss1 + loss2 +loss3 +loss4
#loss = tf.reduce_mean(tf.abs((net_regression.outputs + 1) - labels_sigma))
with tf.variable_scope('learning_rate'):
lr_v = tf.Variable(lr_init, trainable = False)
### DEFINE OPTIMIZER ###
vgg_vars = tl.layers.get_variables_with_name('VGG', True, True) # ?
t_vars = tl.layers.get_variables_with_name('UNet', True, True) #?
a_vars = tl.layers.get_variables_with_name('Unified', False, True) #
var_list1 = vgg_vars
var_list2 = t_vars
opt1 = tf.train.AdamOptimizer(lr_v*0.1*0.1)
opt2 = tf.train.AdamOptimizer(lr_v*0.1)
grads = tf.gradients(loss, var_list1 + var_list2)
grads1 = grads[:len(var_list1)]
grads2 = grads[len(var_list1):]
train_op1 = opt1.apply_gradients(zip(grads1, var_list1))
train_op2 = opt2.apply_gradients(zip(grads2, var_list2))
train_op = tf.group(train_op1, train_op2)
sess = tf.Session(config=tf.ConfigProto(allow_soft_placement = True, log_device_placement = False))
print "initializing global variable..."
tl.layers.initialize_global_variables(sess)
print "initializing global variable...DONE"
### initial checkpoint ###
checkpoint_dir2 = "test_checkpoint/PG_CUHK/"
tl.files.load_ckpt(sess=sess, mode_name='SA_net_PG_CUHK.ckpt', save_dir=checkpoint_dir2, var_list=a_vars, is_latest=True)
### START TRAINING ###
sess.run(tf.assign(lr_v, lr_init))
global_step = 0
new_lr_decay=1
prev_train_blur_imgs =train_blur_imgs
prev_train_classification_mask =train_mask_imgs
for epoch in range(0, n_epoch + 1):
## update learning rate
if epoch !=0 and (epoch % decay_every == 0):
new_lr_decay = lr_decay ** (epoch // decay_every)
#new_lr_decay = new_lr_decay * lr_decay
sess.run(tf.assign(lr_v, lr_init * new_lr_decay))
log = " ** new learning rate: %f" % (lr_init * new_lr_decay)
print(log)
elif epoch == 0:
sess.run(tf.assign(lr_v, lr_init))
log = " ** init lr: %f decay_every_init: %d, lr_decay: %f" % (lr_init, decay_every, lr_decay)
print(log)
epoch_time = time.time()
total_loss, n_iter = 0, 0
new_batch_size = batch_size #batchsize 50->40 + 10(augmented)
#data suffle***
suffle_index = np.arange(len(prev_train_blur_imgs))
np.random.shuffle(suffle_index)
print len(train_blur_imgs)
#print suffle_index
train_blur_imgs = []
train_classification_mask =[]
for i in range(0,len(suffle_index),1):
train_blur_imgs.append(prev_train_blur_imgs[suffle_index[i]] )
train_classification_mask.append(prev_train_classification_mask[suffle_index[i]] )
for idx in range(0, len(train_blur_imgs) , new_batch_size):
step_time = time.time()
augmentation_list = [0,1]
augmentation= random.choice(augmentation_list)
if(augmentation ==0):
images_and_score = tl.prepro.threading_data([_ for _ in zip(train_blur_imgs[idx: idx + new_batch_size],train_classification_mask[idx: idx + new_batch_size])],fn=crop_sub_img_and_classification_fn)
elif (augmentation==1):
images_and_score = tl.prepro.threading_data([_ for _ in zip(train_blur_imgs[idx: idx + new_batch_size],train_classification_mask[idx: idx + new_batch_size])],fn=crop_sub_img_and_classification_fn_aug)
#print images_and_score.shape
imlist, clist= images_and_score.transpose((1,0,2,3,4))
#print clist.shape
clist = clist[:, :, :, 0]
#print clist.shape
clist = np.expand_dims(clist, axis=3)
#print imlist.shape, clist.shape
err,l1,l2,l3,l4, _ ,outmap= sess.run([loss,loss1,loss2,loss3,loss4, train_op,out], {patches_blurred: imlist, classification_map: clist})
outmap1 = np.squeeze(outmap[1,:,:,0])
outmap2 = np.squeeze(outmap[1, :, :, 1])
outmap3 = np.squeeze(outmap[1, :, :, 2])
if(idx%100 ==0):
scipy.misc.imsave(save_dir_sample + '/input_mask.png', np.squeeze(clist[1, :, :, 0]))
scipy.misc.imsave(save_dir_sample + '/input.png', np.squeeze(imlist[1,:,:,:]))
scipy.misc.imsave(save_dir_sample + '/im.png', outmap1)
scipy.misc.imsave(save_dir_sample + '/im1.png', outmap2)
scipy.misc.imsave(save_dir_sample + '/im2.png', outmap3)
print("Epoch [%2d/%2d] %4d time: %4.4fs, err: %.6f, loss1: %.6f,loss2: %.6f,loss3: %.6f,loss4: %.6f" % (epoch, n_epoch, n_iter, time.time() - step_time, err,l1,l2,l3,l4))
total_loss += err
n_iter += 1
global_step += 1
log = "[*] Epoch: [%2d/%2d] time: %4.4fs, total_err: %.8f" % (epoch, n_epoch, time.time() - epoch_time, total_loss/n_iter)
print(log)
## save model
if epoch % 10== 0:
tl.files.save_ckpt(sess=sess, mode_name='SA_net_{}.ckpt'.format(tl.global_flag['mode']), save_dir = checkpoint_dir, var_list = a_vars, global_step = global_step, printable = False)
def main_train():
mask_perc = tl.global_flag['maskperc']
mask_name = tl.global_flag['mask']
model_name = tl.global_flag['model']
mode_lossfunciton = tl.global_flag['lossfunciton']
# mode_normalization = tl.global_flag['normalization']
mode_penalty = tl.global_flag['penalty']
# =================================== BASIC CONFIGS =================================== #
print('[*] run basic configs ... ')
log_dir = "logs/log_{}_{}_{}".format(model_name, mask_name, mask_perc)
tl.files.exists_or_mkdir(log_dir)
log_all, log_eval, log_50, log_all_filename, log_eval_filename, log_50_filename = logging_setup(
log_dir)
checkpoint_dir = "checkpoints/checkpoint_{}_{}_{}".format(
model_name, mask_name, mask_perc)
tl.files.exists_or_mkdir(checkpoint_dir)
save_dir = "samples/samples_{}_{}_{}".format(model_name, mask_name,
mask_perc)
tl.files.exists_or_mkdir(save_dir)
tensorboard_logs = "tensorboard_logs/tensorboard_logs_{}_{}_{}".format(
model_name, mask_name, mask_perc)
tl.files.exists_or_mkdir(tensorboard_logs)
# configs
batch_size = config.TRAIN.batch_size
early_stopping_num = config.TRAIN.early_stopping_num
g_alpha = config.TRAIN.g_alpha # weight for pixel loss
g_beta = config.TRAIN.g_beta # weight for frequency loss
g_gamma = config.TRAIN.g_gamma # weight for perceptual loss
g_adv = config.TRAIN.g_adv # weight for generate network loss
lr = config.TRAIN.lr # lr
lr_decay = config.TRAIN.lr_decay
decay_every = config.TRAIN.decay_every
beta1 = config.TRAIN.beta1
beta2 = config.TRAIN.beta2
n_epoch = config.TRAIN.n_epoch
sample_size = config.TRAIN.sample_size
n_critic = config.TRAIN.n_critic
log_config(log_all_filename, config)
log_config(log_eval_filename, config)
log_config(log_50_filename, config)
# ==================================== PREPARE DATA ==================================== #
print('[*] load data ... ')
training_data_path = config.TRAIN.training_data_path
val_data_path = config.TRAIN.val_data_path
testing_data_path = config.TRAIN.testing_data_path
with open(training_data_path, 'rb') as f:
X_train = pickle.load(f)
with open(val_data_path, 'rb') as f:
X_val = pickle.load(f)
with open(testing_data_path, 'rb') as f:
X_test = pickle.load(f)
print('X_train shape/min/max: ', X_train.shape, X_train.min(),
X_train.max()) # (11820, 256, 256, 1) -1.0 1.0
print('X_val shape/min/max: ', X_val.shape, X_val.min(),
X_val.max()) # (3897, 256, 256, 1) -1.0 1.0
print('X_test shape/min/max: ', X_test.shape, X_test.min(),
X_test.max()) # (7914, 256, 256, 1) -1.0 1.0
print('[*] loading mask ... ')
if mask_name == "radialcartesi":
mask = \
loadmat(
os.path.join(config.TRAIN.mask_RadialApproxOnCartesi_path, "RadialApproxOnCartesi_{}.mat".format(mask_perc)))[
'mask']
elif mask_name == "gaussian2d":
mask = \
loadmat(
os.path.join(config.TRAIN.mask_Gaussian2D_path, "GaussianDistribution2DMask_{}.mat".format(mask_perc)))[
'maskRS2']
elif mask_name == "gaussian1d":
mask = \
loadmat(
os.path.join(config.TRAIN.mask_Gaussian1D_path, "GaussianDistribution1DMask_{}.mat".format(mask_perc)))[
'maskRS1']
elif mask_name == "poisson2d":
mask = \
loadmat(
os.path.join(config.TRAIN.mask_Gaussian1D_path, "PoissonDistributionMask_{}.mat".format(mask_perc)))[
'population_matrix']
else:
raise ValueError("no such mask exists: {}".format(mask_name))
# ==================================== DEFINE MODEL ==================================== #
print('[*] define model ... ')
nw, nh, nz = X_train.shape[1:]
t_image_good = tf.placeholder('float32', [batch_size, nw, nh, nz],
name='good_image')
t_image_good_samples = tf.placeholder('float32', [sample_size, nw, nh, nz],
name='good_image_samples')
t_image_bad = tf.placeholder('float32', [batch_size, nw, nh, nz],
name='bad_image')
t_image_bad_samples = tf.placeholder('float32', [sample_size, nw, nh, nz],
name='bad_image_samples')
t_gen = tf.placeholder('float32', [batch_size, nw, nh, nz],
name='generated_image_for_test')
t_gen_sample = tf.placeholder('float32', [sample_size, nw, nh, nz],
name='generated_sample_image_for_test')
t_image_good_224 = tf.placeholder('float32', [batch_size, 224, 224, 3],
name='vgg_good_image')
# define generator network
if tl.global_flag['model'] == 'unet':
net = u_net_bn(t_image_bad,
is_train=True,
reuse=False,
is_refine=False)
net_test = u_net_bn(t_image_bad,
is_train=False,
reuse=True,
is_refine=False)
net_test_sample = u_net_bn(t_image_bad_samples,
is_train=False,
reuse=True,
is_refine=False)
elif tl.global_flag['model'] == 'unet_refine':
net = u_net_bn(t_image_bad, is_train=True, reuse=False, is_refine=True)
net_test = u_net_bn(t_image_bad,
is_train=False,
reuse=True,
is_refine=True)
net_test_sample = u_net_bn(t_image_bad_samples,
is_train=False,
reuse=True,
is_refine=True)
else:
raise Exception("unknown model")
# calculate generator and discriminator loss
net_d, logits_fake = discriminator(net.outputs,
lossfunciton=mode_lossfunciton,
is_train=True,
reuse=False)
_, logits_real = discriminator(t_image_good,
lossfunciton=mode_lossfunciton,
is_train=True,
reuse=True)
# define VGG network
net_vgg_conv4_good, net_vgg = vgg19_simple_api(t_image_good_224,
reuse=False)
net_vgg_conv4_gen, _ = vgg19_simple_api(tf.tile(
tf.image.resize_images(net.outputs, [224, 224]), [1, 1, 1, 3]),
reuse=True)
# ==================================== DEFINE LOSS ==================================== #
print('[*] define loss functions ... ')
if mode_lossfunciton == 'gan_loss':
d_loss1 = tl.cost.sigmoid_cross_entropy(logits_real,
tf.ones_like(logits_real),
name='d1')
d_loss2 = tl.cost.sigmoid_cross_entropy(logits_fake,
tf.zeros_like(logits_fake),
name='d2')
# discriminator loss (adversarial)
d_loss = d_loss1 + d_loss2
# generator loss (adversarial)
g_loss = tl.cost.sigmoid_cross_entropy(logits_fake,
tf.ones_like(logits_fake),
name='g')
elif mode_lossfunciton == 'wgan_loss':
# discriminator loss (adversarial)
d_loss = tf.reduce_mean(logits_fake) - tf.reduce_mean(logits_real)
# generator loss (adversarial)
g_loss = -tf.reduce_mean(logits_fake)
elif mode_lossfunciton == 'lsgan_loss':
d_loss_real = tf.reduce_mean(tf.square(logits_real - 1.0))
d_loss_fake = tf.reduce_mean(tf.square(logits_fake))
# discriminator loss (adversarial)
d_loss = 0.5 * (d_loss_real + d_loss_fake)
g_loss = 0.5 * tf.reduce_mean(tf.square(logits_fake - 1.0))
else:
raise Exception("Unknow lossfunction")
# ==================================== DEFINE Penality ==================================== #
print('[*] define loss functions ... ')
if mode_penalty == 'no_penalty':
penalty_loss = 0.0
d_loss = d_loss + penalty_loss
elif mode_penalty == 'wgangp_penalty':
epsilon = tf.random_uniform([], minval=0.0, maxval=1.0)
x_hat = t_image_good * epsilon + (1 - epsilon) * (net.outputs)
_, d_hat = discriminator(x_hat,
lossfunciton=mode_lossfunciton,
is_train=True,
reuse=True)
gradients = tf.gradients(d_hat, [x_hat])[0] # calculate the gradients
slopes = tf.sqrt(
tf.reduce_sum(tf.square(gradients), reduction_indices=[1]))
gradient_penalty = 1 * tf.reduce_mean((slopes - 1.0)**2)
d_loss = d_loss + gradient_penalty
elif mode_penalty == 'dragan_penalty':
_, var = tf.nn.moments(t_image_good,
axes=list(range(len(t_image_good.get_shape()))))
std = tf.sqrt(var)
x_noisy = t_image_good + std * (tf.random_uniform(t_image_good.shape) -
0.5)
x_noisy = tf.clip_by_value(x_noisy, 0.0, 1.0)
_, d_hat = discriminator(x_noisy,
lossfunciton=mode_lossfunciton,
is_train=True,
reuse=True)
gradients = tf.gradients(d_hat,
[x_noisy])[0] # calculate the gradients
slopes = tf.sqrt(
tf.reduce_sum(tf.square(gradients), reduction_indices=[1]))
gradient_penalty = 10 * tf.reduce_mean((slopes - 1.0)**2)
d_loss = d_loss + gradient_penalty
else:
raise Exception("Unknow penalty")
# generator loss (perceptual)
print("=================hee==============")
print(net_vgg_conv4_good.outputs.shape)
print(net_vgg_conv4_gen.outputs.shape)
g_perceptual = tl.cost.mean_squared_error(net_vgg_conv4_good.outputs,
net_vgg_conv4_gen.outputs,
is_mean=True)
# generator loss (pixel-wise) 1/2||Xt-Xu(~)||2
g_nmse_a = tf.sqrt(
tf.reduce_sum(tf.squared_difference(net.outputs, t_image_good),
axis=[1, 2, 3]))
g_nmse_b = tf.sqrt(tf.reduce_sum(tf.square(t_image_good), axis=[1, 2, 3]))
g_nmse = tf.reduce_mean(g_nmse_a / g_nmse_b)
# generator loss (frequency)
fft_good_abs = tf.map_fn(fft_abs_for_map_fn, t_image_good)
fft_gen_abs = tf.map_fn(fft_abs_for_map_fn, net.outputs)
g_fft = tf.reduce_mean(
tf.reduce_mean(tf.squared_difference(fft_good_abs, fft_gen_abs),
axis=[1, 2]))
# generator loss (total)总共损失
g_loss = g_adv * g_loss + g_alpha * g_nmse + g_gamma * g_perceptual + g_beta * g_fft
# g_loss = g_adv * g_loss + g_alpha * g_nmse
nmse_a_0_1 = tf.sqrt(
tf.reduce_sum(tf.squared_difference(t_gen, t_image_good),
axis=[1, 2, 3]))
nmse_b_0_1 = tf.sqrt(tf.reduce_sum(tf.square(t_image_good), axis=[1, 2,
3]))
nmse_0_1 = nmse_a_0_1 / nmse_b_0_1
nmse_a_0_1_sample = tf.sqrt(
tf.reduce_sum(tf.squared_difference(t_gen_sample,
t_image_good_samples),
axis=[1, 2, 3]))
nmse_b_0_1_sample = tf.sqrt(
tf.reduce_sum(tf.square(t_image_good_samples), axis=[1, 2, 3]))
nmse_0_1_sample = nmse_a_0_1_sample / nmse_b_0_1_sample
# Wasserstein GAN Loss
with tf.name_scope('WGAN_GP/D'):
tf.summary.scalar('d_loss', d_loss)
with tf.name_scope('WGAN_GP/G'):
tf.summary.scalar('g_loss', g_loss)
merged = tf.summary.merge_all()
# ==================================== DEFINE TRAIN OPTS ==================================== #
print('[*] define training options ... ')
g_vars = tl.layers.get_variables_with_name('u_net', True, True)
d_vars = tl.layers.get_variables_with_name('discriminator', True, True)
with tf.variable_scope('learning_rate'):
lr_v = tf.Variable(lr, trainable=False)
g_optim = tf.train.AdamOptimizer(lr_v, beta1=beta1,
beta2=beta2).minimize(g_loss,
var_list=g_vars)
d_optim = tf.train.AdamOptimizer(lr_v, beta1=beta1,
beta2=beta2).minimize(d_loss,
var_list=d_vars)
# ==================================== TRAINING ==================================== #
sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
loss_writer = tf.summary.FileWriter(tensorboard_logs, sess.graph)
tl.layers.initialize_global_variables(sess)
sess.run(tf.global_variables_initializer())
# load generator and discriminator weights (for continuous training purpose)
tl.files.load_and_assign_npz(
sess=sess,
name=os.path.join(checkpoint_dir, tl.global_flag['model']) + '.npz',
network=net)
tl.files.load_and_assign_npz(
sess=sess,
name=os.path.join(checkpoint_dir, tl.global_flag['model']) + '_d.npz',
network=net_d)
net_vgg_conv4_path = config.TRAIN.VGG19_path
if not os.path.isfile(net_vgg_conv4_path):
print(
"Please download vgg19.npz from : https://github.com/machrisaa/tensorflow-vgg"
)
exit() # 退出
npz = np.load(net_vgg_conv4_path, encoding='latin1').item()
# 这里几点注意
params = []
for val in sorted(npz.items()):
W = np.asarray(val[1][0])
b = np.asarray(val[1][1])
print(" Loading %s: %s, %s" % (val[0], W.shape, b.shape))
params.extend([W, b])
tl.files.assign_params(sess, params, net_vgg)
net_vgg.print_params(False)
n_training_examples = len(X_train)
n_step_epoch = round(n_training_examples / batch_size)
# sample testing images X_test:7914
idex = tl.utils.get_random_int(min_v=0,
max_v=len(X_test) - 1,
number=sample_size,
seed=config.TRAIN.seed)
X_samples_good = X_test[idex]
# X_samples_good shape: (50,256,256,1)
X_samples_bad = threading_data(X_samples_good, fn=to_bad_img, mask=mask)
x_good_sample_rescaled = (X_samples_good + 1) / 2
x_bad_sample_rescaled = (X_samples_bad + 1) / 2
tl.visualize.save_images(X_samples_good, [5, 10],
os.path.join(save_dir, "sample_image_good.png"))
tl.visualize.save_images(X_samples_bad, [5, 10],
os.path.join(save_dir, "sample_image_bad.png"))
tl.visualize.save_images(
np.abs(X_samples_good - X_samples_bad), [5, 10],
os.path.join(save_dir, "sample_image_diff_abs.png"))
tl.visualize.save_images(
np.sqrt(
np.abs(X_samples_good - X_samples_bad) / 2 + config.TRAIN.epsilon),
[5, 10], os.path.join(save_dir, "sample_image_diff_sqrt_abs.png"))
tl.visualize.save_images(
np.clip(10 * np.abs(X_samples_good - X_samples_bad) / 2, 0, 1),
[5, 10],
os.path.join(save_dir, "sample_image_diff_sqrt_abs_10_clip.png"))
tl.visualize.save_images(threading_data(X_samples_good, fn=distort_img),
[5, 10],
os.path.join(save_dir, "sample_image_aug.png"))
scipy.misc.imsave(os.path.join(save_dir, "mask.png"), mask * 255)
print('[*] start training ... ')
best_nmse = np.inf
best_epoch = 1
esn = early_stopping_num
for epoch in range(0, n_epoch):
# learning rate decay
if epoch != 0 and (epoch % decay_every == 0):
new_lr_decay = lr_decay**(epoch // decay_every)
sess.run(tf.assign(lr_v, lr * new_lr_decay))
log = " ** new learning rate: %f" % (lr * new_lr_decay)
print(log)
log_all.debug(log)
elif epoch == 0:
log = " ** init lr: %f decay_every_epoch: %d, lr_decay: %f" % (
lr, decay_every, lr_decay)
print(log)
log_all.debug(log)
for step in range(n_step_epoch):
step_time = time.time()
# now train the generator once!
idex = tl.utils.get_random_int(min_v=0,
max_v=n_training_examples - 1,
number=batch_size)
X_good = X_train[idex]
X_good_aug = threading_data(X_good, fn=distort_img)
X_good_224 = threading_data(X_good_aug, fn=vgg_prepro)
X_bad = threading_data(X_good_aug, fn=to_bad_img, mask=mask)
sess.run(d_optim, {t_image_good: X_good_aug, t_image_bad: X_bad})
sess.run(
g_optim, {
t_image_good_224: X_good_224,
t_image_good: X_good_aug,
t_image_bad: X_bad
})
errG, errG_perceptual, errG_nmse, errG_fft, errD, summary = sess.run(
[g_loss, g_perceptual, g_nmse, g_fft, d_loss, merged], {
t_image_good_224: X_good_224,
t_image_good: X_good_aug,
t_image_bad: X_bad
})
loss_writer.add_summary(summary, n_epoch)
log = "Epoch[{:3}/{:3}] step={:3} d_loss={:5} g_loss={:5} g_perceptual_loss={:5} g_mse={:5} g_freq={:5} took {:3}s".format(
epoch + 1, n_epoch, step, round(float(errD), 3),
round(float(errG), 3), round(float(errG_perceptual), 3),
round(float(errG_nmse), 3), round(float(errG_fft), 3),
round(time.time() - step_time, 2))
print(log)
log_all.debug(log)
# evaluation for training data
total_nmse_training = 0
total_ssim_training = 0
total_psnr_training = 0
num_training_temp = 0
for batch in tl.iterate.minibatches(inputs=X_train,
targets=X_train,
batch_size=batch_size,
shuffle=False):
x_good, _ = batch
# x_bad = threading_data(x_good, fn=to_bad_img, mask=mask)
x_bad = threading_data(x_good, fn=to_bad_img, mask=mask)
x_gen = sess.run(net_test.outputs, {t_image_bad: x_bad})
x_good_0_1 = (x_good + 1) / 2
x_gen_0_1 = (x_gen + 1) / 2
nmse_res = sess.run(nmse_0_1, {
t_gen: x_gen_0_1,
t_image_good: x_good_0_1
})
ssim_res = threading_data([_ for _ in zip(x_good_0_1, x_gen_0_1)],
fn=ssim)
psnr_res = threading_data([_ for _ in zip(x_good_0_1, x_gen_0_1)],
fn=psnr)
total_nmse_training += np.sum(nmse_res)
total_ssim_training += np.sum(ssim_res)
total_psnr_training += np.sum(psnr_res)
num_training_temp += batch_size
total_nmse_training /= num_training_temp
total_ssim_training /= num_training_temp
total_psnr_training /= num_training_temp
log = "Epoch: {}\nNMSE training: {:8}, SSIM training: {:8}, PSNR training: {:8}".format(
epoch + 1, total_nmse_training, total_ssim_training,
total_psnr_training)
print(log)
log_all.debug(log)
log_eval.info(log)
# evaluation for validation data
total_nmse_val = 0
total_ssim_val = 0
total_psnr_val = 0
num_val_temp = 0
for batch in tl.iterate.minibatches(inputs=X_val,
targets=X_val,
batch_size=batch_size,
shuffle=False):
x_good, _ = batch
# x_bad = threading_data(x_good, fn=to_bad_img, mask=mask)
x_bad = threading_data(x_good, fn=to_bad_img, mask=mask)
x_gen = sess.run(net_test.outputs, {t_image_bad: x_bad})
x_good_0_1 = (x_good + 1) / 2
x_gen_0_1 = (x_gen + 1) / 2
nmse_res = sess.run(nmse_0_1, {
t_gen: x_gen_0_1,
t_image_good: x_good_0_1
})
ssim_res = threading_data([_ for _ in zip(x_good_0_1, x_gen_0_1)],
fn=ssim)
psnr_res = threading_data([_ for _ in zip(x_good_0_1, x_gen_0_1)],
fn=psnr)
total_nmse_val += np.sum(nmse_res)
total_ssim_val += np.sum(ssim_res)
total_psnr_val += np.sum(psnr_res)
num_val_temp += batch_size
total_nmse_val /= num_val_temp
total_ssim_val /= num_val_temp
total_psnr_val /= num_val_temp
log = "Epoch: {}\nNMSE val: {:8}, SSIM val: {:8}, PSNR val: {:8}".format(
epoch + 1, total_nmse_val, total_ssim_val, total_psnr_val)
print(log)
log_all.debug(log)
log_eval.info(log)
img = sess.run(net_test_sample.outputs,
{t_image_bad_samples: X_samples_bad})
tl.visualize.save_images(
img, [5, 10], os.path.join(save_dir, "image_{}.png".format(epoch)))
if total_nmse_val < best_nmse:
esn = early_stopping_num
best_nmse = total_nmse_val
best_epoch = epoch + 1
# save current best model
tl.files.save_npz(
net.all_params,
name=os.path.join(checkpoint_dir, tl.global_flag['model']) +
'.npz',
sess=sess)
tl.files.save_npz(
net_d.all_params,
name=os.path.join(checkpoint_dir, tl.global_flag['model']) +
'_d.npz',
sess=sess)
print("[*] Save checkpoints SUCCESS!")
else:
esn -= 1
log = "Best NMSE result: {} at {} epoch".format(best_nmse, best_epoch)
log_eval.info(log)
log_all.debug(log)
print(log)
if esn == 0:
log_eval.info(log)
tl.files.load_and_assign_npz(
sess=sess,
name=os.path.join(checkpoint_dir, tl.global_flag['model']) +
'.npz',
network=net)
# evluation for test data
x_gen = sess.run(net_test_sample.outputs,
{t_image_bad_samples: X_samples_bad})
x_gen_0_1 = (x_gen + 1) / 2
savemat(save_dir + '/test_random_50_generated.mat',
{'x_gen_0_1': x_gen_0_1})
nmse_res = sess.run(
nmse_0_1_sample, {
t_gen_sample: x_gen_0_1,
t_image_good_samples: x_good_sample_rescaled
})
ssim_res = threading_data(
[_ for _ in zip(x_good_sample_rescaled, x_gen_0_1)], fn=ssim)
psnr_res = threading_data(
[_ for _ in zip(x_good_sample_rescaled, x_gen_0_1)], fn=psnr)
log = "NMSE testing: {}\nSSIM testing: {}\nPSNR testing: {}\n\n".format(
nmse_res, ssim_res, psnr_res)
log_50.debug(log)
log = "NMSE testing average: {}\nSSIM testing average: {}\nPSNR testing average: {}\n\n".format(
np.mean(nmse_res), np.mean(ssim_res), np.mean(psnr_res))
log_50.debug(log)
log = "NMSE testing std: {}\nSSIM testing std: {}\nPSNR testing std: {}\n\n".format(
np.std(nmse_res), np.std(ssim_res), np.std(psnr_res))
log_50.debug(log)
# evaluation for zero-filled (ZF) data,
nmse_res_zf = sess.run(
nmse_0_1_sample, {
t_gen_sample: x_bad_sample_rescaled,
t_image_good_samples: x_good_sample_rescaled
})
ssim_res_zf = threading_data([
_ for _ in zip(x_good_sample_rescaled, x_bad_sample_rescaled)
],
fn=ssim)
psnr_res_zf = threading_data([
_ for _ in zip(x_good_sample_rescaled, x_bad_sample_rescaled)
],
fn=psnr)
log = "NMSE ZF testing: {}\nSSIM ZF testing: {}\nPSNR ZF testing: {}\n\n".format(
nmse_res_zf, ssim_res_zf, psnr_res_zf)
log_50.debug(log)
log = "NMSE ZF average testing: {}\nSSIM ZF average testing: {}\nPSNR ZF average testing: {}\n\n".format(
np.mean(nmse_res_zf), np.mean(ssim_res_zf),
np.mean(psnr_res_zf))
log_50.debug(log)
log = "NMSE ZF std testing: {}\nSSIM ZF std testing: {}\nPSNR ZF std testing: {}\n\n".format(
np.std(nmse_res_zf), np.std(ssim_res_zf), np.std(psnr_res_zf))
log_50.debug(log)
# sample testing images
tl.visualize.save_images(
x_gen, [5, 10],
os.path.join(save_dir, "final_generated_image.png"))
tl.visualize.save_images(
np.clip(10 * np.abs(X_samples_good - x_gen) / 2, 0, 1),
[5, 10],
os.path.join(save_dir,
"final_generated_image_diff_abs_10_clip.png"))
tl.visualize.save_images(
np.clip(10 * np.abs(X_samples_good - X_samples_bad) / 2, 0, 1),
[5, 10],
os.path.join(save_dir, "final_bad_image_diff_abs_10_clip.png"))
print("[*] Job finished!")
break
def main_evaluate():
# load mask
mask_perc = tl.global_flag['maskperc']
mask_name = tl.global_flag['mask']
model_name = tl.global_flag['model']
# load the parameters
log_dir = "logs/log_{}_{}_{}".format(model_name, mask_name, mask_perc)
tl.files.exists_or_mkdir(log_dir)
log_all, log_eval, log_50, log_all_filename, log_eval_filename, log_50_filename = logging_setup(
log_dir)
log_config(log_50_filename, config)
checkpoint_dir = "checkpoints/checkpoint_{}_{}_{}".format(
model_name, mask_name, mask_perc)
# samples the save of the evaluation
evaluate_data_sample_path = config.TRAIN.evaluate_data_sample_path
tl.files.exists_or_mkdir(evaluate_data_sample_path)
evaluate_data_save_path = config.TRAIN.evaluate_data_save_path
tl.files.exists_or_mkdir(evaluate_data_save_path)
print('[*] loading mask ... ')
if mask_name == "radialcartesi":
mask = \
loadmat(
os.path.join(config.TRAIN.mask_RadialApproxOnCartesi_path,
"RadialApproxOnCartesi_{}.mat".format(mask_perc)))[
'mask']
elif mask_name == "gaussian2d":
mask = \
loadmat(
os.path.join(config.TRAIN.mask_Gaussian2D_path, "GaussianDistribution2DMask_{}.mat".format(mask_perc)))[
'maskRS2']
elif mask_name == "gaussian1d":
mask = \
loadmat(
os.path.join(config.TRAIN.mask_Gaussian1D_path, "GaussianDistribution1DMask_{}.mat".format(mask_perc)))[
'maskRS1']
elif mask_name == "poisson2d":
mask = \
loadmat(
os.path.join(config.TRAIN.mask_Gaussian1D_path, "PoissonDistributionMask_{}.mat".format(mask_perc)))[
'population_matrix']
else:
raise ValueError("no such mask exists: {}".format(mask_name))
print('[*] load evaluate data ... ')
evaluate_gd_img_list = sorted(
tl.files.load_file_list(path=evaluate_data_sample_path,
regx='.*.png',
printable=False))
evaluate_gd_imgs = tl.vis.read_images(evaluate_gd_img_list,
path=evaluate_data_sample_path,
n_threads=1)
###========================== DEFINE MODEL ============================###
# imid = 0 # the first image
# evaluate_gd_img = evaluate_gd_imgs[imid]
print("=================test===============")
test_length = len(evaluate_gd_imgs)
test_good_image = []
for i in range(test_length):
evaluate_gd_img = evaluate_gd_imgs[i]
evaluate_gd_img = evaluate_gd_img[:, :, np.newaxis]
# print(evaluate_gd_img.shape)
evaluate_gd_img = evaluate_gd_img / 127.5 - 1
test_good_image.append(evaluate_gd_img)
# print("lens of test", len(test_list))
evaluate_samples_bad = threading_data(test_good_image,
fn=to_bad_img,
mask=mask)
x_bad_sample_rescaled = (evaluate_samples_bad + 1) / 2
print("====================test length==================")
size = evaluate_samples_bad.shape
print(size)
print(len(evaluate_samples_bad))
evaluate_image = tf.placeholder('float32',
[None, size[1], size[2], size[3]],
name='input_image')
t_gen_sample = tf.placeholder('float32', [None, size[1], size[2], size[3]],
name='generated_sample_image_for_test')
t_image_good_samples = tf.placeholder('float32',
[None, size[1], size[2], size[3]],
name='good_image_samples')
nmse_a_0_1_sample = tf.sqrt(
tf.reduce_sum(tf.squared_difference(t_gen_sample,
t_image_good_samples),
axis=[1, 2, 3]))
nmse_b_0_1_sample = tf.sqrt(
tf.reduce_sum(tf.square(t_image_good_samples), axis=[1, 2, 3]))
nmse_0_1_sample = nmse_a_0_1_sample / nmse_b_0_1_sample
###========================== RESTORE G =============================###
# define generator network
if tl.global_flag['model'] == 'unet':
net = u_net_bn(evaluate_image,
is_train=False,
reuse=False,
is_refine=False)
elif tl.global_flag['model'] == 'unet_refine':
net = u_net_bn(evaluate_image,
is_train=False,
reuse=False,
is_refine=True)
else:
raise Exception("unknown model")
sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True,
log_device_placement=False))
tl.layers.initialize_global_variables(sess)
tl.files.load_and_assign_npz(
sess=sess,
name=os.path.join(checkpoint_dir, tl.global_flag['model']) + '.npz',
network=net)
###======================= EVALUATION =============================###
start_time = time.time()
evaluate_restore_img = sess.run(net.outputs,
{evaluate_image: evaluate_samples_bad})
print("took: %4.4fs" % (time.time() - start_time))
print("zero filled size : %s / restore image size: %s" %
(size, evaluate_restore_img.shape))
for i in range(test_length):
if i >= 0 and i < 9:
# tl.vis.save_image(evaluate_restore_img[i], evaluate_data_save_path + '/resore/00{}.png'.format(i + 1))
savemat(evaluate_data_save_path + '/resore/00{}.mat'.format(i + 1),
{'dagangp': evaluate_restore_img[i]})
# tl.vis.save_image(evaluate_samples_bad[i], evaluate_data_save_path + '/zerofilled/00{}.png'.format(i + 1))
savemat(
evaluate_data_save_path + '/zerofilled/00{}.mat'.format(i + 1),
{'zerofill': evaluate_samples_bad[i]})
elif i >= 9 and i < 99:
# tl.vis.save_image(evaluate_restore_img[i], evaluate_data_save_path + '/resore/00{}.png'.format(i + 1))
savemat(evaluate_data_save_path + '/resore/0{}.mat'.format(i + 1),
{'dagangp': evaluate_restore_img[i]})
# tl.vis.save_image(evaluate_samples_bad[i], evaluate_data_save_path + '/zerofilled/00{}.png'.format(i + 1))
savemat(
evaluate_data_save_path + '/zerofilled/0{}.mat'.format(i + 1),
{'zerofill': evaluate_samples_bad[i]})
else:
# tl.vis.save_image(evaluate_restore_img[i], evaluate_data_save_path + '/resore/00{}.png'.format(i + 1))
savemat(evaluate_data_save_path + '/resore/{}.mat'.format(i + 1),
{'dagangp': evaluate_restore_img[i]})
# tl.vis.save_image(evaluate_samples_bad[i], evaluate_data_save_path + '/zerofilled/00{}.png'.format(i + 1))
savemat(
evaluate_data_save_path + '/zerofilled/{}.mat'.format(i + 1),
{'zerofill': evaluate_samples_bad[i]})
# calculatge the NMSE, ssim, psnr
evaluate_restore_img_rescaled = (evaluate_restore_img + 1) / 2
evaluate_gd_img_rescaled = []
for i in range(test_length):
test_good_image[i] = (test_good_image[i] + 1) / 2
evaluate_gd_img_rescaled.append(test_good_image[i])
# print("evaluate_restore_img_rescaled shape")
# print(evaluate_restore_img_rescaled.shape)
# print(evaluate_gd_img_rescaled.shape)
print("start evaluate ssim and PSNR")
print("===================start: good-restore================")
ssim_res = threading_data([
_ for _ in zip(evaluate_gd_img_rescaled, evaluate_restore_img_rescaled)
],
fn=ssim)
psnr_res = threading_data([
_ for _ in zip(evaluate_gd_img_rescaled, evaluate_restore_img_rescaled)
],
fn=psnr)
nmse_res = sess.run(
nmse_0_1_sample, {
t_gen_sample: evaluate_restore_img_rescaled,
t_image_good_samples: evaluate_gd_img_rescaled
})
log = "NMSE testing: {}\nSSIM testing: {}\nPSNR testing: {}\n\n".format(
nmse_res, ssim_res, psnr_res)
log_50.debug(log)
log = "NMSE testing average: {}\nSSIM testing average: {}\nPSNR testing average: {}\n\n".format(
np.mean(nmse_res), np.mean(ssim_res), np.mean(psnr_res))
log_50.debug(log)
log = "NMSE testing std: {}\nSSIM testing std: {}\nPSNR testing std: {}\n\n".format(
np.std(nmse_res), np.std(ssim_res), np.std(psnr_res))
#
log_50.debug(log)
print("===================start: good-zf ================")
ssim_res_zf = threading_data(
[_ for _ in zip(evaluate_gd_img_rescaled, x_bad_sample_rescaled)],
fn=ssim)
psnr_res_zf = threading_data(
[_ for _ in zip(evaluate_gd_img_rescaled, x_bad_sample_rescaled)],
fn=psnr)
nmse_res_zf = sess.run(
nmse_0_1_sample, {
t_gen_sample: x_bad_sample_rescaled,
t_image_good_samples: evaluate_gd_img_rescaled
})
log = "NMSE ZF testing: {}\nSSIM ZF testing: {}\nPSNR ZF testing: {}\n\n".format(
nmse_res_zf, ssim_res_zf, psnr_res_zf)
log_50.debug(log)
log = "NMSE ZF average testing: {}\nSSIM ZF average testing: {}\nPSNR ZF average testing: {}\n\n".format(
np.mean(nmse_res_zf), np.mean(ssim_res_zf), np.mean(psnr_res_zf))
log_50.debug(log)
log = "NMSE ZF std testing: {}\nSSIM ZF std testing: {}\nPSNR ZF std testing: {}\n\n".format(
np.std(nmse_res_zf), np.std(ssim_res_zf), np.std(psnr_res_zf))
log_50.debug(log)
print("job finished")
from config import CURRENT_EXP_DIR, config, get_logger, log_config
logger = get_logger(exp_dir=CURRENT_EXP_DIR)
log_config(logger)
import numpy as np
import sherpa
import sherpa.algorithms.bayesian_optimization as bayesian_optimization
from cade.metrics.comparative import intersection_nn, lncs2, moving_lncs2
from gensim.models.word2vec import Word2Vec
from scipy.spatial.distance import cosine
from sklearn.metrics import (
accuracy_score,
f1_score,
precision_score,
recall_score,
)
def fitness(threshold: float, topn: int, t: float):
similarity = 3
# Task 1 - Binary Classification
predictions = []
for word in binary_truth[:, 0]:
if similarity == 0: # cosine similarity
prediction = (0 if
1 - cosine(model1[word], model2[word]) >= threshold
else 1)
elif similarity == 1: # lncs2 similarity
prediction = (0 if lncs2(word, model1, model2, topn) >= threshold
parser.add_argument('-es',
'--epoch_start',
type=int,
default=config.epoch_start,
help='whether to apply semantic replacement')
args = parser.parse_args()
config = get_config(project, args.mode)
config.is_train = args.is_train
config.delete_log = args.delete_log
config.lr_init = args.lr_init
config.thread_num = args.thread_num
config.batch_size = args.batch_size
config.model = args.model
config.load_pretrained = args.load_pretrained
config.epoch_start = args.epoch_start
handle_directory(config, args.delete_log)
print_config(config)
log_config(config.LOG_DIR.config, config)
if config.is_train:
trainer = Trainer(config)
trainer.train()
# else:
# return
# ==============================================================================
import multiprocessing
import os, sys, time
from config import config, log_config
import util
AGENT_COUNT = config["agent_config"]["count"]
EVALUATOR_COUNT = config["evaluator_config"]["count"]
MODEL_AUGMENTED = config["model_config"] is not False
if config["resume"]:
ROOT_PATH = "output/" + config["env"]["name"] + "/" + config["name"]
else:
ROOT_PATH = util.create_and_wipe_directory("output/" + config["env"]["name"] + "/" + config["name"])
log_config()
import learner, agent, valuerl_learner
if MODEL_AUGMENTED: import worldmodel_learner
if __name__ == '__main__':
all_procs = set([])
interaction_procs = set([])
# lock
policy_lock = multiprocessing.Lock()
model_lock = multiprocessing.Lock() if MODEL_AUGMENTED else None
# queue
policy_replay_frame_queue = multiprocessing.Queue(1)
model_replay_frame_queue = multiprocessing.Queue(1) if MODEL_AUGMENTED else None
def main_train():
mask_perc = tl.global_flag['maskperc']
mask_name = tl.global_flag['mask']
model_name = tl.global_flag['model']
# =================================== BASIC CONFIGS =================================== #
print('[*] run basic configs ... ')
log_dir = "log_{}_{}_{}".format(model_name, mask_name, mask_perc)
tl.files.exists_or_mkdir(log_dir)
log_all, log_eval, log_50, log_all_filename, log_eval_filename, log_50_filename = logging_setup(
log_dir)
checkpoint_dir = "checkpoint_{}_{}_{}".format(model_name, mask_name,
mask_perc)
tl.files.exists_or_mkdir(checkpoint_dir)
save_dir = "samples_{}_{}_{}".format(model_name, mask_name, mask_perc)
tl.files.exists_or_mkdir(save_dir)
# configs
batch_size = config.TRAIN.batch_size
early_stopping_num = config.TRAIN.early_stopping_num
g_alpha = config.TRAIN.g_alpha
g_beta = config.TRAIN.g_beta
g_gamma = config.TRAIN.g_gamma
g_adv = config.TRAIN.g_adv
lr = config.TRAIN.lr
lr_decay = config.TRAIN.lr_decay
decay_every = config.TRAIN.decay_every
beta1 = config.TRAIN.beta1
n_epoch = config.TRAIN.n_epoch
sample_size = config.TRAIN.sample_size
log_config(log_all_filename, config)
log_config(log_eval_filename, config)
log_config(log_50_filename, config)
# ==================================== PREPARE DATA ==================================== #
print('[*] load data ... ')
# training_data_path = config.TRAIN.training_data_path
# val_data_path = config.TRAIN.val_data_path
# testing_data_path = config.TRAIN.testing_data_path
training_target_data_path = config.TRAIN.training_target_data_path
training_blurry_data_path = config.TRAIN.training_blurry_data_path
val_target_data_path = config.TRAIN.val_target_data_path
val_blurry_data_path = config.TRAIN.val_blurry_data_path
testing_target_data_path = config.TRAIN.testing_target_data_path
testing_blurry_data_path = config.TRAIN.testing_blurry_data_path
with open(training_target_data_path, 'rb') as f:
X_train_target = pickle.load(f)
with open(training_blurry_data_path, 'rb') as f:
X_train_blurry = pickle.load(f)
with open(val_target_data_path, 'rb') as f:
X_val_target = pickle.load(f)
with open(val_blurry_data_path, 'rb') as f:
X_val_blurry = pickle.load(f)
with open(testing_target_data_path, 'rb') as f:
X_test_target = pickle.load(f)
with open(testing_blurry_data_path, 'rb') as f:
X_test_blurry = pickle.load(f)
print('X_train_target shape/min/max: ', X_train_target.shape,
X_train_target.min(), X_train_target.max())
print('X_train_blurry shape/min/max: ', X_train_blurry.shape,
X_train_blurry.min(), X_train_blurry.max())
print('X_val_target shape/min/max: ', X_val_target.shape,
X_val_target.min(), X_val_target.max())
print('X_val_blurry shape/min/max: ', X_val_blurry.shape,
X_val_blurry.min(), X_val_blurry.max())
print('X_test_target shape/min/max: ', X_test_target.shape,
X_test_target.min(), X_test_target.max())
print('X_test_blurry shape/min/max: ', X_test_blurry.shape,
X_test_blurry.min(), X_test_blurry.max())
print('[*] loading mask ... ')
if mask_name == "gaussian2d":
mask = \
loadmat(
os.path.join(config.TRAIN.mask_Gaussian2D_path, "GaussianDistribution2DMask_{}.mat".format(mask_perc)))[
'maskRS2']
elif mask_name == "gaussian1d":
mask = \
loadmat(
os.path.join(config.TRAIN.mask_Gaussian1D_path, "GaussianDistribution1DMask_{}.mat".format(mask_perc)))[
'maskRS1']
elif mask_name == "poisson2d":
mask = \
loadmat(
os.path.join(config.TRAIN.mask_Gaussian1D_path, "PoissonDistributionMask_{}.mat".format(mask_perc)))[
'population_matrix']
else:
raise ValueError("no such mask exists: {}".format(mask_name))
# ==================================== DEFINE MODEL ==================================== #
print('[*] define model ... ')
nw, nh, nz = X_train_target.shape[1:]
# define placeholders
print("Model - define placeholders")
t_image_good = tf.placeholder('float32', [batch_size, nw, nh, nz],
name='good_image')
t_image_good_samples = tf.placeholder('float32', [sample_size, nw, nh, nz],
name='good_image_samples')
t_image_bad = tf.placeholder('float32', [batch_size, nw, nh, nz],
name='bad_image')
t_image_bad_samples = tf.placeholder('float32', [sample_size, nw, nh, nz],
name='bad_image_samples')
t_gen = tf.placeholder('float32', [batch_size, nw, nh, nz],
name='generated_image_for_test')
t_gen_sample = tf.placeholder('float32', [sample_size, nw, nh, nz],
name='generated_sample_image_for_test')
t_image_good_244 = tf.placeholder('float32', [batch_size, 244, 244, 3],
name='vgg_good_image')
# define generator network
print("Model - define generator network")
if tl.global_flag['model'] == 'unet':
net = u_net_bn(t_image_bad,
is_train=True,
reuse=False,
is_refine=False)
net_test = u_net_bn(t_image_bad,
is_train=False,
reuse=True,
is_refine=False)
net_test_sample = u_net_bn(t_image_bad_samples,
is_train=False,
reuse=True,
is_refine=False)
elif tl.global_flag['model'] == 'unet_refine':
net = u_net_bn(t_image_bad, is_train=True, reuse=False, is_refine=True)
net_test = u_net_bn(t_image_bad,
is_train=False,
reuse=True,
is_refine=True)
net_test_sample = u_net_bn(t_image_bad_samples,
is_train=False,
reuse=True,
is_refine=True)
else:
raise Exception("unknown model")
# define discriminator network
print("Model - define discriminator network")
net_d, logits_fake = discriminator(net.outputs, is_train=True, reuse=False)
_, logits_real = discriminator(t_image_good, is_train=True, reuse=True)
# define VGG network
print("Model - define VGG network")
net_vgg_conv4_good, _ = vgg16_cnn_emb(t_image_good_244, reuse=False)
net_vgg_conv4_gen, _ = vgg16_cnn_emb(tf.tile(
tf.image.resize_images(net.outputs, [244, 244]), [1, 1, 1, 3]),
reuse=True)
# ==================================== DEFINE LOSS ==================================== #
print('[*] define loss functions ... ')
# discriminator loss
d_loss1 = tl.cost.sigmoid_cross_entropy(logits_real,
tf.ones_like(logits_real),
name='d1')
d_loss2 = tl.cost.sigmoid_cross_entropy(logits_fake,
tf.zeros_like(logits_fake),
name='d2')
d_loss = d_loss1 + d_loss2
# generator loss (adversarial)
g_loss = tl.cost.sigmoid_cross_entropy(logits_fake,
tf.ones_like(logits_fake),
name='g')
# generator loss (perceptual)
g_perceptual = tf.reduce_mean(
tf.reduce_mean(tf.squared_difference(net_vgg_conv4_good.outputs,
net_vgg_conv4_gen.outputs),
axis=[1, 2, 3]))
# generator loss (pixel-wise)
g_nmse_a = tf.sqrt(
tf.reduce_sum(tf.squared_difference(net.outputs, t_image_good),
axis=[1, 2, 3]))
g_nmse_b = tf.sqrt(tf.reduce_sum(tf.square(t_image_good), axis=[1, 2, 3]))
g_nmse = tf.reduce_mean(g_nmse_a / g_nmse_b)
# generator loss (frequency)
fft_good_abs = tf.map_fn(fft_abs_for_map_fn, t_image_good)
fft_gen_abs = tf.map_fn(fft_abs_for_map_fn, net.outputs)
g_fft = tf.reduce_mean(
tf.reduce_mean(tf.squared_difference(fft_good_abs, fft_gen_abs),
axis=[1, 2]))
# generator loss (total)
g_loss = g_adv * g_loss + g_alpha * g_nmse + g_gamma * g_perceptual + g_beta * g_fft
# nmse metric for testing purpose
nmse_a_0_1 = tf.sqrt(
tf.reduce_sum(tf.squared_difference(t_gen, t_image_good),
axis=[1, 2, 3]))
nmse_b_0_1 = tf.sqrt(tf.reduce_sum(tf.square(t_image_good), axis=[1, 2,
3]))
nmse_0_1 = nmse_a_0_1 / nmse_b_0_1
nmse_a_0_1_sample = tf.sqrt(
tf.reduce_sum(tf.squared_difference(t_gen_sample,
t_image_good_samples),
axis=[1, 2, 3]))
nmse_b_0_1_sample = tf.sqrt(
tf.reduce_sum(tf.square(t_image_good_samples), axis=[1, 2, 3]))
nmse_0_1_sample = nmse_a_0_1_sample / nmse_b_0_1_sample
# ==================================== DEFINE TRAIN OPTS ==================================== #
print('[*] define training options ... ')
g_vars = tl.layers.get_variables_with_name('u_net', True, True)
d_vars = tl.layers.get_variables_with_name('discriminator', True, True)
with tf.variable_scope('learning_rate'):
lr_v = tf.Variable(lr, trainable=False)
g_optim = tf.train.AdamOptimizer(lr_v,
beta1=beta1).minimize(g_loss,
var_list=g_vars)
d_optim = tf.train.AdamOptimizer(lr_v,
beta1=beta1).minimize(d_loss,
var_list=d_vars)
# ==================================== TRAINING ==================================== #
sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
tl.layers.initialize_global_variables(sess)
# load generator and discriminator weights (for continuous training purpose)
tl.files.load_and_assign_npz(
sess=sess,
name=os.path.join(checkpoint_dir, tl.global_flag['model']) + '.npz',
network=net)
tl.files.load_and_assign_npz(
sess=sess,
name=os.path.join(checkpoint_dir, tl.global_flag['model']) + '_d.npz',
network=net_d)
# load vgg weights
net_vgg_conv4_path = config.TRAIN.VGG16_path
npz = np.load(net_vgg_conv4_path)
assign_op = []
for idx, val in enumerate(sorted(npz.items())[0:20]):
print(" Loading pretrained VGG16, CNN part %s" % str(val[1].shape))
assign_op.append(net_vgg_conv4_good.all_params[idx].assign(val[1]))
sess.run(assign_op)
net_vgg_conv4_good.print_params(False)
n_training_examples = len(X_train_target)
n_step_epoch = round(n_training_examples / batch_size)
# sample testing images
idex = tl.utils.get_random_int(min=0,
max=len(X_test_target) - 1,
number=sample_size,
seed=config.TRAIN.seed)
X_samples_good = X_test_target[idex]
# X_samples_bad = threading_data(X_samples_good, fn=to_bad_img, mask=mask)
X_samples_bad = X_test_blurry[idex]
x_good_sample_rescaled = (X_samples_good + 1) / 2
x_bad_sample_rescaled = (X_samples_bad + 1) / 2
tl.visualize.save_images(X_samples_good, [5, 10],
os.path.join(save_dir, "sample_image_good.png"))
tl.visualize.save_images(X_samples_bad, [5, 10],
os.path.join(save_dir, "sample_image_bad.png"))
tl.visualize.save_images(
np.abs(X_samples_good - X_samples_bad), [5, 10],
os.path.join(save_dir, "sample_image_diff_abs.png"))
tl.visualize.save_images(
np.sqrt(
np.abs(X_samples_good - X_samples_bad) / 2 + config.TRAIN.epsilon),
[5, 10], os.path.join(save_dir, "sample_image_diff_sqrt_abs.png"))
tl.visualize.save_images(
np.clip(10 * np.abs(X_samples_good - X_samples_bad) / 2, 0, 1),
[5, 10],
os.path.join(save_dir, "sample_image_diff_sqrt_abs_10_clip.png"))
tl.visualize.save_images(threading_data(X_samples_good, fn=distort_img),
[5, 10],
os.path.join(save_dir, "sample_image_aug.png"))
scipy.misc.imsave(os.path.join(save_dir, "mask.png"), mask * 255)
print('[*] start training ... ')
best_nmse = np.inf
best_epoch = 1
esn = early_stopping_num
for epoch in range(0, n_epoch):
# learning rate decay
if epoch != 0 and (epoch % decay_every == 0):
new_lr_decay = lr_decay**(epoch // decay_every)
sess.run(tf.assign(lr_v, lr * new_lr_decay))
log = " ** new learning rate: %f" % (lr * new_lr_decay)
print(log)
log_all.debug(log)
elif epoch == 0:
log = " ** init lr: %f decay_every_epoch: %d, lr_decay: %f" % (
lr, decay_every, lr_decay)
print(log)
log_all.debug(log)
for step in range(n_step_epoch):
step_time = time.time()
idex = tl.utils.get_random_int(min=0,
max=n_training_examples - 1,
number=batch_size)
# X_good = X_train[idex]
# X_good_aug = threading_data(X_good, fn=distort_img)
# X_good_244 = threading_data(X_good_aug, fn=vgg_prepro)
# X_bad = threading_data(X_good_aug, fn=to_bad_img, mask=mask)
X_good = X_train_target[idex]
X_good_aug = X_good
X_good_244 = threading_data(X_good_aug, fn=vgg_prepro)
X_bad = X_train_blurry[idex]
errD, _ = sess.run([d_loss, d_optim], {
t_image_good: X_good_aug,
t_image_bad: X_bad
})
# errD = sess.run([d_loss], {t_image_good: X_good_aug, t_image_bad: X_bad})
# errD = errD[0]
errG, errG_perceptual, errG_nmse, errG_fft, _ = sess.run(
[g_loss, g_perceptual, g_nmse, g_fft, g_optim], {
t_image_good_244: X_good_244,
t_image_good: X_good_aug,
t_image_bad: X_bad
})
log = "Epoch[{:3}/{:3}] step={:3} d_loss={:5} g_loss={:5} g_perceptual_loss={:5} g_mse={:5} g_freq={:5} took {:3}s".format(
epoch + 1, n_epoch, step, round(float(errD), 3),
round(float(errG), 3), round(float(errG_perceptual), 3),
round(float(errG_nmse), 3), round(float(errG_fft), 3),
round(time.time() - step_time, 2))
print(log)
log_all.debug(log)
################################## evaluation (part of training loop) ###############################
# evaluation for training data
total_nmse_training = 0
total_ssim_training = 0
total_psnr_training = 0
num_training_temp = 0
for batch in tl.iterate.minibatches(inputs=X_train_blurry,
targets=X_train_target,
batch_size=batch_size,
shuffle=False):
# x_good, _ = batch
# # x_bad = threading_data(x_good, fn=to_bad_img, mask=mask)
# x_bad = threading_data(
# x_good,
# fn=to_bad_img,
# mask=mask)
x_bad, x_good = batch
x_gen = sess.run(net_test.outputs, {t_image_bad: x_bad})
x_good_0_1 = (x_good + 1) / 2
x_gen_0_1 = (x_gen + 1) / 2
nmse_res = sess.run(nmse_0_1, {
t_gen: x_gen_0_1,
t_image_good: x_good_0_1
})
ssim_res = threading_data([_ for _ in zip(x_good_0_1, x_gen_0_1)],
fn=ssim)
psnr_res = threading_data([_ for _ in zip(x_good_0_1, x_gen_0_1)],
fn=psnr)
total_nmse_training += np.sum(nmse_res)
total_ssim_training += np.sum(ssim_res)
total_psnr_training += np.sum(psnr_res)
num_training_temp += batch_size
total_nmse_training /= num_training_temp
total_ssim_training /= num_training_temp
total_psnr_training /= num_training_temp
log = "Epoch: {}\nNMSE training: {:8}, SSIM training: {:8}, PSNR training: {:8}".format(
epoch + 1, total_nmse_training, total_ssim_training,
total_psnr_training)
print(log)
log_all.debug(log)
log_eval.info(log)
# evaluation for validation data
total_nmse_val = 0
total_ssim_val = 0
total_psnr_val = 0
num_val_temp = 0
for batch in tl.iterate.minibatches(inputs=X_val_blurry,
targets=X_val_target,
batch_size=batch_size,
shuffle=False):
# x_good, _ = batch
# # x_bad = threading_data(x_good, fn=to_bad_img, mask=mask)
# x_bad = threading_data(
# x_good,
# fn=to_bad_img,
# mask=mask)
x_bad, x_good = batch
x_gen = sess.run(net_test.outputs, {t_image_bad: x_bad})
x_good_0_1 = (x_good + 1) / 2
x_gen_0_1 = (x_gen + 1) / 2
nmse_res = sess.run(nmse_0_1, {
t_gen: x_gen_0_1,
t_image_good: x_good_0_1
})
ssim_res = threading_data([_ for _ in zip(x_good_0_1, x_gen_0_1)],
fn=ssim)
psnr_res = threading_data([_ for _ in zip(x_good_0_1, x_gen_0_1)],
fn=psnr)
total_nmse_val += np.sum(nmse_res)
total_ssim_val += np.sum(ssim_res)
total_psnr_val += np.sum(psnr_res)
num_val_temp += batch_size
total_nmse_val /= num_val_temp
total_ssim_val /= num_val_temp
total_psnr_val /= num_val_temp
log = "Epoch: {}\nNMSE val: {:8}, SSIM val: {:8}, PSNR val: {:8}".format(
epoch + 1, total_nmse_val, total_ssim_val, total_psnr_val)
print(log)
log_all.debug(log)
log_eval.info(log)
img = sess.run(net_test_sample.outputs,
{t_image_bad_samples: X_samples_bad})
tl.visualize.save_images(
img, [5, 10], os.path.join(save_dir, "image_{}.png".format(epoch)))
if total_nmse_val < best_nmse:
esn = early_stopping_num # reset early stopping num
best_nmse = total_nmse_val
best_epoch = epoch + 1
# save current best model
tl.files.save_npz(
net.all_params,
name=os.path.join(checkpoint_dir, tl.global_flag['model']) +
'.npz',
sess=sess)
tl.files.save_npz(
net_d.all_params,
name=os.path.join(checkpoint_dir, tl.global_flag['model']) +
'_d.npz',
sess=sess)
print("[*] Save checkpoints SUCCESS!")
else:
esn -= 1
log = "Best NMSE result: {} at {} epoch".format(best_nmse, best_epoch)
log_eval.info(log)
log_all.debug(log)
print(log)
# early stopping triggered
if esn == 0:
log_eval.info(log)
tl.files.load_and_assign_npz(
sess=sess,
name=os.path.join(checkpoint_dir, tl.global_flag['model']) +
'.npz',
network=net)
# evluation for test data
x_gen = sess.run(net_test_sample.outputs,
{t_image_bad_samples: X_samples_bad})
x_gen_0_1 = (x_gen + 1) / 2
savemat(save_dir + '/test_random_50_generated.mat',
{'x_gen_0_1': x_gen_0_1})
nmse_res = sess.run(
nmse_0_1_sample, {
t_gen_sample: x_gen_0_1,
t_image_good_samples: x_good_sample_rescaled
})
ssim_res = threading_data(
[_ for _ in zip(x_good_sample_rescaled, x_gen_0_1)], fn=ssim)
psnr_res = threading_data(
[_ for _ in zip(x_good_sample_rescaled, x_gen_0_1)], fn=psnr)
log = "NMSE testing: {}\nSSIM testing: {}\nPSNR testing: {}\n\n".format(
nmse_res, ssim_res, psnr_res)
log_50.debug(log)
log = "NMSE testing average: {}\nSSIM testing average: {}\nPSNR testing average: {}\n\n".format(
np.mean(nmse_res), np.mean(ssim_res), np.mean(psnr_res))
log_50.debug(log)
log = "NMSE testing std: {}\nSSIM testing std: {}\nPSNR testing std: {}\n\n".format(
np.std(nmse_res), np.std(ssim_res), np.std(psnr_res))
log_50.debug(log)
# evaluation for zero-filled (ZF) data
nmse_res_zf = sess.run(
nmse_0_1_sample, {
t_gen_sample: x_bad_sample_rescaled,
t_image_good_samples: x_good_sample_rescaled
})
ssim_res_zf = threading_data([
_ for _ in zip(x_good_sample_rescaled, x_bad_sample_rescaled)
],
fn=ssim)
psnr_res_zf = threading_data([
_ for _ in zip(x_good_sample_rescaled, x_bad_sample_rescaled)
],
fn=psnr)
log = "NMSE ZF testing: {}\nSSIM ZF testing: {}\nPSNR ZF testing: {}\n\n".format(
nmse_res_zf, ssim_res_zf, psnr_res_zf)
log_50.debug(log)
log = "NMSE ZF average testing: {}\nSSIM ZF average testing: {}\nPSNR ZF average testing: {}\n\n".format(
np.mean(nmse_res_zf), np.mean(ssim_res_zf),
np.mean(psnr_res_zf))
log_50.debug(log)
log = "NMSE ZF std testing: {}\nSSIM ZF std testing: {}\nPSNR ZF std testing: {}\n\n".format(
np.std(nmse_res_zf), np.std(ssim_res_zf), np.std(psnr_res_zf))
log_50.debug(log)
# sample testing images
tl.visualize.save_images(
x_gen, [5, 10],
os.path.join(save_dir, "final_generated_image.png"))
tl.visualize.save_images(
np.clip(10 * np.abs(X_samples_good - x_gen) / 2, 0, 1),
[5, 10],
os.path.join(save_dir,
"final_generated_image_diff_abs_10_clip.png"))
tl.visualize.save_images(
np.clip(10 * np.abs(X_samples_good - X_samples_bad) / 2, 0, 1),
[5, 10],
os.path.join(save_dir, "final_bad_image_diff_abs_10_clip.png"))
print("[*] Job finished!")
break
# ==============================================================================
import multiprocessing
import os, sys, time
from config import config, log_config
import util
AGENT_COUNT = config["agent_config"]["count"]
EVALUATOR_COUNT = config["evaluator_config"]["count"]
MODEL_AUGMENTED = config["model_config"] is not False
if config["resume"]:
ROOT_PATH = "output/" + config["env"]["name"] + "/" + config["name"]
else:
ROOT_PATH = util.create_and_wipe_directory("output/" + config["env"]["name"] + "/" + config["name"])
log_config()
import learner, agent, valuerl_learner
if MODEL_AUGMENTED: import worldmodel_learner
if __name__ == '__main__':
all_procs = set([])
interaction_procs = set([])
# lock
policy_lock = multiprocessing.Lock()
model_lock = multiprocessing.Lock() if MODEL_AUGMENTED else None
# queue
policy_replay_frame_queue = multiprocessing.Queue(1)
model_replay_frame_queue = multiprocessing.Queue(1) if MODEL_AUGMENTED else None
def train():
## CREATE DIRECTORIES
mode_dir = config.TRAIN.root_dir + '{}'.format(tl.global_flag['mode'])
ckpt_dir = mode_dir + '/checkpoint'
init_dir = mode_dir + '/init'
log_dir_scalar_init = mode_dir + '/log/scalar_init'
log_dir_image_init = mode_dir + '/log/image_init'
log_dir_scalar = mode_dir + '/log/scalar'
log_dir_image = mode_dir + '/log/image'
sample_dir = mode_dir + '/samples/0_train'
config_dir = mode_dir + '/config'
if tl.global_flag['delete_log']:
shutil.rmtree(ckpt_dir, ignore_errors=True)
if tl.global_flag['is_pretrain']:
shutil.rmtree(log_dir_scalar_init, ignore_errors=True)
shutil.rmtree(log_dir_image_init, ignore_errors=True)
shutil.rmtree(log_dir_scalar, ignore_errors=True)
shutil.rmtree(log_dir_image, ignore_errors=True)
shutil.rmtree(sample_dir, ignore_errors=True)
shutil.rmtree(config_dir, ignore_errors=True)
tl.files.exists_or_mkdir(ckpt_dir)
tl.files.exists_or_mkdir(init_dir)
tl.files.exists_or_mkdir(log_dir_scalar_init)
tl.files.exists_or_mkdir(log_dir_image_init)
tl.files.exists_or_mkdir(log_dir_scalar)
tl.files.exists_or_mkdir(log_dir_image)
tl.files.exists_or_mkdir(sample_dir)
tl.files.exists_or_mkdir(config_dir)
log_config(config_dir, config)
## DEFINE SESSION
sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True,
log_device_placement=False))
## READ DATASET LIST
# train
train_real_img_list = np.array(
sorted(
tl.files.load_file_list(path=config.TRAIN.real_img_path,
regx='.*',
printable=False)))
train_real_binary_map_list = np.array(
sorted(
tl.files.load_file_list(path=config.TRAIN.real_binary_map_path,
regx='.*',
printable=False)))
train_real_img_no_label_list = np.array(
sorted(
tl.files.load_file_list(path=config.TRAIN.real_img_no_label_path,
regx='.*',
printable=False)))
# test
test_blur_img_list = np.array(
sorted(
tl.files.load_file_list(path=config.TEST.cuhk_img_path,
regx='.*',
printable=False)))
test_gt_list = np.array(
sorted(
tl.files.load_file_list(path=config.TEST.cuhk_binary_map_path,
regx='.*',
printable=False)))
test_blur_imgs = read_all_imgs(test_blur_img_list,
path=config.TEST.cuhk_img_path,
mode='RGB')
test_gt_imgs = read_all_imgs(test_gt_list,
path=config.TEST.cuhk_binary_map_path,
mode='GRAY')
## DEFINE MODEL
# input
with tf.variable_scope('input'):
patches_synthetic = tf.placeholder('float32', [None, h, w, 3],
name='input_synthetic')
labels_synthetic_defocus = tf.placeholder(
'float32', [None, h, w, 1], name='labels_synthetic_defocus')
patches_real = tf.placeholder('float32', [None, h, w, 3],
name='input_real')
labels_real_binary = tf.placeholder('float32', [None, h, w, 1],
name='labels_real_binary')
patches_real_no_label = tf.placeholder('float32', [None, h, w, 3],
name='input_real_no_label')
patches_real_test = tf.placeholder('float32', [None, h, w, 3],
name='input_real')
labels_real_binary_test = tf.placeholder('float32', [None, h, w, 1],
name='labels_real_binary')
# model
with tf.variable_scope('main_net') as scope:
with tf.variable_scope('defocus_net') as scope:
with tf.variable_scope('encoder') as scope:
net_vgg, feats_synthetic_down, _, _ = VGG19_down(
patches_synthetic, reuse=False, scope=scope)
_, feats_real_down, _, _ = VGG19_down(patches_real,
reuse=True,
scope=scope)
_, feats_real_no_label_down, _, _ = VGG19_down(
patches_real_no_label, reuse=True, scope=scope)
with tf.variable_scope('decoder') as scope:
output_synthetic_defocus, feats_synthetic_up_aux, feats_synthetic_da, _ = UNet_up(
patches_synthetic,
feats_synthetic_down,
is_train=True,
reuse=False,
scope=scope)
output_real_defocus, _, feats_real_da, _ = UNet_up(
patches_real,
feats_real_down,
is_train=True,
reuse=True,
scope=scope)
output_real_no_label_defocus, _, feats_real_no_label_da, _ = UNet_up(
patches_real_no_label,
feats_real_no_label_down,
is_train=True,
reuse=True,
scope=scope)
with tf.variable_scope('binary_net') as scope:
output_real_binary_logits, output_real_binary = Binary_Net(
output_real_defocus, is_train=True, reuse=False, scope=scope)
with tf.variable_scope('discriminator') as scope:
with tf.variable_scope('feature') as scope:
d_feature_logits_synthetic, d_feature_synthetic = feature_discriminator(
feats_synthetic_da, is_train=True, reuse=False, scope=scope)
d_feature_logits_real, d_feature_real = feature_discriminator(
feats_real_da, is_train=True, reuse=True, scope=scope)
d_feature_logits_real_no_label, d_feature_real_no_label = feature_discriminator(
feats_real_no_label_da, is_train=True, reuse=True, scope=scope)
with tf.variable_scope('perceptual') as scope:
output_synthetic_defocus_3c = tf.concat([
output_synthetic_defocus, output_synthetic_defocus,
output_synthetic_defocus
],
axis=3)
net_vgg_perceptual, _, perceptual_synthetic_out, logits_perceptual_out = VGG19_down(
output_synthetic_defocus_3c, reuse=False, scope=scope)
labels_synthetic_defocus_3c = tf.concat([
labels_synthetic_defocus, labels_synthetic_defocus,
labels_synthetic_defocus
],
axis=3)
_, _, perceptual_synthetic_label, logits_perceptual_label = VGG19_down(
labels_synthetic_defocus_3c, reuse=True, scope=scope)
# for test
with tf.variable_scope('main_net') as scope:
with tf.variable_scope('defocus_net') as scope:
with tf.variable_scope('encoder') as scope:
_, feats_real_down_test, _, _ = VGG19_down(patches_real_test,
reuse=True,
scope=scope)
with tf.variable_scope('decoder') as scope:
output_real_defocus_test, _, _, _ = UNet_up(
patches_real,
feats_real_down_test,
is_train=True,
reuse=True,
scope=scope)
## DEFINE LOSS
with tf.variable_scope('loss'):
with tf.variable_scope('discriminator'):
with tf.variable_scope('feature'):
loss_synthetic_d_feature = tl.cost.sigmoid_cross_entropy(
d_feature_logits_synthetic,
tf.ones_like(d_feature_logits_synthetic),
name='synthetic')
loss_real_d_feature = tl.cost.sigmoid_cross_entropy(
d_feature_logits_real,
tf.zeros_like(d_feature_logits_real),
name='real')
loss_real_no_label_d_feature = tl.cost.sigmoid_cross_entropy(
d_feature_logits_real_no_label,
tf.zeros_like(d_feature_logits_real_no_label),
name='real')
loss_d_feature = tf.identity(
(2 * loss_synthetic_d_feature + loss_real_d_feature +
loss_real_no_label_d_feature) / 2. * lambda_adv,
name='total')
loss_d = tf.identity(loss_d_feature, name='total')
with tf.variable_scope('generator'):
with tf.variable_scope('feature'):
loss_real_g_feature = tl.cost.sigmoid_cross_entropy(
d_feature_logits_real,
tf.ones_like(d_feature_logits_real),
name='real')
loss_real_no_label_g_feature = tl.cost.sigmoid_cross_entropy(
d_feature_logits_real_no_label,
tf.ones_like(d_feature_logits_real_no_label),
name='real')
loss_g_feature = tf.identity(
(loss_real_g_feature + loss_real_no_label_g_feature) / 2.,
name='total')
loss_g = tf.identity(loss_g_feature * lambda_adv, name='total')
with tf.variable_scope('defocus'):
loss_defocus = tl.cost.mean_squared_error(output_synthetic_defocus,
labels_synthetic_defocus,
is_mean=True,
name='synthetic')
with tf.variable_scope('auxilary'):
labels_layer = InputLayer(labels_synthetic_defocus)
loss_aux_1 = tl.cost.mean_squared_error(
feats_synthetic_up_aux[0],
DownSampling2dLayer(labels_layer, (int(h / 16), int(w / 16)),
is_scale=False,
method=1,
align_corners=True).outputs,
is_mean=True,
name='aux1')
loss_aux_2 = tl.cost.mean_squared_error(
feats_synthetic_up_aux[1],
DownSampling2dLayer(labels_layer, (int(h / 8), int(w / 8)),
is_scale=False,
method=1,
align_corners=True).outputs,
is_mean=True,
name='aux2')
loss_aux_3 = tl.cost.mean_squared_error(
feats_synthetic_up_aux[2],
DownSampling2dLayer(labels_layer, (int(h / 4), int(w / 4)),
is_scale=False,
method=1,
align_corners=True).outputs,
is_mean=True,
name='aux3')
loss_aux_4 = tl.cost.mean_squared_error(
feats_synthetic_up_aux[3],
DownSampling2dLayer(labels_layer, (int(h / 2), int(w / 2)),
is_scale=False,
method=1,
align_corners=True).outputs,
is_mean=True,
name='aux4')
loss_aux_5 = tl.cost.mean_squared_error(feats_synthetic_up_aux[4],
labels_synthetic_defocus,
is_mean=True,
name='aux5')
loss_aux = tf.identity(loss_aux_1 + loss_aux_2 + loss_aux_3 +
loss_aux_4 + loss_aux_5,
name='total')
with tf.variable_scope('perceptual'):
loss_synthetic_perceptual = tl.cost.mean_squared_error(
perceptual_synthetic_out,
perceptual_synthetic_label,
is_mean=True,
name='synthetic')
loss_perceptual = tf.identity(loss_synthetic_perceptual *
lambda_perceptual,
name='total')
with tf.variable_scope('binary'):
loss_real_binary = tl.cost.sigmoid_cross_entropy(
output_real_binary_logits, labels_real_binary, name='real')
loss_binary = tf.identity(loss_real_binary * lambda_binary,
name='total')
loss_main = tf.identity(loss_defocus + loss_binary + loss_perceptual +
loss_aux + loss_g,
name='total')
loss_init = tf.identity(loss_defocus, name='loss_init')
## DEFINE OPTIMIZER
# variables to save / train
d_vars = tl.layers.get_variables_with_name('discriminator', True, False)
main_vars = tl.layers.get_variables_with_name('main_net', True, False)
save_vars = tl.layers.get_variables_with_name(
'main_net', False, False) + tl.layers.get_variables_with_name(
'discriminator', False, False)
init_vars = tl.layers.get_variables_with_name('defocus_net', True, False)
save_init_vars = tl.layers.get_variables_with_name('defocus_net', False,
False)
# define optimizer
with tf.variable_scope('Optimizer'):
learning_rate = tf.Variable(lr_init, trainable=False)
learning_rate_init = tf.Variable(lr_init_init, trainable=False)
optim_d = tf.train.AdamOptimizer(learning_rate * lambda_lr_d,
beta1=beta1).minimize(loss_d,
var_list=d_vars)
optim_main = tf.train.AdamOptimizer(
learning_rate, beta1=beta1).minimize(loss_main, var_list=main_vars)
optim_init = tf.train.AdamOptimizer(
learning_rate_init, beta1=beta1).minimize(loss_init,
var_list=init_vars)
## DEFINE SUMMARY
# writer
writer_scalar = tf.summary.FileWriter(log_dir_scalar,
sess.graph,
flush_secs=30,
filename_suffix='.loss_log')
writer_image = tf.summary.FileWriter(log_dir_image,
sess.graph,
flush_secs=30,
filename_suffix='.image_log')
if tl.global_flag['is_pretrain']:
writer_scalar_init = tf.summary.FileWriter(
log_dir_scalar_init,
sess.graph,
flush_secs=30,
filename_suffix='.loss_log_init')
writer_image_init = tf.summary.FileWriter(
log_dir_image_init,
sess.graph,
flush_secs=30,
filename_suffix='.image_log_init')
# for pretrain
loss_sum_list_init = []
with tf.variable_scope('loss_init'):
loss_sum_list_init.append(
tf.summary.scalar('1_total_loss_init', loss_init))
loss_sum_list_init.append(
tf.summary.scalar('2_defocus_loss_init', loss_defocus))
loss_sum_init = tf.summary.merge(loss_sum_list_init)
image_sum_list_init = []
image_sum_list_init.append(
tf.summary.image('1_synthetic_input_init', patches_synthetic))
image_sum_list_init.append(
tf.summary.image('2_synthetic_defocus_out_init',
fix_image_tf(output_synthetic_defocus, 1)))
image_sum_list_init.append(
tf.summary.image('3_synthetic_defocus_gt_init',
fix_image_tf(labels_synthetic_defocus, 1)))
image_sum_init = tf.summary.merge(image_sum_list_init)
# for train
loss_sum_g_list = []
with tf.variable_scope('loss_generator'):
loss_sum_g_list.append(tf.summary.scalar('1_total', loss_main))
loss_sum_g_list.append(tf.summary.scalar('2_g', loss_g))
loss_sum_g_list.append(tf.summary.scalar('3_defocus', loss_defocus))
loss_sum_g_list.append(
tf.summary.scalar('4_perceptual', loss_perceptual))
loss_sum_g_list.append(tf.summary.scalar('5_auxilary', loss_aux))
loss_sum_g_list.append(tf.summary.scalar('6_binary', loss_binary))
loss_sum_g = tf.summary.merge(loss_sum_g_list)
loss_sum_d_list = []
with tf.variable_scope('loss_discriminator'):
loss_sum_d_list.append(tf.summary.scalar('1_d', loss_d_feature))
loss_sum_d = tf.summary.merge(loss_sum_d_list)
image_sum_list = []
image_sum_list.append(
tf.summary.image('1_synthetic_input', patches_synthetic))
image_sum_list.append(
tf.summary.image('2_synthetic_defocus_out',
fix_image_tf(output_synthetic_defocus, 1)))
image_sum_list.append(
tf.summary.image('3_synthetic_defocus_gt',
fix_image_tf(labels_synthetic_defocus, 1)))
image_sum_list.append(tf.summary.image('4_real_input', patches_real))
image_sum_list.append(
tf.summary.image('5_real_defocus_out',
fix_image_tf(output_real_defocus, 1)))
image_sum_list.append(
tf.summary.image('6_real_binary_out',
fix_image_tf(output_real_binary, 1)))
image_sum_list.append(
tf.summary.image('7_real_binary_gt',
fix_image_tf(labels_real_binary, 1)))
image_sum_list.append(
tf.summary.image('8_real_binary_gt_no_label', patches_real_no_label))
image_sum_list.append(
tf.summary.image('9_real_defocus_out_no_label',
fix_image_tf(output_real_no_label_defocus, 1)))
image_sum = tf.summary.merge(image_sum_list)
## INITIALIZE SESSION
sess.run(tf.global_variables_initializer())
## LOAD VGG
vgg19_npy_path = "pretrained/vgg19.npy"
if not os.path.isfile(vgg19_npy_path):
print(
"Please download vgg19.npz from : https://github.com/machrisaa/tensorflow-vgg"
)
exit()
npz = np.load(vgg19_npy_path, encoding='latin1').item()
params = []
for val in sorted(npz.items()):
if val[0] == 'fc6':
break
W = np.asarray(val[1][0])
b = np.asarray(val[1][1])
print(" Loading %s: %s, %s" % (val[0], W.shape, b.shape))
params.extend([W, b])
tl.files.assign_params(sess, params, net_vgg)
tl.files.assign_params(sess, params, net_vgg_perceptual)
tl.files.load_and_assign_npz_dict(
name=init_dir + '/{}_init.npz'.format(tl.global_flag['mode']),
sess=sess)
if tl.global_flag['is_pretrain']:
print('*****************************************')
print(' PRE-TRAINING START')
print('*****************************************')
global_step = 0
for epoch in range(0, n_epoch_init):
total_loss_init, n_iter = 0, 0
# reload image list
train_synthetic_img_list = np.array(
sorted(
tl.files.load_file_list(
path=config.TRAIN.synthetic_img_path,
regx='.*',
printable=False)))
train_defocus_map_list = np.array(
sorted(
tl.files.load_file_list(path=config.TRAIN.defocus_map_path,
regx='.*',
printable=False)))
# shuffle datasets
shuffle_index = np.arange(len(train_synthetic_img_list))
np.random.shuffle(shuffle_index)
train_synthetic_img_list = train_synthetic_img_list[shuffle_index]
train_defocus_map_list = train_defocus_map_list[shuffle_index]
epoch_time = time.time()
#for idx in range(0, len(train_synthetic_img_list), batch_size_init):
for idx in range(0, 10):
step_time = time.time()
## READ DATA
# read synthetic data
b_idx = (idx + np.arange(batch_size_init)
) % len(train_synthetic_img_list)
synthetic_images_blur = read_all_imgs(
train_synthetic_img_list[b_idx],
path=config.TRAIN.synthetic_img_path,
mode='RGB')
synthetic_defocus_maps = read_all_imgs(
train_defocus_map_list[b_idx],
path=config.TRAIN.defocus_map_path,
mode='DEPTH')
synthetic_images_blur, synthetic_defocus_maps = \
crop_pair_with_different_shape_images(synthetic_images_blur, synthetic_defocus_maps, [h, w], is_gaussian_noise = tl.global_flag['is_noise'])
err_init, lr, _ = \
sess.run([loss_init, learning_rate_init, optim_init], {patches_synthetic: synthetic_images_blur, labels_synthetic_defocus: synthetic_defocus_maps})
print('[%s] Ep [%2d/%2d] %4d/%4d time: %4.2fs, err_init: %1.2e, lr: %1.2e' % \
(tl.global_flag['mode'], epoch, n_epoch_init, n_iter, len(train_synthetic_img_list)/batch_size_init, time.time() - step_time, err_init, lr))
if global_step % config.TRAIN.write_log_every == 0:
summary_loss_init, summary_image_init = sess.run(
[loss_sum_init, image_sum_init], {
patches_synthetic: synthetic_images_blur,
labels_synthetic_defocus: synthetic_defocus_maps
})
writer_scalar_init.add_summary(summary_loss_init,
global_step)
writer_image_init.add_summary(summary_image_init,
global_step)
total_loss_init += err_init
n_iter += 1
global_step += 1
if epoch % config.TRAIN.refresh_image_log_every and epoch != n_epoch_init == 0:
writer_image_init.close()
remove_file_end_with(log_dir_image_init, '*.image_log')
writer_image_init.reopen()
if epoch % 2 or epoch == n_epoch_init - 1:
tl.files.save_npz_dict(
save_init_vars,
name=init_dir +
'/{}_init.npz'.format(tl.global_flag['mode']),
sess=sess)
writer_image_init.close()
writer_scalar_init.close()
## START TRAINING
print('*****************************************')
print(' TRAINING START')
print('*****************************************')
global_step = 0
for epoch in range(0, n_epoch + 1):
total_loss, n_iter = 0, 0
# reload synthetic datasets
train_synthetic_img_list = np.array(
sorted(
tl.files.load_file_list(path=config.TRAIN.synthetic_img_path,
regx='.*',
printable=False)))
train_defocus_map_list = np.array(
sorted(
tl.files.load_file_list(path=config.TRAIN.defocus_map_path,
regx='.*',
printable=False)))
# shuffle datasets
shuffle_index = np.arange(len(train_synthetic_img_list))
np.random.shuffle(shuffle_index)
train_synthetic_img_list = train_synthetic_img_list[shuffle_index]
train_defocus_map_list = train_defocus_map_list[shuffle_index]
shuffle_index = np.arange(len(train_real_img_list))
np.random.shuffle(shuffle_index)
train_real_img_list = train_real_img_list[shuffle_index]
train_real_binary_map_list = train_real_binary_map_list[shuffle_index]
shuffle_index = np.arange(len(train_real_img_no_label_list))
np.random.shuffle(shuffle_index)
train_real_img_no_label_list = train_real_img_no_label_list[
shuffle_index]
# update learning rate
if epoch != 0 and (epoch % decay_every == 0):
new_lr_decay = lr_decay**(epoch // decay_every)
sess.run(tf.assign(learning_rate, lr_init * new_lr_decay))
elif epoch == 0:
sess.run(tf.assign(learning_rate, lr_init))
epoch_time = time.time()
for idx in range(0, len(train_synthetic_img_list), batch_size):
step_time = time.time()
## READ DATA
# read synthetic data
b_idx = (idx +
np.arange(batch_size)) % len(train_synthetic_img_list)
synthetic_images_blur = read_all_imgs(
train_synthetic_img_list[b_idx],
path=config.TRAIN.synthetic_img_path,
mode='RGB')
synthetic_defocus_maps = read_all_imgs(
train_defocus_map_list[b_idx],
path=config.TRAIN.defocus_map_path,
mode='DEPTH')
synthetic_images_blur, synthetic_defocus_maps = crop_pair_with_different_shape_images(
synthetic_images_blur,
synthetic_defocus_maps, [h, w],
is_gaussian_noise=tl.global_flag['is_noise'])
# read real data #
b_idx = (idx % len(train_real_img_list) +
np.arange(batch_size)) % len(train_real_img_list)
real_images_blur = read_all_imgs(train_real_img_list[b_idx],
path=config.TRAIN.real_img_path,
mode='RGB')
real_binary_maps = read_all_imgs(
train_real_binary_map_list[b_idx],
path=config.TRAIN.real_binary_map_path,
mode='GRAY')
real_images_blur, real_binary_maps = crop_pair_with_different_shape_images(
real_images_blur, real_binary_maps, [h, w])
real_images_no_label_blur = read_all_imgs(
train_real_img_no_label_list[b_idx],
path=config.TRAIN.real_img_no_label_path,
mode='RGB')
real_images_no_label_blur = random_crop(real_images_no_label_blur,
[h, w])
## RUN NETWORK
#discriminator
feed_dict = {
patches_synthetic: synthetic_images_blur,
patches_real: real_images_blur,
labels_synthetic_defocus: synthetic_defocus_maps,
patches_real_no_label: real_images_no_label_blur
}
_ = sess.run(optim_d, feed_dict)
#generator
feed_dict = {
patches_synthetic: synthetic_images_blur,
labels_synthetic_defocus: synthetic_defocus_maps,
patches_real: real_images_blur,
labels_real_binary: real_binary_maps,
patches_real_no_label: real_images_no_label_blur
}
_ = sess.run(optim_main, feed_dict)
#log
err_main, err_g, err_d, d_synthetic, d_real, lr = \
sess.run([loss_main, loss_g, loss_d, d_feature_synthetic, d_feature_real, learning_rate], feed_dict)
d_acc = get_disc_accuracy([d_synthetic, d_real], [0, 1])
g_acc = get_disc_accuracy([d_synthetic, d_real], [1, 0])
print('[%s] Ep [%2d/%2d] %4d/%4d time: %4.2fs, err[main: %1.2e, g(acc): %1.2e(%1.2f), d(acc): %1.2e(%1.2f)], lr: %1.2e' % \
(tl.global_flag['mode'], epoch, n_epoch, n_iter, len(train_synthetic_img_list)/batch_size, time.time() - step_time, err_main, err_g, g_acc, err_d, d_acc, lr))
## SAVE LOGS
# save loss & image log
if global_step % config.TRAIN.write_log_every == 0:
summary_loss_g, summary_loss_d, summary_image = sess.run(
[loss_sum_g, loss_sum_d, image_sum], {
patches_synthetic: synthetic_images_blur,
labels_synthetic_defocus: synthetic_defocus_maps,
patches_real: real_images_blur,
labels_real_binary: real_binary_maps,
patches_real_no_label: real_images_no_label_blur
})
writer_scalar.add_summary(summary_loss_d, global_step)
writer_scalar.add_summary(summary_loss_g, global_step)
writer_image.add_summary(summary_image, global_step)
# save samples
# if global_step != 0 and global_step % config.TRAIN.write_sample_every == 0:
# synthetic_defocus_out, real_defocus_out, real_binary_out = sess.run([output_synthetic_defocus, output_real_defocus, output_real_binary], {patches_synthetic: synthetic_images_blur, patches_real: real_images_blur, labels_real_binary: real_binary_maps })
# save_images(synthetic_images_blur, [ni, ni], sample_dir + '/{}_{}_1_synthetic_input.png'.format(epoch, global_step))
# save_images(norm_image(synthetic_defocus_out), [ni, ni], sample_dir + '/{}_{}_2_synthetic_defocus_out_norm.png'.format(epoch, global_step))
# save_images(norm_image(synthetic_defocus_maps), [ni, ni], sample_dir + '/{}_{}_3_synthetic_defocus_gt.png'.format(epoch, global_step))
# save_images(real_images_blur, [ni, ni], sample_dir + '/{}_{}_4_real_input.png'.format(epoch, global_step))
# save_images(norm_image(real_defocus_out), [ni, ni], sample_dir + '/{}_{}_5_real_defocus_out_norm.png'.format(epoch, global_step))
# save_images(real_binary_out, [ni, ni], sample_dir + '/{}_{}_6_real_binary_out.png'.format(epoch, global_step))
# save_images(real_binary_maps, [ni, ni], sample_dir + '/{}_{}_7_real_binary_gt.png'.format(epoch, global_step))
total_loss += err_main
n_iter += 1
global_step += 1
print('[TRAIN] Epoch: [%2d/%2d] time: %4.4fs, total_err: %1.2e' %
(epoch, n_epoch, time.time() - epoch_time, total_loss / n_iter))
# reset image log
if epoch % config.TRAIN.refresh_image_log_every == 0:
writer_image.close()
remove_file_end_with(log_dir_image, '*.image_log')
writer_image.reopen()
if epoch % config.TRAIN.write_ckpt_every == 0:
#remove_file_end_with(ckpt_dir, '*.npz')
tl.files.save_npz_dict(
save_vars,
name=ckpt_dir +
'/{}_{}.npz'.format(tl.global_flag['mode'], epoch),
sess=sess)
def session_1():
log_config('./log_file.txt', config)
def config():
args = parse_args()
dir_config(args)
log_config(args, 'train')
return args
def train():
## CREATE DIRECTORIES
mode_dir = config.root_dir + '{}'.format(tl.global_flag['mode'])
ckpt_dir = mode_dir + '/checkpoint'
init_dir = mode_dir + '/init'
log_dir_scalar = mode_dir + '/log/scalar'
log_dir_image = mode_dir + '/log/image'
sample_dir = mode_dir + '/samples/0_train'
config_dir = mode_dir + '/config'
img_dir = mode_dir + '/img'
if tl.global_flag['delete_log']:
shutil.rmtree(ckpt_dir, ignore_errors = True)
shutil.rmtree(log_dir_scalar, ignore_errors = True)
shutil.rmtree(log_dir_image, ignore_errors = True)
shutil.rmtree(sample_dir, ignore_errors = True)
shutil.rmtree(config_dir, ignore_errors = True)
shutil.rmtree(img_dir, ignore_errors = True)
tl.files.exists_or_mkdir(ckpt_dir)
tl.files.exists_or_mkdir(log_dir_scalar)
tl.files.exists_or_mkdir(log_dir_image)
tl.files.exists_or_mkdir(sample_dir)
tl.files.exists_or_mkdir(config_dir)
tl.files.exists_or_mkdir(img_dir)
log_config(config_dir, config)
## DEFINE SESSION
sess = tf.Session(config = tf.ConfigProto(allow_soft_placement = True, log_device_placement = False))
## DEFINE MODEL
# input
with tf.variable_scope('input'):
patches_stab = tf.placeholder('float32', [batch_size, h, w, 3], name = 'input_frames_stab')
patches_unstab = tf.placeholder('float32', [batch_size, h, w, 3], name = 'input_frames_unstab')
inputs = {
'input_a': tf.image.resize_bilinear(patch_unstab,tf.constant([192,256])),
'input_b': tf.image.resize_bilinear(patch_stab,tf.constant([192,256]))
}
training_schedule = LONG_SCHEDULE
predictions = flownetS.model(inputs, training_schedule)
pred_flow = predictions['flow']
patch_warped = tf_warp(tf.image.resize_bilinear(patch_unstab,tf.constant([192,256])),-1*pred_flow,192,256)
## define loss
def masked_MSE(pred, gt, mask):
pred_mask = pred * mask
gt_mask = gt * mask
MSE = tf.reduce_sum(tf.squared_difference(pred_mask, gt_mask), axis = [1, 2, 3])
safe_mask = tf.cast(tf.where(tf.equal(mask, tf.zeros_like(mask)), mask + tf.constant(1e-8), mask), tf.float32)
MSE = MSE / tf.reduce_sum(safe_mask, axis = [1, 2, 3])
return tf.reduce_mean(MSE)
def lossterm(predict_flow,stab_image,unstab_image):
size = [predict_flow.shape[1], predict_flow.shape[2]]
downsampled_stab = tf.image.resize_images(stab_image, size )#downsample(gt_flow, size)
downsampled_unstab = tf.image.resize_images(unstab_image, size )#downsample(gt_flow, size)
warpedimg = tf_warp(downsampled_unstab, predict_flow, predict_flow.shape[1], predict_flow.shape[2])
#loss6 = tl.cost.mean_squared_error(downsampled_stab6, warpedimg6, is_mean = True, name = 'loss6') #* 0.32
mask = tf_warp(tf.ones_like(downsampled_unstab), predict_flow, predict_flow.shape[1], predict_flow.shape[2])
return masked_MSE(warpedimg, downsampled_stab, mask),warpedimg
with tf.variable_scope('loss'):
loss6,warpedimg6 = lossterm(-20 * predictions['predict_flow6'],patches_stab,patches_unstab)
loss5,warpedimg5 = lossterm(-20 * predictions['predict_flow5'],patches_stab,patches_unstab)
loss4,warpedimg4 = lossterm(-20 * predictions['predict_flow4'],patches_stab,patches_unstab)
loss3,warpedimg3 = lossterm(-20 * predictions['predict_flow3'],patches_stab,patches_unstab)
loss2,warpedimg2 = lossterm(-20 * predictions['predict_flow2'],patches_stab,patches_unstab)
loss1,warpedimg1 = lossterm(-1 * predictions['flow'],patches_stab,patches_unstab)
loss_main = tf.identity( loss6+loss5+loss4+loss3+loss2+loss1, name = 'total')
## DEFINE OPTIMIZER
# variables to save / train
main_vars = tl.layers.get_variables_with_name('FlowNetS', True, False)
save_vars = tl.layers.get_variables_with_name('FlowNetS', False, False)
# define optimizer
with tf.variable_scope('Optimizer'):
learning_rate = tf.Variable(lr_init, trainable = False)
optim_main = tf.train.AdamOptimizer(learning_rate, beta1 = beta1).minimize(loss_main, var_list = main_vars)
## DEFINE SUMMARY
# writer
writer_scalar = tf.summary.FileWriter(log_dir_scalar, sess.graph, flush_secs=30, filename_suffix = '.loss_log')
writer_image = tf.summary.FileWriter(log_dir_image, sess.graph, flush_secs=30, filename_suffix = '.image_log')
loss_sum_main_list = []
with tf.variable_scope('loss'):
loss_sum_main_list.append(tf.summary.scalar('loss6', loss6))
loss_sum_main_list.append(tf.summary.scalar('loss5', loss5))
loss_sum_main_list.append(tf.summary.scalar('loss4', loss4))
loss_sum_main_list.append(tf.summary.scalar('loss3', loss3))
loss_sum_main_list.append(tf.summary.scalar('loss2', loss2))
loss_sum_main_list.append(tf.summary.scalar('loss1', loss1))
loss_sum_main_list.append(tf.summary.scalar('loss_main', loss_main))
loss_sum_main = tf.summary.merge(loss_sum_main_list)
image_sum_list = []
image_sum_list.append(tf.summary.image('patch_unstab_source', fix_image_tf(patch_unstab, 1)))
image_sum_list.append(tf.summary.image('patch_stab_source', fix_image_tf(patch_stab, 1)))
image_sum_list.append(tf.summary.image('patch_warped', fix_image_tf(patch_warped, 1)))
image_sum_list.append(tf.summary.image('warpedimg6', fix_image_tf(warpedimg6, 1)))
image_sum_list.append(tf.summary.image('warpedimg5', fix_image_tf(warpedimg5, 1)))
image_sum_list.append(tf.summary.image('warpedimg4', fix_image_tf(warpedimg4, 1)))
image_sum_list.append(tf.summary.image('warpedimg3', fix_image_tf(warpedimg3, 1)))
image_sum_list.append(tf.summary.image('warpedimg2', fix_image_tf(warpedimg2, 1)))
image_sum = tf.summary.merge(image_sum_list)
## INITIALIZE SESSION
tl.layers.initialize_global_variables(sess)
flownetSaver = tf.train.Saver(tl.layers.get_variables_with_name('FlowNetS'))
flownetSaver.restore(sess, './flownet/checkpoints/FlowNetS/flownet-S.ckpt-0')
#tl.files.load_and_assign_npz_dict(name = '/Jarvis/logs/LJH/deep_video_stabilization/VS_addhomo_Limit_checkpoint/checkpoint/VS_addhomo_Limit.npz', sess = sess)
## START TRAINING
print '*****************************************'
print ' TRAINING START'
print '*****************************************'
global_step = 0
def train():
# ==================== setup config ==========================
parser = config_parser()
args = parser.parse_args()
setup_runtime(args)
log_config(args)
# ==================== create NeRF model =====================
output_ch = 5 if args.N_importance > 0 else 4
embed_pos, ch_pos = get_embedder(args.multires)
embed_dir, ch_dir = get_embedder(args.multires_views)
net_coarse = NeRF(layers=args.netdepth,
hidden_dims=args.netwidth,
input_ch=ch_pos,
input_ch_views=ch_dir,
output_ch=output_ch,
use_viewdirs=True)
net_fine = NeRF(layers=args.netdepth_fine,
hidden_dims=args.netwidth_fine,
input_ch=ch_pos,
input_ch_views=ch_dir,
output_ch=output_ch,
use_viewdirs=True)
params = list(net_coarse.parameters())
params += list(net_fine.parameters())
optimizer = torch.optim.Adam(params=params,
lr=args.lrate,
betas=(0.9, 0.999))
neural_renderer = Renderer(embed_pos,
embed_dir,
net_coarse,
net_fine,
cfg=args)
mem_coarse = cal_model_params(net_coarse)
mem_fine = cal_model_params(net_fine)
print(
f'memory usage: net_coarse:{mem_coarse:.4f} MB net_fine:{mem_fine:.4f} MB'
)
# ==================== load pretrained model =========================================================
start = 0
if args.checkpoint is not None:
start = load_checkpoint(args.checkpoint, net_coarse, net_fine,
optimizer)
log_dir = os.path.join(args.basedir, args.expname)
ckpts = [
os.path.join(log_dir, f) for f in sorted(os.listdir(log_dir))
if 'tar' in f
]
if len(ckpts) > 0:
print('Found checkpoints', ckpts[-1])
start = load_checkpoint(ckpts[-1], net_coarse, net_fine, optimizer)
# ==================== load data ========================================================================
images, poses, render_poses, hwf, i_split = load_blender_data(
args.datadir, args.half_res, args.testskip)
print('Loaded blender images={} render_poses={} intrinsics={}'.format(
images.shape, render_poses.shape, hwf))
i_train, i_val, i_test = i_split
# Cast intrinsics to right types
H, W, focal = hwf
H, W = int(H), int(W)
hwf = [H, W, focal]
if args.render_test:
render_poses = np.array(poses[i_test])
render_poses = torch.Tensor(render_poses).to(device)
images = torch.Tensor(images).to(device)
poses = torch.Tensor(poses).to(device)
# ==================== train ===========================================================================
global_step = start
for i in trange(start, args.num_iter):
img_i = np.random.choice(i_train)
target = images[img_i]
pose = poses[img_i, :3, :4]
rgb, disp, acc, extras = neural_renderer.render(H,
W,
focal,
c2w=pose,
target_img=target)
img_loss = mse(rgb, extras['target_rgb'])
loss = img_loss
psnr = mse2psnr(img_loss)
if 'rgb0' in extras:
img_loss0 = mse(extras['rgb0'], extras['target_rgb'])
loss = loss + img_loss0
psnr0 = mse2psnr(img_loss0)
optimizer.zero_grad()
loss.backward()
optimizer.step()
# update learning rate
decay_rate = 0.1
decay_steps = args.lrate_decay * 1000
new_lr = args.lrate * (decay_rate**(global_step / decay_steps))
for param_group in optimizer.param_groups:
param_group['lr'] = new_lr
if global_step % args.i_print == 0:
mem = torch.cuda.max_memory_cached() / (1024**2)
tqdm.write(
f"[TRAIN] iter{global_step}: loss:{loss.item()} PSNR:{psnr.item()} lr:{new_lr} mem:{mem} MB"
)
if global_step % args.i_weights == 0:
path = os.path.join(args.basedir, args.expname,
'{:06d}.tar'.format(i))
torch.save(
{
'global_step': global_step,
'net_coarse': net_coarse.state_dict(),
'net_fine': net_fine.state_dict(),
'optimizer': optimizer.state_dict()
}, path)
print('Saved checkpoint at', path)
if global_step % args.i_img == 0:
img_i = np.random.choice(i_val)
pose = poses[img_i, :3, :4]
with torch.no_grad():
rgb, disp, acc, extras = neural_renderer.render(H,
W,
focal,
c2w=pose)
rgb_img = to8b(rgb.cpu().numpy())
imageio.imwrite(
os.path.join(args.basedir, args.expname,
f"{global_step}.png"), rgb_img)
global_step += 1
def main_train():
mask_perc = tl.global_flag['maskperc']
mask_name = tl.global_flag['mask']
# =================================== BASIC CONFIGS =================================== #
print('[*] run basic configs ... ')
log_dir = "log_{}_{}".format(mask_name, mask_perc)
tl.files.exists_or_mkdir(log_dir)
log_all, log_eval, log_50, log_all_filename, log_eval_filename, log_50_filename = logging_setup(
log_dir)
save_dir = "samples_{}_{}".format(mask_name, mask_perc)
tl.files.exists_or_mkdir(save_dir)
# configs
batch_size = config.TRAIN.batch_size
early_stopping_num = config.TRAIN.early_stopping_num
g_alpha = config.TRAIN.g_alpha
g_beta = config.TRAIN.g_beta
g_gamma = config.TRAIN.g_gamma
g_adv = config.TRAIN.g_adv
lr = config.TRAIN.lr
lr_decay = config.TRAIN.lr_decay
decay_every = config.TRAIN.decay_every
n_epoch = config.TRAIN.n_epoch
sample_size = config.TRAIN.sample_size
log_config(log_all_filename, config)
log_config(log_eval_filename, config)
log_config(log_50_filename, config)
# ==================================== PREPARE DATA ==================================== #
print('[*] load data ... ')
training_data_path = config.TRAIN.training_data_path
val_data_path = config.TRAIN.val_data_path
testing_data_path = config.TRAIN.testing_data_path
with open(training_data_path, 'rb') as f:
X_train = pickle.load(f)
with open(val_data_path, 'rb') as f:
X_val = pickle.load(f)
with open(testing_data_path, 'rb') as f:
X_test = pickle.load(f)
print('[*] X_train shape/min/max: ', X_train.shape, X_train.min(),
X_train.max())
print('[*] X_val shape/min/max: ', X_val.shape, X_val.min(), X_val.max())
print('[*] X_test shape/min/max: ', X_test.shape, X_test.min(),
X_test.max())
print('[*] loading mask ... ')
if mask_name == "spiral":
mask = \
loadmat(
os.path.join(config.TRAIN.mask_Spiral_path, "spiral_{}.mat".format(mask_perc)))[
'spiral_mask']
elif mask_name == "radial":
mask = \
loadmat(
os.path.join(config.TRAIN.mask_Radial_path, "radial_{}.mat".format(mask_perc)))[
'radial_mask']
elif mask_name == "cartesian":
mask = \
loadmat(
os.path.join(config.TRAIN.mask_Cartesian_path, "cartesian_{}.mat".format(mask_perc)))[
'cartesian_mask']
elif mask_name == "random":
mask = \
loadmat(
os.path.join(config.TRAIN.mask_Random_path, "random_{}.mat".format(mask_perc)))[
'random_mask']
else:
raise ValueError("no such mask exists: {}".format(mask_name))
# ==================================== DEFINE MODEL ==================================== #
print('[*] define model ... ')
nw, nh, nz = X_train.shape[1:]
# define placeholders
t_image_good = tf.placeholder('float32', [batch_size, nw, nh, nz],
name='good_image')
t_image_good_samples = tf.placeholder('float32', [sample_size, nw, nh, nz],
name='good_image_samples')
t_image_bad = tf.placeholder('float32', [batch_size, nw, nh, nz],
name='bad_image')
t_image_bad_samples = tf.placeholder('float32', [sample_size, nw, nh, nz],
name='bad_image_samples')
t_gen = tf.placeholder('float32', [batch_size, nw, nh, nz],
name='generated_image_for_test')
t_gen_sample = tf.placeholder('float32', [sample_size, nw, nh, nz],
name='generated_sample_image_for_test')
# define generator network
net = generator(t_image_bad, is_train=True, reuse=False)
net_test = generator(t_image_bad, is_train=False, reuse=True)
net_test_sample = generator(t_image_bad_samples,
is_train=False,
reuse=True)
# define discriminator network
net_d, logits_fake = discriminator(net.outputs, is_train=True, reuse=False)
_, logits_real = discriminator(t_image_good, is_train=True, reuse=True)
# ==================================== DEFINE LOSS ==================================== #
print('[*] define loss functions ... ')
# discriminator loss
d_loss1 = tl.cost.sigmoid_cross_entropy(logits_real,
tf.ones_like(logits_real),
name='d1')
d_loss2 = tl.cost.sigmoid_cross_entropy(logits_fake,
tf.zeros_like(logits_fake),
name='d2')
d_loss = d_loss1 + d_loss2
# generator loss (adversarial)
g_adver = tl.cost.sigmoid_cross_entropy(logits_fake,
tf.ones_like(logits_fake),
name='g')
# generator loss (pixel-wise)
g_MAE = tf.reduce_mean(tf.abs(t_image_good - net.outputs), name='g_MAE')
g_ms_ssim = tf_ms_ssim(t_image_good, net.outputs)
# generator loss (L_grad)
good_dif1 = t_image_good[1:, :, :] - t_image_good[:-1, :, :]
good_dif2 = t_image_good[:, 1:, :] - t_image_good[:, :-1, :]
gen_dif1 = net.outputs[1:, :, :] - net.outputs[:-1, :, :]
gen_dif2 = net.outputs[:, 1:, :] - net.outputs[:, :-1, :]
dif1 = tf.reduce_mean(
tf.reduce_mean(tf.squared_difference(good_dif1, gen_dif1), axis=[1,
2]))
dif2 = tf.reduce_mean(
tf.reduce_mean(tf.squared_difference(good_dif2, gen_dif2), axis=[1,
2]))
g_grad = dif1 + dif2
# generator loss (total)
g_loss = g_adv * g_adver + g_alpha * g_MAE + g_beta * g_ms_ssim + g_gamma * g_grad
# nmse metric for testing purpose
nmse_a_0_1 = tf.sqrt(
tf.reduce_sum(tf.squared_difference(t_gen, t_image_good),
axis=[1, 2, 3]))
nmse_b_0_1 = tf.sqrt(tf.reduce_sum(tf.square(t_image_good), axis=[1, 2,
3]))
nmse_0_1 = nmse_a_0_1 / nmse_b_0_1
nmse_a_0_1_sample = tf.sqrt(
tf.reduce_sum(tf.squared_difference(t_gen_sample,
t_image_good_samples),
axis=[1, 2, 3]))
nmse_b_0_1_sample = tf.sqrt(
tf.reduce_sum(tf.square(t_image_good_samples), axis=[1, 2, 3]))
nmse_0_1_sample = nmse_a_0_1_sample / nmse_b_0_1_sample
# ==================================== DEFINE TRAIN OPTS ==================================== #
print('[*] define training options ... ')
g_vars = tl.layers.get_variables_with_name('Generator', True, True)
d_vars = tl.layers.get_variables_with_name('discriminator', True, True)
Total_parameters_g = count_trainable_params('Generator')
Total_parameters_d = count_trainable_params('discriminator')
Total_parameters = Total_parameters_g + Total_parameters_d
with tf.variable_scope('learning_rate'):
lr_v = tf.Variable(lr, trainable=False)
g_optim = tf.train.AdamOptimizer(lr_v).minimize(g_loss, var_list=g_vars)
d_optim = tf.train.AdamOptimizer(lr_v).minimize(d_loss, var_list=d_vars)
# ==================================== TRAINING ==================================== #
sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
sess.run(tf.global_variables_initializer())
print("g Total training params: %.2fM" % (Total_parameters_g / 1e6))
log_all.debug("g Total training params: %.2fM\n" %
(Total_parameters_g / 1e6))
print("d Total training params: %.2fM" % (Total_parameters_d / 1e6))
log_all.debug("d Total training params: %.2fM\n" %
(Total_parameters_d / 1e6))
print("Total training params: %.2fM" % (Total_parameters / 1e6))
log_all.debug("Total training params: %.2fM\n" % (Total_parameters / 1e6))
n_training_examples = len(X_train)
n_step_epoch = round(n_training_examples / batch_size)
# sample testing images
idex = tl.utils.get_random_int(min_v=0,
max_v=len(X_test) - 1,
number=sample_size,
seed=config.TRAIN.seed)
X_samples_good = X_test[idex]
X_samples_bad = threading_data(X_samples_good, fn=to_bad_img, mask=mask)
x_good_sample_rescaled = (X_samples_good + 1) / 2
x_bad_sample_rescaled = (X_samples_bad + 1) / 2
tl.visualize.save_images(X_samples_good, [5, 10],
os.path.join(save_dir, "sample_image_good.png"))
tl.visualize.save_images(X_samples_bad, [5, 10],
os.path.join(save_dir, "sample_image_bad.png"))
tl.visualize.save_images(threading_data(X_samples_good, fn=distort_img),
[5, 10],
os.path.join(save_dir, "sample_image_aug.png"))
scipy.misc.imsave(os.path.join(save_dir, "mask.png"), mask * 255)
print('[*] start training ... ')
best_nmse = np.inf
best_epoch = 1
esn = early_stopping_num
training_NMSE = []
training_PSNR = []
training_SSIM = []
val_NMSE = []
val_PSNR = []
val_SSIM = []
for epoch in range(0, n_epoch):
# learning rate decay
if epoch != 0 and (epoch % decay_every == 0):
new_lr_decay = lr_decay**(epoch // decay_every)
sess.run(tf.assign(lr_v, lr * new_lr_decay))
log = " ** new learning rate: %f" % (lr * new_lr_decay)
print(log)
log_all.debug(log)
elif epoch == 0:
log = " ** init lr: %f decay_every_epoch: %d, lr_decay: %f" % (
lr, decay_every, lr_decay)
print(log)
log_all.debug(log)
for step in range(n_step_epoch):
step_time = time.time()
idex = tl.utils.get_random_int(min_v=0,
max_v=n_training_examples - 1,
number=batch_size)
X_good = X_train[idex]
X_good_aug = threading_data(X_good, fn=distort_img)
X_bad = threading_data(X_good_aug, fn=to_bad_img, mask=mask)
errD, _ = sess.run([d_loss, d_optim], {
t_image_good: X_good_aug,
t_image_bad: X_bad
})
errG, errG_MAE, errG_msssim, errG_grad, _ = sess.run(
[g_loss, g_MAE, g_ms_ssim, g_grad, g_optim], {
t_image_good: X_good_aug,
t_image_bad: X_bad
})
log = "Epoch[{:3}/{:3}] step={:3} d_loss={:5} g_loss={:5} g_mae={:5} g_ms_ssim={:5} g_grad={:5} took {:3}s" \
.format(
epoch + 1,
n_epoch,
step,
round(float(errD), 3),
round(float(errG), 3),
round(float(errG_MAE), 3),
round(float(errG_msssim), 3),
round(float(errG_grad), 3),
round(time.time() - step_time, 2))
print(log)
log_all.debug(log)
# evaluation for training data
total_nmse_training = 0
total_ssim_training = 0
total_psnr_training = 0
num_training_temp = 0
for batch in tl.iterate.minibatches(inputs=X_train,
targets=X_train,
batch_size=batch_size,
shuffle=False):
x_good, _ = batch
# x_bad = threading_data(x_good, fn=to_bad_img, mask=mask)
x_bad = threading_data(x_good, fn=to_bad_img, mask=mask)
x_gen = sess.run(net_test.outputs, {t_image_bad: x_bad})
x_good_0_1 = (x_good + 1) / 2
x_gen_0_1 = (x_gen + 1) / 2
nmse_res = sess.run(nmse_0_1, {
t_gen: x_gen_0_1,
t_image_good: x_good_0_1
})
ssim_res = threading_data([_ for _ in zip(x_good_0_1, x_gen_0_1)],
fn=ssim)
psnr_res = threading_data([_ for _ in zip(x_good_0_1, x_gen_0_1)],
fn=psnr)
total_nmse_training += np.sum(nmse_res)
total_ssim_training += np.sum(ssim_res)
total_psnr_training += np.sum(psnr_res)
num_training_temp += batch_size
total_nmse_training /= num_training_temp
total_ssim_training /= num_training_temp
total_psnr_training /= num_training_temp
training_NMSE.append(total_nmse_training)
training_PSNR.append(total_psnr_training)
training_SSIM.append(total_ssim_training)
log = "Epoch: {}\nNMSE training: {:8}, SSIM training: {:8}, PSNR training: {:8}".format(
epoch + 1, total_nmse_training, total_ssim_training,
total_psnr_training)
print(log)
log_all.debug(log)
log_eval.info(log)
# evaluation for validation data
total_nmse_val = 0
total_ssim_val = 0
total_psnr_val = 0
num_val_temp = 0
for batch in tl.iterate.minibatches(inputs=X_val,
targets=X_val,
batch_size=batch_size,
shuffle=False):
x_good, _ = batch
# x_bad = threading_data(x_good, fn=to_bad_img, mask=mask)
x_bad = threading_data(x_good, fn=to_bad_img, mask=mask)
x_gen = sess.run(net_test.outputs, {t_image_bad: x_bad})
x_good_0_1 = (x_good + 1) / 2
x_gen_0_1 = (x_gen + 1) / 2
nmse_res = sess.run(nmse_0_1, {
t_gen: x_gen_0_1,
t_image_good: x_good_0_1
})
ssim_res = threading_data([_ for _ in zip(x_good_0_1, x_gen_0_1)],
fn=ssim)
psnr_res = threading_data([_ for _ in zip(x_good_0_1, x_gen_0_1)],
fn=psnr)
total_nmse_val += np.sum(nmse_res)
total_ssim_val += np.sum(ssim_res)
total_psnr_val += np.sum(psnr_res)
num_val_temp += batch_size
total_nmse_val /= num_val_temp
total_ssim_val /= num_val_temp
total_psnr_val /= num_val_temp
val_NMSE.append(total_nmse_val)
val_PSNR.append(total_psnr_val)
val_SSIM.append(total_ssim_val)
log = "Epoch: {}\nNMSE val: {:8}, SSIM val: {:8}, PSNR val: {:8}".format(
epoch + 1, total_nmse_val, total_ssim_val, total_psnr_val)
print(log)
log_all.debug(log)
log_eval.info(log)
img = sess.run(net_test_sample.outputs,
{t_image_bad_samples: X_samples_bad})
tl.visualize.save_images(
img, [5, 10],
os.path.join(save_dir, "image_{}.png".format(epoch + 1)))
if total_nmse_val < best_nmse:
esn = early_stopping_num # reset early stopping num
best_nmse = total_nmse_val
best_epoch = epoch + 1
tl.files.save_ckpt(sess=sess,
mode_name='model.ckpt',
save_dir='checkpoint')
print("[*] Save checkpoints SUCCESS!")
else:
esn -= 1
log = "Best NMSE result: {} at {} epoch".format(best_nmse, best_epoch)
log_eval.info(log)
log_all.debug(log)
print(log)
# early stopping triggered
if esn == 0:
log_eval.info(log)
log_eval.info("\ntraining_NMSE: {}".format(training_NMSE))
log_eval.info("\ntraining_PSNR: {}".format(training_PSNR))
log_eval.info("\ntraining_SSIM: {}".format(training_SSIM))
log_eval.info("\nval_NMSE: {}".format(val_NMSE))
log_eval.info("\nval_PSNR: {}".format(val_PSNR))
log_eval.info("\nval_SSIM: {}".format(val_SSIM))
tl.files.load_ckpt(sess=sess,
mode_name='model.ckpt',
save_dir='checkpoint')
# evluation for test data
test_start_time = time.time()
x_gen = sess.run(net_test_sample.outputs,
{t_image_bad_samples: X_samples_bad})
test_duration = (time.time() - test_start_time) / sample_size
x_gen_0_1 = (x_gen + 1) / 2
savemat(
save_dir + '/test_random_{}_generated.mat'.format(sample_size),
{'x_gen_0_1': x_gen_0_1})
nmse_res = sess.run(
nmse_0_1_sample, {
t_gen_sample: x_gen_0_1,
t_image_good_samples: x_good_sample_rescaled
})
ssim_res = threading_data(
[_ for _ in zip(x_good_sample_rescaled, x_gen_0_1)], fn=ssim)
psnr_res = threading_data(
[_ for _ in zip(x_good_sample_rescaled, x_gen_0_1)], fn=psnr)
log = "NMSE testing: {}\nSSIM testing: {}\nPSNR testing: {}\n\n".format(
nmse_res, ssim_res, psnr_res)
log_50.debug(log)
log = "NMSE testing average: {}\nSSIM testing average: {}\nPSNR testing average: {}\n\n".format(
np.mean(nmse_res), np.mean(ssim_res), np.mean(psnr_res))
log_50.debug(log)
log = "NMSE testing std: {}\nSSIM testing std: {}\nPSNR testing std: {}\n\n".format(
np.std(nmse_res), np.std(ssim_res), np.std(psnr_res))
log_50.debug(log)
log = "\nAverage test time: {}\n".format(test_duration)
log_50.debug(log)
# evaluation for zero-filled (ZF) data
nmse_res_zf = sess.run(
nmse_0_1_sample, {
t_gen_sample: x_bad_sample_rescaled,
t_image_good_samples: x_good_sample_rescaled
})
ssim_res_zf = threading_data([
_ for _ in zip(x_good_sample_rescaled, x_bad_sample_rescaled)
],
fn=ssim)
psnr_res_zf = threading_data([
_ for _ in zip(x_good_sample_rescaled, x_bad_sample_rescaled)
],
fn=psnr)
log = "NMSE ZF testing: {}\nSSIM ZF testing: {}\nPSNR ZF testing: {}\n\n".format(
nmse_res_zf, ssim_res_zf, psnr_res_zf)
log_50.debug(log)
log = "NMSE ZF average testing: {}\nSSIM ZF average testing: {}\nPSNR ZF average testing: {}\n\n".format(
np.mean(nmse_res_zf), np.mean(ssim_res_zf),
np.mean(psnr_res_zf))
log_50.debug(log)
log = "NMSE ZF std testing: {}\nSSIM ZF std testing: {}\nPSNR ZF std testing: {}\n\n".format(
np.std(nmse_res_zf), np.std(ssim_res_zf), np.std(psnr_res_zf))
log_50.debug(log)
# sample testing images
tl.visualize.save_images(
x_gen, [5, 10],
os.path.join(save_dir, "final_generated_image.png"))
print("[*] Job finished!")
break
def main_train():
mask_perc = tl.global_flag['maskperc']
mask_name = tl.global_flag['mask']
model_name = tl.global_flag['model']
# =================================== BASIC CONFIGS =================================== #
print('[*] run basic configs ... ')
log_dir = "log_{}_{}_{}".format(model_name, mask_name, mask_perc)
tl.files.exists_or_mkdir(log_dir)
log_all, log_eval, log_50, log_all_filename, log_eval_filename, log_50_filename = logging_setup(log_dir)
checkpoint_dir = "checkpoint_{}_{}_{}".format(model_name, mask_name, mask_perc)
tl.files.exists_or_mkdir(checkpoint_dir)
save_dir = "samples_{}_{}_{}".format(model_name, mask_name, mask_perc)
tl.files.exists_or_mkdir(save_dir)
# configs
batch_size = config.TRAIN.batch_size
early_stopping_num = config.TRAIN.early_stopping_num
g_alpha = config.TRAIN.g_alpha
g_beta = config.TRAIN.g_beta
g_gamma = config.TRAIN.g_gamma
g_adv = config.TRAIN.g_adv
lr = config.TRAIN.lr
lr_decay = config.TRAIN.lr_decay
decay_every = config.TRAIN.decay_every
beta1 = config.TRAIN.beta1
n_epoch = config.TRAIN.n_epoch
sample_size = config.TRAIN.sample_size
log_config(log_all_filename, config)
log_config(log_eval_filename, config)
log_config(log_50_filename, config)
# ==================================== PREPARE DATA ==================================== #
print('[*] load data ... ')
training_data_path = config.TRAIN.training_data_path
val_data_path = config.TRAIN.val_data_path
testing_data_path = config.TRAIN.testing_data_path
with open(training_data_path, 'rb') as f:
X_train = pickle.load(f)
with open(val_data_path, 'rb') as f:
X_val = pickle.load(f)
with open(testing_data_path, 'rb') as f:
X_test = pickle.load(f)
print('X_train shape/min/max: ', X_train.shape, X_train.min(), X_train.max())
print('X_val shape/min/max: ', X_val.shape, X_val.min(), X_val.max())
print('X_test shape/min/max: ', X_test.shape, X_test.min(), X_test.max())
print('[*] loading mask ... ')
if mask_name == "gaussian2d":
mask = \
loadmat(
os.path.join(config.TRAIN.mask_Gaussian2D_path, "GaussianDistribution2DMask_{}.mat".format(mask_perc)))[
'maskRS2']
elif mask_name == "gaussian1d":
mask = \
loadmat(
os.path.join(config.TRAIN.mask_Gaussian1D_path, "GaussianDistribution1DMask_{}.mat".format(mask_perc)))[
'maskRS1']
elif mask_name == "poisson2d":
mask = \
loadmat(
os.path.join(config.TRAIN.mask_Gaussian1D_path, "PoissonDistributionMask_{}.mat".format(mask_perc)))[
'population_matrix']
else:
raise ValueError("no such mask exists: {}".format(mask_name))
# ==================================== DEFINE MODEL ==================================== #
print('[*] define model ... ')
nw, nh, nz = X_train.shape[1:]
# define placeholders
t_image_good = tf.placeholder('float32', [batch_size, nw, nh, nz], name='good_image')
t_image_good_samples = tf.placeholder('float32', [sample_size, nw, nh, nz], name='good_image_samples')
t_image_bad = tf.placeholder('float32', [batch_size, nw, nh, nz], name='bad_image')
t_image_bad_samples = tf.placeholder('float32', [sample_size, nw, nh, nz], name='bad_image_samples')
t_gen = tf.placeholder('float32', [batch_size, nw, nh, nz], name='generated_image_for_test')
t_gen_sample = tf.placeholder('float32', [sample_size, nw, nh, nz], name='generated_sample_image_for_test')
t_image_good_244 = tf.placeholder('float32', [batch_size, 244, 244, 3], name='vgg_good_image')
# define generator network
if tl.global_flag['model'] == 'unet':
net = u_net_bn(t_image_bad, is_train=True, reuse=False, is_refine=False)
net_test = u_net_bn(t_image_bad, is_train=False, reuse=True, is_refine=False)
net_test_sample = u_net_bn(t_image_bad_samples, is_train=False, reuse=True, is_refine=False)
elif tl.global_flag['model'] == 'unet_refine':
net = u_net_bn(t_image_bad, is_train=True, reuse=False, is_refine=True)
net_test = u_net_bn(t_image_bad, is_train=False, reuse=True, is_refine=True)
net_test_sample = u_net_bn(t_image_bad_samples, is_train=False, reuse=True, is_refine=True)
else:
raise Exception("unknown model")
# define discriminator network
net_d, logits_fake = discriminator(net.outputs, is_train=True, reuse=False)
_, logits_real = discriminator(t_image_good, is_train=True, reuse=True)
# define VGG network
net_vgg_conv4_good, _ = vgg16_cnn_emb(t_image_good_244, reuse=False)
net_vgg_conv4_gen, _ = vgg16_cnn_emb(tf.tile(tf.image.resize_images(net.outputs, [244, 244]), [1, 1, 1, 3]), reuse=True)
# ==================================== DEFINE LOSS ==================================== #
print('[*] define loss functions ... ')
# discriminator loss
d_loss1 = tl.cost.sigmoid_cross_entropy(logits_real, tf.ones_like(logits_real), name='d1')
d_loss2 = tl.cost.sigmoid_cross_entropy(logits_fake, tf.zeros_like(logits_fake), name='d2')
d_loss = d_loss1 + d_loss2
# generator loss (adversarial)
g_loss = tl.cost.sigmoid_cross_entropy(logits_fake, tf.ones_like(logits_fake), name='g')
# generator loss (perceptual)
g_perceptual = tf.reduce_mean(tf.reduce_mean(tf.squared_difference(
net_vgg_conv4_good.outputs,
net_vgg_conv4_gen.outputs),
axis=[1, 2, 3]))
# generator loss (pixel-wise)
g_nmse_a = tf.sqrt(tf.reduce_sum(tf.squared_difference(net.outputs, t_image_good), axis=[1, 2, 3]))
g_nmse_b = tf.sqrt(tf.reduce_sum(tf.square(t_image_good), axis=[1, 2, 3]))
g_nmse = tf.reduce_mean(g_nmse_a / g_nmse_b)
# generator loss (frequency)
fft_good_abs = tf.map_fn(fft_abs_for_map_fn, t_image_good)
fft_gen_abs = tf.map_fn(fft_abs_for_map_fn, net.outputs)
g_fft = tf.reduce_mean(tf.reduce_mean(tf.squared_difference(fft_good_abs, fft_gen_abs), axis=[1, 2]))
# generator loss (total)
g_loss = g_adv * g_loss + g_alpha * g_nmse + g_gamma * g_perceptual + g_beta * g_fft
# nmse metric for testing purpose
nmse_a_0_1 = tf.sqrt(tf.reduce_sum(tf.squared_difference(t_gen, t_image_good), axis=[1, 2, 3]))
nmse_b_0_1 = tf.sqrt(tf.reduce_sum(tf.square(t_image_good), axis=[1, 2, 3]))
nmse_0_1 = nmse_a_0_1 / nmse_b_0_1
nmse_a_0_1_sample = tf.sqrt(tf.reduce_sum(tf.squared_difference(t_gen_sample, t_image_good_samples), axis=[1, 2, 3]))
nmse_b_0_1_sample = tf.sqrt(tf.reduce_sum(tf.square(t_image_good_samples), axis=[1, 2, 3]))
nmse_0_1_sample = nmse_a_0_1_sample / nmse_b_0_1_sample
# ==================================== DEFINE TRAIN OPTS ==================================== #
print('[*] define training options ... ')
g_vars = tl.layers.get_variables_with_name('u_net', True, True)
d_vars = tl.layers.get_variables_with_name('discriminator', True, True)
with tf.variable_scope('learning_rate'):
lr_v = tf.Variable(lr, trainable=False)
g_optim = tf.train.AdamOptimizer(lr_v, beta1=beta1).minimize(g_loss, var_list=g_vars)
d_optim = tf.train.AdamOptimizer(lr_v, beta1=beta1).minimize(d_loss, var_list=d_vars)
# ==================================== TRAINING ==================================== #
sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
tl.layers.initialize_global_variables(sess)
# load generator and discriminator weights (for continuous training purpose)
tl.files.load_and_assign_npz(sess=sess,
name=os.path.join(checkpoint_dir, tl.global_flag['model']) + '.npz',
network=net)
tl.files.load_and_assign_npz(sess=sess,
name=os.path.join(checkpoint_dir, tl.global_flag['model']) + '_d.npz',
network=net_d)
# load vgg weights
net_vgg_conv4_path = config.TRAIN.VGG16_path
npz = np.load(net_vgg_conv4_path)
assign_op = []
for idx, val in enumerate(sorted(npz.items())[0:20]):
print(" Loading pretrained VGG16, CNN part %s" % str(val[1].shape))
assign_op.append(net_vgg_conv4_good.all_params[idx].assign(val[1]))
sess.run(assign_op)
net_vgg_conv4_good.print_params(False)
n_training_examples = len(X_train)
n_step_epoch = round(n_training_examples / batch_size)
# sample testing images
idex = tl.utils.get_random_int(min=0, max=len(X_test) - 1, number=sample_size, seed=config.TRAIN.seed)
X_samples_good = X_test[idex]
X_samples_bad = threading_data(X_samples_good, fn=to_bad_img, mask=mask)
x_good_sample_rescaled = (X_samples_good + 1) / 2
x_bad_sample_rescaled = (X_samples_bad + 1) / 2
tl.visualize.save_images(X_samples_good,
[5, 10],
os.path.join(save_dir, "sample_image_good.png"))
tl.visualize.save_images(X_samples_bad,
[5, 10],
os.path.join(save_dir, "sample_image_bad.png"))
tl.visualize.save_images(np.abs(X_samples_good - X_samples_bad),
[5, 10],
os.path.join(save_dir, "sample_image_diff_abs.png"))
tl.visualize.save_images(np.sqrt(np.abs(X_samples_good - X_samples_bad) / 2 + config.TRAIN.epsilon),
[5, 10],
os.path.join(save_dir, "sample_image_diff_sqrt_abs.png"))
tl.visualize.save_images(np.clip(10 * np.abs(X_samples_good - X_samples_bad) / 2, 0, 1),
[5, 10],
os.path.join(save_dir, "sample_image_diff_sqrt_abs_10_clip.png"))
tl.visualize.save_images(threading_data(X_samples_good, fn=distort_img),
[5, 10],
os.path.join(save_dir, "sample_image_aug.png"))
scipy.misc.imsave(os.path.join(save_dir, "mask.png"), mask * 255)
print('[*] start training ... ')
best_nmse = np.inf
best_epoch = 1
esn = early_stopping_num
for epoch in range(0, n_epoch):
# learning rate decay
if epoch != 0 and (epoch % decay_every == 0):
new_lr_decay = lr_decay ** (epoch // decay_every)
sess.run(tf.assign(lr_v, lr * new_lr_decay))
log = " ** new learning rate: %f" % (lr * new_lr_decay)
print(log)
log_all.debug(log)
elif epoch == 0:
log = " ** init lr: %f decay_every_epoch: %d, lr_decay: %f" % (lr, decay_every, lr_decay)
print(log)
log_all.debug(log)
for step in range(n_step_epoch):
step_time = time.time()
idex = tl.utils.get_random_int(min=0, max=n_training_examples - 1, number=batch_size)
X_good = X_train[idex]
X_good_aug = threading_data(X_good, fn=distort_img)
X_good_244 = threading_data(X_good_aug, fn=vgg_prepro)
X_bad = threading_data(X_good_aug, fn=to_bad_img, mask=mask)
errD, _ = sess.run([d_loss, d_optim], {t_image_good: X_good_aug, t_image_bad: X_bad})
errG, errG_perceptual, errG_nmse, errG_fft, _ = sess.run([g_loss, g_perceptual, g_nmse, g_fft, g_optim],
{t_image_good_244: X_good_244,
t_image_good: X_good_aug,
t_image_bad: X_bad})
log = "Epoch[{:3}/{:3}] step={:3} d_loss={:5} g_loss={:5} g_perceptual_loss={:5} g_mse={:5} g_freq={:5} took {:3}s".format(
epoch + 1,
n_epoch,
step,
round(float(errD), 3),
round(float(errG), 3),
round(float(errG_perceptual), 3),
round(float(errG_nmse), 3),
round(float(errG_fft), 3),
round(time.time() - step_time, 2))
print(log)
log_all.debug(log)
# evaluation for training data
total_nmse_training = 0
total_ssim_training = 0
total_psnr_training = 0
num_training_temp = 0
for batch in tl.iterate.minibatches(inputs=X_train, targets=X_train, batch_size=batch_size, shuffle=False):
x_good, _ = batch
# x_bad = threading_data(x_good, fn=to_bad_img, mask=mask)
x_bad = threading_data(
x_good,
fn=to_bad_img,
mask=mask)
x_gen = sess.run(net_test.outputs, {t_image_bad: x_bad})
x_good_0_1 = (x_good + 1) / 2
x_gen_0_1 = (x_gen + 1) / 2
nmse_res = sess.run(nmse_0_1, {t_gen: x_gen_0_1, t_image_good: x_good_0_1})
ssim_res = threading_data([_ for _ in zip(x_good_0_1, x_gen_0_1)], fn=ssim)
psnr_res = threading_data([_ for _ in zip(x_good_0_1, x_gen_0_1)], fn=psnr)
total_nmse_training += np.sum(nmse_res)
total_ssim_training += np.sum(ssim_res)
total_psnr_training += np.sum(psnr_res)
num_training_temp += batch_size
total_nmse_training /= num_training_temp
total_ssim_training /= num_training_temp
total_psnr_training /= num_training_temp
log = "Epoch: {}\nNMSE training: {:8}, SSIM training: {:8}, PSNR training: {:8}".format(
epoch + 1,
total_nmse_training,
total_ssim_training,
total_psnr_training)
print(log)
log_all.debug(log)
log_eval.info(log)
# evaluation for validation data
total_nmse_val = 0
total_ssim_val = 0
total_psnr_val = 0
num_val_temp = 0
for batch in tl.iterate.minibatches(inputs=X_val, targets=X_val, batch_size=batch_size, shuffle=False):
x_good, _ = batch
# x_bad = threading_data(x_good, fn=to_bad_img, mask=mask)
x_bad = threading_data(
x_good,
fn=to_bad_img,
mask=mask)
x_gen = sess.run(net_test.outputs, {t_image_bad: x_bad})
x_good_0_1 = (x_good + 1) / 2
x_gen_0_1 = (x_gen + 1) / 2
nmse_res = sess.run(nmse_0_1, {t_gen: x_gen_0_1, t_image_good: x_good_0_1})
ssim_res = threading_data([_ for _ in zip(x_good_0_1, x_gen_0_1)], fn=ssim)
psnr_res = threading_data([_ for _ in zip(x_good_0_1, x_gen_0_1)], fn=psnr)
total_nmse_val += np.sum(nmse_res)
total_ssim_val += np.sum(ssim_res)
total_psnr_val += np.sum(psnr_res)
num_val_temp += batch_size
total_nmse_val /= num_val_temp
total_ssim_val /= num_val_temp
total_psnr_val /= num_val_temp
log = "Epoch: {}\nNMSE val: {:8}, SSIM val: {:8}, PSNR val: {:8}".format(
epoch + 1,
total_nmse_val,
total_ssim_val,
total_psnr_val)
print(log)
log_all.debug(log)
log_eval.info(log)
img = sess.run(net_test_sample.outputs, {t_image_bad_samples: X_samples_bad})
tl.visualize.save_images(img,
[5, 10],
os.path.join(save_dir, "image_{}.png".format(epoch)))
if total_nmse_val < best_nmse:
esn = early_stopping_num # reset early stopping num
best_nmse = total_nmse_val
best_epoch = epoch + 1
# save current best model
tl.files.save_npz(net.all_params,
name=os.path.join(checkpoint_dir, tl.global_flag['model']) + '.npz',
sess=sess)
tl.files.save_npz(net_d.all_params,
name=os.path.join(checkpoint_dir, tl.global_flag['model']) + '_d.npz',
sess=sess)
print("[*] Save checkpoints SUCCESS!")
else:
esn -= 1
log = "Best NMSE result: {} at {} epoch".format(best_nmse, best_epoch)
log_eval.info(log)
log_all.debug(log)
print(log)
# early stopping triggered
if esn == 0:
log_eval.info(log)
tl.files.load_and_assign_npz(sess=sess,
name=os.path.join(checkpoint_dir, tl.global_flag['model']) + '.npz',
network=net)
# evluation for test data
x_gen = sess.run(net_test_sample.outputs, {t_image_bad_samples: X_samples_bad})
x_gen_0_1 = (x_gen + 1) / 2
savemat(save_dir + '/test_random_50_generated.mat', {'x_gen_0_1': x_gen_0_1})
nmse_res = sess.run(nmse_0_1_sample, {t_gen_sample: x_gen_0_1, t_image_good_samples: x_good_sample_rescaled})
ssim_res = threading_data([_ for _ in zip(x_good_sample_rescaled, x_gen_0_1)], fn=ssim)
psnr_res = threading_data([_ for _ in zip(x_good_sample_rescaled, x_gen_0_1)], fn=psnr)
log = "NMSE testing: {}\nSSIM testing: {}\nPSNR testing: {}\n\n".format(
nmse_res,
ssim_res,
psnr_res)
log_50.debug(log)
log = "NMSE testing average: {}\nSSIM testing average: {}\nPSNR testing average: {}\n\n".format(
np.mean(nmse_res),
np.mean(ssim_res),
np.mean(psnr_res))
log_50.debug(log)
log = "NMSE testing std: {}\nSSIM testing std: {}\nPSNR testing std: {}\n\n".format(np.std(nmse_res),
np.std(ssim_res),
np.std(psnr_res))
log_50.debug(log)
# evaluation for zero-filled (ZF) data
nmse_res_zf = sess.run(nmse_0_1_sample,
{t_gen_sample: x_bad_sample_rescaled, t_image_good_samples: x_good_sample_rescaled})
ssim_res_zf = threading_data([_ for _ in zip(x_good_sample_rescaled, x_bad_sample_rescaled)], fn=ssim)
psnr_res_zf = threading_data([_ for _ in zip(x_good_sample_rescaled, x_bad_sample_rescaled)], fn=psnr)
log = "NMSE ZF testing: {}\nSSIM ZF testing: {}\nPSNR ZF testing: {}\n\n".format(
nmse_res_zf,
ssim_res_zf,
psnr_res_zf)
log_50.debug(log)
log = "NMSE ZF average testing: {}\nSSIM ZF average testing: {}\nPSNR ZF average testing: {}\n\n".format(
np.mean(nmse_res_zf),
np.mean(ssim_res_zf),
np.mean(psnr_res_zf))
log_50.debug(log)
log = "NMSE ZF std testing: {}\nSSIM ZF std testing: {}\nPSNR ZF std testing: {}\n\n".format(
np.std(nmse_res_zf),
np.std(ssim_res_zf),
np.std(psnr_res_zf))
log_50.debug(log)
# sample testing images
tl.visualize.save_images(x_gen,
[5, 10],
os.path.join(save_dir, "final_generated_image.png"))
tl.visualize.save_images(np.clip(10 * np.abs(X_samples_good - x_gen) / 2, 0, 1),
[5, 10],
os.path.join(save_dir, "final_generated_image_diff_abs_10_clip.png"))
tl.visualize.save_images(np.clip(10 * np.abs(X_samples_good - X_samples_bad) / 2, 0, 1),
[5, 10],
os.path.join(save_dir, "final_bad_image_diff_abs_10_clip.png"))
print("[*] Job finished!")
break
config.TRAIN.PRETRAIN.loss_applied, args.pcoef_high
) if args.pcoef_high is not None else config.TRAIN.coef_high
config.TRAIN.PRETRAIN.coef_init = get_dict_with_list(
config.TRAIN.PRETRAIN.loss_applied, args.pcoef_init
) if args.pcoef_init is not None else config.TRAIN.coef_init
config.TRAIN.max_ckpt_num = args.max_ckpt_num
print(toWhite('============== CONFIG =============='))
print_config(config)
print(toWhite('Creating log directories...'))
handle_directory(config.TRAIN, config.delete_log)
print(toWhite('Saving config...'))
log_config(config.TRAIN.LOG_DIR.config, config)
else:
config.EVAL.skip_length = args.eval_skip_length
config.EVAL.load_ckpt_by_score = t_or_f(args.load_ckpt_by_score)
config.EVAL.eval_mode = args.eval_mode
print(toWhite('Creating log directories...'))
handle_directory(config.EVAL, False)
tl.logging.set_verbosity(tl.logging.DEBUG)
#tl.logging.set_verbosity(tl.logging.INFO)
if config.is_train:
print(toYellow('\n[TRAINING {}]\n'.format(config.mode)))
train(config.TRAIN, config.mode)
else:
print(toYellow('\n[TESTING {}]\n'.format(config.mode)))
