def run_test(self, model_index=0):
# Create a session
gpu_options = tf.GPUOptions(
allow_growth=True
) # Does not pre-allocate large, increase if needed
config = tf.ConfigProto(
allow_soft_placement=True, gpu_options=gpu_options
) # soft_placement allows to work on CPUs if GPUs are not available
sess = tf.Session(config=config)
# Saver
log_name = args.loss_type
summary_writer = tf.summary.FileWriter(args.log_dir, sess.graph)
train_saver = tf.train.Saver(max_to_keep=20) # Keep maximum 20 models
# Initialize
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
# Threads coordinator
coordinator = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coordinator)
# Restore
print(args.model_dir)
train_saver.restore(sess, tf.train.latest_checkpoint(args.model_dir))
# Index of the image want to display
index = 0
h_total_loss_value = 0
rec_total_loss_value = 0
ssim_total_loss_value = 0
l1_total_loss_value = 0
total_num_fail_value = 0
# Start testing
h_losses_array = []
total_num_fail = 0
for step in range(self.num_total_steps):
num_fail_value, h_loss_value, rec_loss_value, ssim_loss_value, l1_loss_value, l1_smooth_loss_value, pred_I2_value, I1_aug_value, I2_aug_value, I_value, I_prime_value, pts1_value, gt_value, pred_h4p_value = sess.run(
[
self.total_num_fail, self.total_h_loss,
self.total_rec_loss, self.total_ssim_loss,
self.total_l1_loss, self.total_l1_smooth_loss,
self.pred_I2, self.I1_aug, self.I2_aug, self.I,
self.I_prime, self.pts1, self.gt, self.pred_h4p
])
h_total_loss_value += h_loss_value
rec_total_loss_value += rec_loss_value
ssim_total_loss_value += ssim_loss_value
l1_total_loss_value += l1_loss_value
total_num_fail_value += num_fail_value
h_losses_array.append(h_loss_value)
if args.save_visual:
I_sample = utils.denorm_img(I_value[0]).astype(np.uint8)
I_prime_sample = utils.denorm_img(I_prime_value[0]).astype(
np.uint8)
pts1_sample = pts1_value[0].reshape([4, 2]).astype(np.float32)
gt_h4p_sample = gt_value[0].reshape([4, 2]).astype(np.float32)
pts2_sample = pts1_sample + gt_h4p_sample
pred_h4p_sample = pred_h4p_value[0].reshape([4, 2]).astype(
np.float32)
pred_pts2_sample = pts1_sample + pred_h4p_sample
# Save
visual_file_name = str(
step *
args.batch_size) + '_' + args.loss_type + '_loss_' + str(
h_loss_value) + '.jpg'
utils.save_correspondences_img(I_prime_sample, I_sample,
pts1_sample, pts2_sample,
pred_pts2_sample,
args.results_dir,
visual_file_name)
if step % 10 == 0:
print('===> This iteration num Fail: %d \n' % num_fail_value)
total_time = utils.progress_bar(
step, self.num_total_steps,
'Test, h_loss %4.3f, rec_loss %4.3f, ssim_loss %4.3f, l1_loss %4.3f, fail_percent %4.4f'
% (h_total_loss_value / (step + 1), rec_total_loss_value /
(step + 1), ssim_total_loss_value /
(step + 1), l1_total_loss_value /
(step + 1), total_num_fail_value /
(step + 1) / args.batch_size))
if args.visual and step % 10 == 0:
plt.subplot(2, 1, 1)
plt.imshow(np.concatenate([
pred_I2_value[index, :, :, 0].astype(np.uint8),
I2_aug_value[index, :, :, 0]
], 1),
cmap='gray')
plt.title('Pred I2 vs I2')
plt.subplot(2, 1, 2)
plt.imshow(I1_aug_value[index, :, :, 0],
I2_aug_value[index, :, :, 0],
cmap='gray')
plt.title('I1_aug vs I2_aug')
plt.show()
plt.pause(0.05)
# Final result
total_time = utils.progress_bar(
step, self.num_total_steps,
'Test, h_loss %4.3f, rec_loss %4.3f, ssim_loss %4.3f, l1_loss %4.3f, fail_percent %4.4f'
%
(h_total_loss_value / (step + 1), rec_total_loss_value /
(step + 1), ssim_total_loss_value /
(step + 1), l1_total_loss_value /
(step + 1), total_num_fail_value / (step + 1) / args.batch_size))
# Summarize results
print('====> Result for RHO:', RHO, ' loss ', args.loss_type,
' noise ', args.do_augment)
print('|Steps | h_loss | l1_loss | Fail percent |')
print(step, h_total_loss_value / (step + 1),
l1_total_loss_value / (step + 1),
100 * total_num_fail_value / (step + 1) / args.batch_size)
tops_list = utils.find_percentile(h_losses_array)
print('===> Percentile Values: (20, 50, 80, 100):')
print(tops_list)
print('======> End! ====================================')
def test():
with tf.Graph().as_default(), tf.device('/gpu:0'): # Use GPU 0
# Training parameters
# Count the number of training & eval data
num_data = utils.count_text_lines(args.test_filenames_file)
print('===> Train: There are totally %d test files' % (num_data))
steps_per_epoch = np.ceil(num_data / args.batch_size).astype(np.int32)
num_total_steps = args.max_epoches * steps_per_epoch
# Load data
data_loader = Dataloader(test_dataloader_params,
shuffle=False) # no shuffle
# Debug test train_dataloader
# test_synthetic_dataloader(data_loader, True)
I1_batch = data_loader.I1_batch
I2_batch = data_loader.I2_batch
I1_aug_batch = data_loader.I1_aug_batch
I2_aug_batch = data_loader.I2_aug_batch
I_batch = data_loader.I_batch
I_prime_batch = data_loader.I_prime_batch
full_I_batch = data_loader.full_I_batch
full_I_prime_batch = data_loader.full_I_prime_batch
pts1_batch = data_loader.pts1_batch
gt_batch = data_loader.gt_batch
patch_indices_batch = data_loader.patch_indices_batch
# Train on multiple GPU:
h_losses = []
# Create a session
gpu_options = tf.GPUOptions(
allow_growth=True
) # Does not pre-allocate large, increase if needed
config = tf.ConfigProto(
allow_soft_placement=True, gpu_options=gpu_options
) # soft_placement allows to work on CPUs if GPUs are not available
sess = tf.Session(config=config)
# Initialize
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
# Threads coordinator
coordinator = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coordinator)
num_samples = 0
total_num_fail = 0
h_losses = []
total_time = 0
# Start test
for step in range(num_total_steps):
full_I_value, full_I_prime_value, I2_value, I1_value, I1_aug_value, I2_aug_value, I_value, I_prime_value, pts_1_value, full_gt_corr_value = sess.run(
[
full_I_batch, full_I_prime_batch, I2_batch, I1_batch,
I1_aug_batch, I2_aug_batch, I_batch, I_prime_batch,
pts1_batch, gt_batch
])
for i in range(args.batch_size):
num_samples += 1
down_ratio = 2
I_sample = utils.denorm_img(full_I_value[i]).astype(np.uint8)
I_prime_sample = utils.denorm_img(
full_I_prime_value[i]).astype(np.uint8)
corr1_sample = full_gt_corr_value[i, 0:8].reshape([
4, 2
]) / down_ratio # (gt is for 480x640. Here, use 240x320)
corr2_sample = full_gt_corr_value[i, 8:16].reshape(
[4, 2]) / down_ratio
# Use RANSAC_homography/ Direct method to find the homography (delta 4 points)
sample_start_time = timeit.default_timer()
if args.method == 'direct':
pred_h4p, _, not_found = direct_h.find_homography(
I_sample,
I_prime_sample,
corr1_sample,
corr2_sample,
visual=args.visual,
method=args.method,
num_iterations=args.num_iterations,
return_h_inv=False)
# RANSAC Methods
else:
pred_h4p, _, not_found = ransac_h.find_homography(
I_sample,
I_prime_sample,
corr1_sample,
corr2_sample,
visual=args.visual,
method=args.method,
min_match_count=args.num_features,
return_h_inv=False)
sample_run_time = timeit.default_timer() - sample_start_time
total_time += sample_run_time
# Set maximum value for every value of delta h4p
pred_h4p[np.where(pred_h4p >= 80)] = 80
pred_h4p[np.where(pred_h4p <= -80)] = -80
pred_corr2_sample = pred_h4p[0] + corr1_sample
h_loss_value = np.sqrt(
np.mean(np.square(pred_corr2_sample - corr2_sample)))
# Evaluate the result
# There are two cases of failure
if not_found: # Cannot find homography
total_num_fail += 1
print('===> Fail case 1: Not found homography')
else:
# H_loss if homography is identity matrix
h_loss_identity = np.sqrt(
np.mean(np.square(corr1_sample - corr2_sample)))
if h_loss_identity < h_loss_value:
print('===> Fail case 2: error > identity')
total_num_fail += 1
h_loss_value = h_loss_identity
h_losses.append(h_loss_value)
_ = utils.progress_bar(
step * args.batch_size + i,
num_total_steps * args.batch_size,
' Test| image %d, h_loss %.3f, h_loss_average %.3f, fail %d/%d, time %.4f'
% (i, h_loss_value, np.mean(h_losses), total_num_fail,
num_samples, sample_run_time))
# Save visualization
if args.save_visual:
# Query full images
img1_with_4pts = I_sample.astype(np.uint8)
img2_with_4pts = I_prime_sample.astype(np.uint8)
# Draw prediction
cv2.polylines(img2_with_4pts,
np.int32([pred_corr2_sample]), 1,
(5, 225, 225), 3)
point_color = (0, 255, 255)
line_color_set = [(255, 102, 255), (51, 153, 255),
(102, 255, 255), (255, 255, 0),
(102, 102, 244), (150, 202, 178),
(153, 240, 142), (102, 0, 51),
(51, 51, 0)]
# Draw 4 points (ground truth)
full_stack_images = utils.draw_matches(
img1_with_4pts,
corr1_sample,
img2_with_4pts,
corr2_sample,
'tmp.jpg',
color_set=line_color_set,
show=False)
# Save image
visual_file_name = os.path.join(
args.results_dir,
str(step * args.batch_size + i) + '_loss_' +
str(h_loss_value) + '.jpg')
#cv2.putText(full_stack_images, 'RMSE %.2f'%h_loss,(800, 100), cv2.FONT_HERSHEY_SIMPLEX, 1,(0,0,255),2)
cv2.imwrite(visual_file_name, full_stack_images)
print('Wrote file %s', visual_file_name)
result_dict = {
'method': args.method,
'h_losses': h_losses,
'h_loss_mu': np.mean(h_losses),
'h_loss_std': np.std(h_losses)
}
import cPickle as pickle
result_file_dir = os.path.join(args.results_dir, 'h_losses.pkl')
with open(result_file_dir, 'wb') as f:
pickle.dump(result_dict, f)
print('===> Successfully write results to %s' % result_file_dir)
print('==========================================================')
mean_h_loss, std_h_loss = np.mean(np.array(h_losses)), np.std(
np.array(h_losses))
print('===> H_loss:', mean_h_loss, '+/-', std_h_loss)
print('Running time:', total_time / num_samples)
fail_percent = total_num_fail * 1.0 / (num_samples)
print('Failure %.3f' % (fail_percent))
output_line = [
num_samples, mean_h_loss, std_h_loss, fail_percent,
total_time / num_samples
]
print('output_line:', output_line)
with open(os.path.join(args.log_dir, 'results.txt'), 'w') as f:
np.savetxt(f, [output_line], delimiter=' ', fmt='%.5f')
print('===> Wrote results to file %s' %
os.path.join(args.log_dir, 'results.txt'))
tops_list = utils.find_percentile(h_losses)
print('===> Percentile Values: (20, 50, 80, 100):')
print(tops_list)
print('======> End! ====================================')
def train():
# Overrides the current default graph for the lifetime of the context
with tf.Graph().as_default(), tf.device('/gpu:0'): # Use GPU 0
global_step = tf.Variable(0, trainable=False)
# Training parameters
# Count the number of training & eval data
num_data = utils.count_text_lines(args.filenames_file)
print('===> Train: There are totally %d training files' % (num_data))
num_total_steps = 150000
# Optimizer. Use exponential decay: decayed_lr = lr* decay_rate^ (global_steps/ decay_steps)
decay_rate = 0.96
decay_steps = (math.log(decay_rate) * num_total_steps) / math.log(
args.min_lr * 1.0 / args.lr)
print('args lr:', args.lr, args.min_lr)
print('===> Decay steps:', decay_steps)
learning_rate = tf.train.exponential_decay(args.lr,
global_step,
int(decay_steps),
decay_rate,
staircase=True)
# Due to slim.batch_norm docs:
# Note: when training, the moving_mean and moving_variance need to be updated.
# By default the update ops are placed in `tf.GraphKeys.UPDATE_OPS`, so they
# need to be added as a dependency to the `train_op`. For example:
# ```python
# update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
# with tf.control_dependencies(update_ops):
# train_op = optimizer.minimize(loss)
# ```
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
opt_step = tf.train.AdamOptimizer(learning_rate)
# Load data
data_loader = Dataloader(train_dataloader_params,
shuffle=True) # shuffle
I1_batch = data_loader.I1_batch
I2_batch = data_loader.I2_batch
I1_aug_batch = data_loader.I1_aug_batch
I2_aug_batch = data_loader.I2_aug_batch
I_batch = data_loader.I_batch
I_prime_batch = data_loader.I_prime_batch
pts1_batch = data_loader.pts1_batch
gt_batch = data_loader.gt_batch
patch_indices_batch = data_loader.patch_indices_batch
# Split on multiple GPU
I1_splits = tf.split(I1_batch, args.num_gpus, 0)
I2_splits = tf.split(I2_batch, args.num_gpus, 0)
I1_aug_splits = tf.split(I1_aug_batch, args.num_gpus, 0)
I2_aug_splits = tf.split(I2_aug_batch, args.num_gpus, 0)
I_splits = tf.split(I_batch, args.num_gpus, 0)
I_prime_splits = tf.split(I_prime_batch, args.num_gpus, 0)
pts1_splits = tf.split(pts1_batch, args.num_gpus, 0)
gt_splits = tf.split(gt_batch, args.num_gpus, 0)
patch_indices_splits = tf.split(patch_indices_batch, args.num_gpus, 0)
# Train on multiple GPU:
multi_grads = []
reuse_variables = None
h_losses = []
rec_losses = []
ssim_losses = []
l1_losses = []
l1_smooth_losses = []
ncc_losses = []
model_params = homography_model_params(
mode=args.mode,
batch_size=int(args.batch_size / args.num_gpus),
patch_size=args.patch_size,
img_h=args.img_h,
img_w=args.img_w,
loss_type=args.loss_type,
use_batch_norm=args.use_batch_norm,
augment_list=args.augment_list,
leftright_consistent_weight=args.leftright_consistent_weight)
# Deal with sharable variables
with tf.variable_scope(tf.get_variable_scope()):
for i in range(args.num_gpus):
with tf.device('/gpu:%d' % i):
model = HomographyModel(model_params,
I1_splits[i],
I2_splits[i],
I1_aug_splits[i],
I2_aug_splits[i],
I_splits[i],
I_prime_splits[i],
pts1_splits[i],
gt_splits[i],
patch_indices_splits[i],
reuse_variables=reuse_variables,
model_index=i)
h_loss = model.h_loss
rec_loss = model.rec_loss
ssim_loss = model.ssim_loss
l1_loss = model.l1_loss
l1_smooth_loss = model.l1_smooth_loss
ncc_loss = model.ncc_loss
pred_I2 = model.pred_I2
I2 = model.I2
H_mat = model.H_mat
I1 = model.I1
I = model.I
I1_aug = model.I1_aug
I2_aug = model.I2_aug
h_losses.append(h_loss)
rec_losses.append(rec_loss)
ssim_losses.append(ssim_loss)
l1_losses.append(l1_loss)
l1_smooth_losses.append(l1_smooth_loss)
ncc_losses.append(ncc_loss)
reuse_variables = True
if args.loss_type == 'h_loss':
grads = opt_step.compute_gradients(h_loss)
elif args.loss_type == 'rec_loss':
grads = opt_step.compute_gradients(rec_loss)
elif args.loss_type == 'ssim_loss':
grads = opt_step.compute_gradients(ssim_loss)
elif args.loss_type == 'l1_loss':
grads = opt_step.compute_gradients(l1_loss)
elif args.loss_type == 'l1_smooth_loss':
grads = opt_step.compute_gradients(l1_smooth_loss)
elif args.loss_type == 'ncc_loss':
grads = opt_step.compute_gradients(ncc_loss)
else:
print('===> Loss type does not exist!')
exit(0)
print('====> Use loss type: ', args.loss_type)
time.sleep(2)
multi_grads.append(grads)
# Take average of the grads
grads = utils.get_average_grads(multi_grads)
apply_grad_opt = opt_step.apply_gradients(grads,
global_step=global_step)
total_h_loss = tf.reduce_mean(h_losses)
total_rec_loss = tf.reduce_mean(rec_losses)
total_ssim_loss = tf.reduce_mean(ssim_losses)
total_l1_loss = tf.reduce_mean(l1_losses)
total_l1_smooth_loss = tf.reduce_mean(l1_smooth_losses)
total_ncc_loss = tf.reduce_mean(ncc_losses)
with tf.name_scope('Losses'):
tf.summary.scalar('Learning_rate', learning_rate)
tf.summary.scalar('Total_h_loss', total_h_loss)
tf.summary.scalar('Total_rec_loss', total_rec_loss)
tf.summary.scalar('Total_ssim_loss', total_ssim_loss)
tf.summary.scalar('Total_l1_loss', total_l1_loss)
tf.summary.scalar('Total_l1_smooth_loss', total_l1_smooth_loss)
tf.summary.scalar('Total_ncc_loss', total_ncc_loss)
summary_opt = tf.summary.merge_all()
# Create a session
gpu_options = tf.GPUOptions(
allow_growth=True
) # Does not pre-allocate large, increase if needed
config = tf.ConfigProto(
allow_soft_placement=True, gpu_options=gpu_options
) # soft_placement allows to work on CPUs if GPUs are not available
sess = tf.Session(config=config)
# Saver
log_name = args.loss_type
summary_writer = tf.summary.FileWriter(args.log_dir, sess.graph)
train_saver = tf.train.Saver(max_to_keep=5) # Keep maximum 5 models
# Initialize
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
# Threads coordinator
coordinator = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coordinator)
# Restore
if args.resume:
train_saver.restore(sess,
tf.train.latest_checkpoint(args.model_dir))
if args.retrain:
sess.run(global_step.assign(0))
# Index of the image want to display
index = 0
h_total_loss_value = 0
rec_total_loss_value = 0
ssim_total_loss_value = 0
l1_total_loss_value = 0
l1_smooth_total_loss_value = 0
ncc_total_loss_value = 0
start_step = global_step.eval(session=sess)
print('===> Start step:', start_step)
# Start training
for step in range(start_step, start_step + num_total_steps):
if args.visual:
_, h_loss_value, rec_loss_value, ssim_loss_value, l1_loss_value, l1_smooth_loss_value, ncc_loss_value, lr_value, H_mat_value, pred_I2_value, I2_value, I1_value, I1_aug_value, I2_aug_value, I_value = sess.run(
[
apply_grad_opt, total_h_loss, total_rec_loss,
total_ssim_loss, total_l1_loss, total_l1_smooth_loss,
total_ncc_loss, learning_rate, H_mat, pred_I2, I2, I1,
I1_aug, I2_aug, I
])
elif args.loss_type == "l1_loss" and not args.visual:
_, h_loss_value, l1_loss_value, l1_smooth_loss_value, lr_value = sess.run(
[
apply_grad_opt, total_h_loss, total_l1_loss,
total_l1_smooth_loss, learning_rate
])
h_total_loss_value += h_loss_value
l1_total_loss_value += l1_loss_value
l1_smooth_total_loss_value += l1_smooth_loss_value
if step % 100 == 0:
total_time = utils.progress_bar(
step, num_total_steps + start_step - 1,
'Train: 1, h_loss %4.3f, l1_loss %.6f, l1_smooth_loss %.6f, lr %.6f'
%
(h_total_loss_value /
(step - start_step + 1), l1_total_loss_value /
(step - start_step + 1), l1_smooth_total_loss_value /
(step - start_step + 1), lr_value))
else:
_, h_loss_value, rec_loss_value, ssim_loss_value, l1_loss_value, l1_smooth_loss_value, ncc_loss_value, lr_value = sess.run(
[
apply_grad_opt, total_h_loss, total_rec_loss,
total_ssim_loss, total_l1_loss, total_l1_smooth_loss,
total_ncc_loss, learning_rate
])
h_total_loss_value += h_loss_value
rec_total_loss_value += rec_loss_value
ssim_total_loss_value += ssim_loss_value
l1_total_loss_value += l1_loss_value
l1_smooth_total_loss_value += l1_smooth_loss_value
ncc_total_loss_value += ncc_loss_value
if step % 100 == 0:
total_time = utils.progress_bar(
step, num_total_steps + start_step - 1,
'Train: 1, h_loss %4.3f, rec_loss %4.3f, ssim_loss %.6f. l1_loss %.6f, l1_smooth_loss %.6f, ncc_loss %.6f, lr %.6f'
%
(h_total_loss_value /
(step - start_step + 1), rec_total_loss_value /
(step - start_step + 1), ssim_total_loss_value /
(step - start_step + 1), l1_total_loss_value /
(step - start_step + 1), l1_smooth_total_loss_value /
(step - start_step + 1), ncc_total_loss_value /
(step - start_step + 1), lr_value))
# Tensorboard
if step % 1000 == 0:
summary_str = sess.run(summary_opt)
summary_writer.add_summary(summary_str, global_step=step)
if step and step % 1000 == 0:
train_saver.save(sess,
args.model_dir + args.model_name,
global_step=step)
if args.visual and step % 1 == 0:
if 'normalize' in args.augment_list:
pred_I2_sample_value = utils.denorm_img(
pred_I2_value[index, :, :, 0]).astype(np.uint8)
I2_sample_value = utils.denorm_img(
I2_value[index, :, :, 0]).astype(np.uint8)
I1_sample_value = utils.denorm_img(
I1_value[index, :, :, 0]).astype(np.uint8)
I1_aug_sample_value = utils.denorm_img(
I1_aug_value[index, :, :, 0]).astype(np.uint8)
I2_aug_sample_value = utils.denorm_img(
I2_aug_value[index, :, :, 0]).astype(np.uint8)
I_sample_value = utils.denorm_img(
I_value[index, ...]).astype(np.uint8)
else:
pred_I2_sample_value = pred_I2_value[index, :, :,
0].astype(np.uint8)
I2_sample_value = I2_value[index, :, :, 0].astype(np.uint8)
I1_sample_value = I1_value[index, :, :, 0].astype(np.uint8)
I1_aug_sample_value = I1_aug_value[index, :, :,
0].astype(np.uint8)
I2_aug_sample_value = I2_aug_value[index, :, :,
0].astype(np.uint8)
I_sample_value = I_value[index, ...].astype(np.uint8)
plt.subplot(3, 1, 1)
plt.imshow(np.concatenate(
[pred_I2_sample_value, I2_sample_value], 1),
cmap='gray')
plt.title('Pred I2 vs I2')
plt.subplot(3, 1, 2)
plt.imshow(np.concatenate(
[I1_aug_sample_value, I2_aug_sample_value], 1),
cmap='gray')
plt.title('I1_aug vs I2_aug')
plt.subplot(3, 1, 3)
plt.imshow(I_sample_value if I_sample_value.shape[2] ==
3 else I_sample_value[:, :, 0])
plt.title('I')
plt.show()
plt.pause(0.05)
# Save the final model
train_saver.save(sess,
args.model_dir + args.model_name,
global_step=step)
def run_test(self, model_index=0):
# Create a session
gpu_options = tf.GPUOptions(
allow_growth=True
) # Does not pre-allocate large, increase if needed
config = tf.ConfigProto(
allow_soft_placement=True, gpu_options=gpu_options
) # soft_placement allows to work on CPUs if GPUs are not available
sess = tf.Session(config=config)
# Saver
log_name = args.loss_type
summary_writer = tf.summary.FileWriter(args.log_dir, sess.graph)
train_saver = tf.train.Saver(max_to_keep=20) # Keep maximum 20 models
# Initialize
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
# Threads coordinator
coordinator = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coordinator)
# Restore
print(args.save_model_dir)
train_saver.restore(sess,
tf.train.latest_checkpoint(args.save_model_dir))
# Index of the image want to display
index = 0
rec_total_loss_value = 0
ssim_total_loss_value = 0
l1_total_loss_value = 0
# Start testing
total_num_fail_value = 0
h_losses = []
num_samples = 0
for step in range(self.num_total_steps):
full_gt_corr_value, pts1_value, pred_h4p_value, rec_loss_value, l1_loss_value, l1_smooth_loss_value, pred_I2_value, I1_aug_value, I2_aug_value, full_I_value, full_I_prime_value = sess.run(
[
self.gt_corr, self.pts1, self.pred_h4p,
self.total_rec_loss, self.total_l1_loss,
self.total_l1_smooth_loss, self.pred_I2, self.I1_aug,
self.I2_aug, self.full_I, self.full_I_prime
])
rec_total_loss_value += rec_loss_value
l1_total_loss_value += l1_loss_value
# Now, use the correspondences that we manually chose (ground truth), compute homography error
# pred_h4p_value: Batch_size x 8
# We cannot just compare these value with the ground truth directly since four points on the first images are different.
for i in range(args.batch_size):
full_I_sample = utils.denorm_img(full_I_value[i]).astype(
np.uint8)
full_I_prime_sample = utils.denorm_img(
full_I_prime_value[i]).astype(np.uint8)
# Ratio between full image and our image used for deep learing
down_ratio = args.full_img_h / args.img_h
delta_h4p_sample = pred_h4p_value[i].reshape([4, 2]).astype(
np.float32)
pts1_sample = pts1_value[i].reshape([4, 2]).astype(np.float32)
# Project onto the full image
full_pts2_sample = (delta_h4p_sample +
pts1_sample) * down_ratio
full_pts1_sample = pts1_sample * down_ratio
# Find homography
full_H = cv2.getPerspectiveTransform(full_pts1_sample,
full_pts2_sample)
full_H_inv = np.linalg.inv(full_H)
# Obtain the chosen correspondences on which we compute the homography error
full_corr1_sample = full_gt_corr_value[i, 0:8].reshape([
4, 2
]) / 2 # Since we chose point on images of size 480 x 640
full_corr2_sample = full_gt_corr_value[i, 8:16].reshape([4, 2
]) / 2
# Apply estimated homography to the corr1
pred_full_corr2_sample = cv2.perspectiveTransform(
np.array([full_corr1_sample]), full_H_inv)[0]
# Find RMSE
h_loss = np.sqrt(
np.mean(
np.square(pred_full_corr2_sample - full_corr2_sample))
) #normalize to be equivalent to SIFT, ORB, direct
h_loss_identity = np.sqrt(
np.mean(np.square(full_corr1_sample - full_corr2_sample)))
if h_loss > h_loss_identity:
print("===> Found %.3f > %.3f = RMSE_identity" %
(h_loss, h_loss_identity))
total_num_fail_value += 1
#TODO: conditionally
h_loss = h_loss_identity
num_samples += 1
h_loss = h_loss
h_losses.append(h_loss)
# VISUALIZATION
# Save visualization
if args.save_visual:
# Query full images
visual_file_name = str(
step * args.batch_size +
i) + '_Unsupervised_' + '_loss_' + str(h_loss) + '.jpg'
utils.save_correspondences_img(
full_I_sample, full_I_prime_sample, full_corr1_sample,
full_corr2_sample, pred_full_corr2_sample,
args.results_dir, visual_file_name)
# Save visualization for other methods - for report
if args.do_report:
for fb_method in ['SIFT']:
SIFT_pred_h4p, _, not_found = ransac_h.find_homography(
full_I_sample,
full_I_prime_sample,
full_corr1_sample,
visual=False,
method='SIFT',
min_match_count=25,
return_h_inv=False)
SIFT_pred_full_corr2_sample = SIFT_pred_h4p[
0] + full_corr1_sample
SIFT_h_loss = np.sqrt(
np.mean(
np.square(SIFT_pred_full_corr2_sample -
full_corr2_sample)))
visual_file_name = str(
step * args.batch_size +
i) + '_' + fb_method + '_loss_' + str(
SIFT_h_loss) + '.jpg'
utils.save_correspondences_img(
full_I_sample, full_I_prime_sample,
full_corr1_sample, full_corr2_sample,
SIFT_pred_full_corr2_sample, args.results_dir,
visual_file_name)
# Direct method
direct_pred_h4p, _, not_found = direct_h.find_homography(
full_I_sample,
full_I_prime_sample,
full_corr1_sample,
visual=False,
method='direct',
num_iterations=1000,
return_h_inv=False)
direct_pred_full_corr2_sample = direct_pred_h4p[
0] + full_corr1_sample
direct_h_loss = np.sqrt(
np.mean(
np.square(direct_pred_full_corr2_sample -
full_corr2_sample)))
visual_file_name = str(step * args.batch_size +
i) + '_direct' + '_loss_' + str(
direct_h_loss) + '.jpg'
utils.save_correspondences_img(
full_I_sample, full_I_prime_sample, full_corr1_sample,
full_corr2_sample, direct_pred_full_corr2_sample,
args.results_dir, visual_file_name)
if step % 1 == 0:
total_time = utils.progress_bar(
step, self.num_total_steps,
'Test, h_loss %4.3f, rec_loss %4.3f, l1_loss %4.3f, fail_percent %4.4f'
% (np.mean(h_losses), rec_total_loss_value /
(step + 1), l1_total_loss_value /
(step + 1), total_num_fail_value /
(step + 1) / args.batch_size))
if args.visual and step % 10 == 0:
plt.subplot(2, 1, 1)
plt.imshow(np.concatenate([
utils.denorm_img(pred_I2_value[index, :, :, 0]).astype(
np.uint8),
utils.denorm_img(I2_aug_value[index, :, :, 0])
], 1),
cmap='gray')
plt.title('Pred I2 vs I2')
plt.subplot(2, 1, 2)
plt.imshow(np.concatenate([
I1_aug_value[index, :, :, 0], I2_aug_value[index, :, :, 0]
], 1),
cmap='gray')
plt.title('I1_aug vs I2_aug')
plt.show()
plt.pause(0.05)
# Final result
total_time = utils.progress_bar(
step, self.num_total_steps,
'Test, h_loss %4.3f, rec_loss %4.3f, ssim_loss %4.3f, l1_loss %4.3f, fail_percent %4.4f'
%
(np.mean(h_losses), rec_total_loss_value /
(step + 1), ssim_total_loss_value /
(step + 1), l1_total_loss_value /
(step + 1), total_num_fail_value / (step + 1) / args.batch_size))
print('====> Result for RHO:', RHO, ' loss ', args.loss_type,
' noise ', args.do_augment)
print('|Steps | h_loss | l1_loss | Fail percent |')
print(step, np.mean(h_losses), l1_total_loss_value / (step + 1),
100 * total_num_fail_value / (step + 1) / args.batch_size)
tops_list = utils.find_percentile(h_losses)
print('===> Percentile Values: (20, 50, 80, 100):')
print(tops_list)
print('======> End! ====================================')
def train(args):
# Load data
TrainDataset = SyntheticDataset(data_path=args.data_path,
mode=args.mode,
img_h=args.img_h,
img_w=args.img_w,
patch_size=args.patch_size,
do_augment=args.do_augment)
train_loader = DataLoader(TrainDataset, batch_size=args.batch_size, shuffle=True, num_workers=4)
print('===> Train: There are totally {} training files'.format(len(TrainDataset)))
net = HomographyModel(args.use_batch_norm)
if args.resume:
model_path = os.path.join(args.model_dir, args.model_name)
ckpt = torch.load(model_path)
net.load_state_dict(ckpt.state_dict())
if torch.cuda.is_available():
net = net.cuda()
optimizer = optim.Adam(net.parameters(), lr=args.lr) # default as 0.0001
decay_rate = 0.96
step_size = (math.log(decay_rate) * args.max_epochs) / math.log(args.min_lr * 1.0 / args.lr)
print('args lr:', args.lr, args.min_lr)
print('===> Decay steps:', step_size)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=int(step_size), gamma=0.96)
print("start training")
writer = SummaryWriter(logdir=args.log_dir, flush_secs=60)
score_print_fre = 100
summary_fre = 1000
model_save_fre = 4000
glob_iter = 0
t0 = time.time()
for epoch in range(args.max_epochs):
net.train()
epoch_start = time.time()
train_l1_loss = 0.0
train_l1_smooth_loss = 0.0
train_h_loss = 0.0
for i, batch_value in enumerate(train_loader):
I1_batch = batch_value[0].float()
I2_batch = batch_value[1].float()
I1_aug_batch = batch_value[2].float()
I2_aug_batch = batch_value[3].float()
I_batch = batch_value[4].float()
I_prime_batch = batch_value[5].float()
pts1_batch = batch_value[6].float()
gt_batch = batch_value[7].float()
patch_indices_batch = batch_value[8].float()
if torch.cuda.is_available():
I1_aug_batch = I1_aug_batch.cuda()
I2_aug_batch = I2_aug_batch.cuda()
I_batch = I_batch.cuda()
pts1_batch = pts1_batch.cuda()
gt_batch = gt_batch.cuda()
patch_indices_batch = patch_indices_batch.cuda()
# forward, backward, update weights
optimizer.zero_grad()
batch_out = net(I1_aug_batch, I2_aug_batch, I_batch, pts1_batch, gt_batch, patch_indices_batch)
h_loss = batch_out['h_loss']
rec_loss = batch_out['rec_loss']
ssim_loss = batch_out['ssim_loss']
l1_loss = batch_out['l1_loss']
l1_smooth_loss = batch_out['l1_smooth_loss']
ncc_loss = batch_out['ncc_loss']
pred_I2 = batch_out['pred_I2']
loss = l1_loss
loss.backward()
optimizer.step()
train_l1_loss += loss.item()
train_l1_smooth_loss += l1_smooth_loss.item()
train_h_loss += h_loss.item()
if (i + 1) % score_print_fre == 0 or (i + 1) == len(train_loader):
print(
"Training: Epoch[{:0>3}/{:0>3}] Iter[{:0>3}]/[{:0>3}] l1 loss: {:.4f} "
"l1 smooth loss: {:.4f} h loss: {:.4f} lr={:.8f}".format(
epoch + 1, args.max_epochs, i + 1, len(train_loader), train_l1_loss / score_print_fre,
train_l1_smooth_loss / score_print_fre, train_h_loss / score_print_fre, scheduler.get_lr()[0]))
train_l1_loss = 0.0
train_l1_smooth_loss = 0.0
train_h_loss = 0.0
if glob_iter % summary_fre == 0:
writer.add_scalar('learning_rate', scheduler.get_lr()[0], glob_iter)
writer.add_scalar('h_loss', h_loss, glob_iter)
writer.add_scalar('rec_loss', rec_loss, glob_iter)
writer.add_scalar('ssim_loss', ssim_loss, glob_iter)
writer.add_scalar('l1_loss', l1_loss, glob_iter)
writer.add_scalar('l1_smooth_loss', l1_smooth_loss, glob_iter)
writer.add_scalar('ncc_loss', ncc_loss, glob_iter)
writer.add_image('I', utils.denorm_img(I_batch[0, ...].cpu().numpy()).astype(np.uint8)[:, :, ::-1],
glob_iter, dataformats='HWC')
writer.add_image('I_prime',
utils.denorm_img(I_prime_batch[0, ...].numpy()).astype(np.uint8)[:, :, ::-1],
glob_iter, dataformats='HWC')
writer.add_image('I1_aug', utils.denorm_img(I1_aug_batch[0, 0, ...].cpu().numpy()).astype(np.uint8),
glob_iter, dataformats='HW')
writer.add_image('I2_aug', utils.denorm_img(I2_aug_batch[0, 0, ...].cpu().numpy()).astype(np.uint8),
glob_iter, dataformats='HW')
writer.add_image('pred_I2',
utils.denorm_img(pred_I2[0, 0, ...].cpu().detach().numpy()).astype(np.uint8),
glob_iter, dataformats='HW')
writer.add_image('I2', utils.denorm_img(I2_batch[0, 0, ...].numpy()).astype(np.uint8), glob_iter,
dataformats='HW')
writer.add_image('I1', utils.denorm_img(I1_batch[0, 0, ...].numpy()).astype(np.uint8), glob_iter,
dataformats='HW')
# save model
if glob_iter % model_save_fre == 0 and glob_iter != 0:
filename = 'model' + '_iter_' + str(glob_iter) + '.pth'
model_save_path = os.path.join(args.model_dir, filename)
torch.save(net, model_save_path)
glob_iter += 1
scheduler.step()
print("Epoch: {} epoch time: {:.1f}s".format(epoch, time.time() - epoch_start))
elapsed_time = time.time() - t0
print("Finished Training in {:.0f}h {:.0f}m {:.0f}s.".format(
elapsed_time // 3600, (elapsed_time % 3600) // 60, (elapsed_time % 3600) % 60))
def test(args):
# Load data
TestDataset = SyntheticDataset(data_path=args.data_path,
mode=args.mode,
img_h=args.img_h,
img_w=args.img_w,
patch_size=args.patch_size,
do_augment=args.do_augment)
test_loader = DataLoader(TestDataset, batch_size=1)
print('===> Test: There are totally {} testing files'.format(len(TestDataset)))
# Load model
net = HomographyModel()
model_path = os.path.join(args.model_dir, args.model_name)
state = torch.load(model_path)
net.load_state_dict(state.state_dict())
if torch.cuda.is_available():
net = net.cuda()
print("start testing")
with torch.no_grad():
net.eval()
test_l1_loss = 0.0
test_h_loss = 0.0
h_losses_array = []
for i, batch_value in enumerate(test_loader):
I1_aug_batch = batch_value[2].float()
I2_aug_batch = batch_value[3].float()
I_batch = batch_value[4].float()
I_prime_batch = batch_value[5].float()
pts1_batch = batch_value[6].float()
gt_batch = batch_value[7].float()
patch_indices_batch = batch_value[8].float()
if torch.cuda.is_available():
I1_aug_batch = I1_aug_batch.cuda()
I2_aug_batch = I2_aug_batch.cuda()
I_batch = I_batch.cuda()
pts1_batch = pts1_batch.cuda()
gt_batch = gt_batch.cuda()
patch_indices_batch = patch_indices_batch.cuda()
batch_out = net(I1_aug_batch, I2_aug_batch, I_batch, pts1_batch, gt_batch, patch_indices_batch)
h_loss = batch_out['h_loss']
rec_loss = batch_out['rec_loss']
ssim_loss = batch_out['ssim_loss']
l1_loss = batch_out['l1_loss']
pred_h4p_value = batch_out['pred_h4p']
test_h_loss += h_loss.item()
test_l1_loss += l1_loss.item()
h_losses_array.append(h_loss.item())
if args.save_visual:
I_sample = utils.denorm_img(I_batch[0].cpu().numpy()).astype(np.uint8)
I_prime_sample = utils.denorm_img(I_prime_batch[0].numpy()).astype(np.uint8)
pts1_sample = pts1_batch[0].cpu().numpy().reshape([4, 2]).astype(np.float32)
gt_h4p_sample = gt_batch[0].cpu().numpy().reshape([4, 2]).astype(np.float32)
pts2_sample = pts1_sample + gt_h4p_sample
pred_h4p_sample = pred_h4p_value[0].cpu().numpy().reshape([4, 2]).astype(np.float32)
pred_pts2_sample = pts1_sample + pred_h4p_sample
# Save
visual_file_name = ('%s' % i).zfill(4) + '.jpg'
utils.save_correspondences_img(I_prime_sample, I_sample, pts1_sample, pts2_sample, pred_pts2_sample,
args.results_dir, visual_file_name)
print("Testing: h_loss: {:4.3f}, rec_loss: {:4.3f}, ssim_loss: {:4.3f}, l1_loss: {:4.3f}".format(
h_loss.item(), rec_loss.item(), ssim_loss.item(), l1_loss.item()
))
print('|Test size | h_loss | l1_loss |')
print(len(test_loader), test_h_loss / len(test_loader), test_l1_loss / len(test_loader))
tops_list = utils.find_percentile(h_losses_array)
print('===> Percentile Values: (20, 50, 80, 100):')
print(tops_list)
print('======> End! ====================================')
def test():
with tf.Graph().as_default(), tf.device('/gpu:0'): # Use GPU 0
# Training parameters
# Count the number of training & eval data
num_data = utils.count_text_lines(args.test_filenames_file)
print('===> Train: There are totally %d test files'%(num_data))
steps_per_epoch = np.ceil(num_data/args.batch_size).astype(np.int32)
num_total_steps = args.max_epoches*steps_per_epoch
# Load data
data_loader = Dataloader(test_dataloader_params, shuffle=False) # no shuffle
# Debug test train_dataloader
# test_synthetic_dataloader(data_loader, True)
I1_batch = data_loader.I1_batch
I2_batch = data_loader.I2_batch
I1_aug_batch = data_loader.I1_aug_batch
I2_aug_batch = data_loader.I2_aug_batch
I_batch = data_loader.I_batch
I_prime_batch = data_loader.I_prime_batch
pts1_batch = data_loader.pts1_batch
gt_batch = data_loader.gt_batch
patch_indices_batch = data_loader.patch_indices_batch
# Train on multiple GPU:
h_losses = []
# Create a session
gpu_options = tf.GPUOptions(allow_growth=True) # Does not pre-allocate large, increase if needed
config = tf.ConfigProto(allow_soft_placement=True, gpu_options=gpu_options) # soft_placement allows to work on CPUs if GPUs are not available
sess = tf.Session(config=config)
# Initialize
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
# Threads coordinator
coordinator = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coordinator)
num_samples = 0
total_num_fail = 0
h_total_loss_value = 0
h_losses = []
total_time = 0
# Start test
for step in range(num_total_steps):
I2_value, I1_value, I1_aug_value, I2_aug_value, I_value, I_prime_value, pts_1_value, gt_value = sess.run([I2_batch, I1_batch, I1_aug_batch, I2_aug_batch, I_batch, I_prime_batch, pts1_batch, gt_batch])
for i in range(args.batch_size):
num_samples += 1
I_sample = utils.denorm_img(I_value[i]).astype(np.uint8)
I_prime_sample = utils.denorm_img(I_prime_value[i]).astype(np.uint8)
pts_1_sample = pts_1_value[i].reshape([4,2])
gt_sample = gt_value[i].reshape([4,2])
pts_2_sample = pts_1_sample + gt_sample
# Use RANSAC_homography to find the homography (delta 4 points)
sample_start_time = timeit.default_timer()
if args.method == 'direct':
pred_h4p, _, not_found = direct_h.find_homography(I_sample, I_prime_sample, pts_1_sample, pts_2_sample, visual=args.visual, method=args.method, num_iterations=args.num_iterations)
else:
pred_h4p, _, not_found = ransac_h.find_homography(I_sample, I_prime_sample, pts_1_sample, pts_2_sample, visual=args.visual, method=args.method, min_match_count=args.num_features)
sample_run_time = timeit.default_timer() - sample_start_time
total_time += sample_run_time
# Set maximum value for every value of delta h4p
pred_h4p[np.where(pred_h4p >= RHO)] = RHO*2
pred_h4p[np.where(pred_h4p <=- RHO)] = -RHO*2
h_loss_value = np.sqrt(np.mean(np.square(pred_h4p[0] - gt_sample)))
# Evaluate the result
# There are two cases of failure
if not_found: # Cannot find homography
total_num_fail += 1
print('===> Fail case 1: Not found homography')
else:
# H_loss if homography is identity matrix
h_loss_identity = np.sqrt(np.mean(np.square(gt_sample)))
if h_loss_identity < h_loss_value:
print('===> Fail case 2: error > identity')
total_num_fail += 1
h_loss_value = h_loss_identity
h_losses.append(h_loss_value)
h_loss_value = np.sqrt(np.mean(np.square(pred_h4p - gt_sample)))
_ = utils.progress_bar(step*args.batch_size+i, num_total_steps*args.batch_size, ' Test| image %d, h_loss %.3f, h_loss_average %.3f, fail %d/%d, time %.4f'%(i, h_loss_value, np.mean(h_losses), total_num_fail, num_samples, sample_run_time))
print ('==========================================================')
mean_h_loss, std_h_loss = np.mean(np.array(h_losses)), np.std(np.array(h_losses))
print ('===> H_loss:', mean_h_loss, '+/-', std_h_loss)
print ('Running time:', total_time/num_samples)
fail_percent = total_num_fail*1.0/(num_samples)
print ('Failure %.3f'%(fail_percent))
output_line = [num_samples, mean_h_loss, std_h_loss, fail_percent, total_time/num_samples]
print ('output_line:', output_line)
with open(os.path.join(args.log_dir, 'results.txt'), 'w') as f:
np.savetxt(f, [output_line], delimiter= ' ', fmt='%.5f')
print('===> Wrote results to file %s'%os.path.join(args.log_dir, 'results.txt'))
tops_list = utils.find_percentile(h_losses)
print('===> Percentile Values: (30, 60, 100):')
print(tops_list)
