Пример #1
0
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! ====================================')
Пример #2
0
if args.visual:
    plt.ion()

train_dataloader_params = dataloader_params(data_path=args.data_path,
                                            filenames_file=args.filenames_file,
                                            pts1_file=args.pts1_file,
                                            gt_file=args.gt_file,
                                            mode=args.mode,
                                            batch_size=args.batch_size,
                                            img_h=args.img_h,
                                            img_w=args.img_w,
                                            patch_size=args.patch_size,
                                            augment_list=args.augment_list,
                                            do_augment=args.do_augment)

num_test_data = utils.count_text_lines(args.test_filenames_file)
# num_train_data = utils.count_text_lines(args.filenames_file)
# print('===> There are totally %d train files'%(num_train_data))
print('===> There are totally %d test files' % (num_test_data))
test_batch_size = np.min([num_test_data, args.batch_size])

test_dataloader_params = extended_dataloader_params(
    data_path=args.data_path,
    filenames_file=args.test_filenames_file,
    pts1_file=args.test_pts1_file,
    gt_file=args.test_gt_file,
    mode='test',
    batch_size=test_batch_size,
    img_h=args.img_h,
    img_w=args.img_w,
    patch_size=args.patch_size,
Пример #3
0
    def __init__(self):
        # Overrides the current default graph for the lifetime of the context
        with tf.device('/gpu:0'):  # Use GPU 0

            # Count the number of eval data
            num_data = utils.count_text_lines(args.test_filenames_file)
            print('===> Test: There are totally %d Test files' % (num_data))

            steps_per_epoch = np.ceil(num_data / args.batch_size).astype(
                np.int32)
            self.num_total_steps = 3 * steps_per_epoch  # Test 3 epoches

            # Load data
            data_loader = Dataloader(test_dataloader_params,
                                     shuffle=True)  # No 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:
            reuse_variables = None
            h_losses = []
            rec_losses = []
            ssim_losses = []
            l1_losses = []
            l1_smooth_losses = []
            num_fails = []
            model_params = homography_model_params(
                mode='test',
                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)
                        # Debug test splits
                        #test_synthetic_dataloader(data_loader, True, I1_splits[i], I2_splits[i], I_splits[i], I_prime_splits[i], pts1_splits[i], gt_splits[i], patch_indices_splits[i])

                        reuse_variables = True
                        # In testing, use bounded_h_loss (under successful condition)
                        h_loss = model.bounded_h_loss
                        rec_loss = model.rec_loss
                        ssim_loss = model.ssim_loss
                        l1_loss = model.l1_loss
                        l1_smooth_loss = model.l1_smooth_loss
                        num_fail = model.num_fail

                        self.pred_I2 = model.pred_I2
                        self.I2 = model.I2
                        self.H_mat = model.H_mat
                        self.I1 = model.I1
                        self.I1_aug = model.I1_aug
                        self.I2_aug = model.I2_aug
                        self.I = model.I
                        self.I_prime = model.I_prime
                        self.pts1 = model.pts_1
                        self.gt = model.gt
                        self.pred_h4p = model.pred_h4p

                        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)
                        num_fails.append(num_fail)
            self.total_h_loss = tf.reduce_mean(h_losses)
            self.total_num_fail = tf.reduce_sum(num_fails)
            self.total_rec_loss = tf.reduce_mean(rec_losses)
            self.total_ssim_loss = tf.reduce_mean(ssim_losses)
            self.total_l1_loss = tf.reduce_mean(l1_losses)
            self.total_l1_smooth_loss = tf.reduce_mean(l1_smooth_losses)
            with tf.name_scope('Losses'):
                tf.summary.scalar('Total_h_loss', self.total_h_loss)
                tf.summary.scalar('Total_rec_loss', self.total_rec_loss)
                tf.summary.scalar('Total_ssim_loss', self.total_ssim_loss)
                tf.summary.scalar('Total_l1_loss', self.total_l1_loss)
                tf.summary.scalar('Total_l1_smooth_loss',
                                  self.total_l1_smooth_loss)
            self.summary_opt = tf.summary.merge_all()
Пример #4
0
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 __init__(self):
        # Overrides the current default graph for the lifetime of the context
        with tf.device('/gpu:0'):  # Use GPU 0

            # Count the number of eval data
            num_data = utils.count_text_lines(args.test_filenames_file)
            print('===> Test: There are totally %d Test files' % (num_data))

            steps_per_epoch = np.ceil(num_data / args.batch_size).astype(
                np.int32)
            self.num_total_steps = 2 * steps_per_epoch  # Test 2 epoches

            # 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

            # 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:
            reuse_variables = None

            rec_losses = []
            ssim_losses = []
            l1_losses = []
            l1_smooth_losses = []
            num_fails = []
            model_params = homography_model_params(
                mode='test',
                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):
                        # Note that ground truth gt_ here is the correspondences between pairs of images
                        # and are different from the actual delta movement of fourpoints that we want to find
                        # This ground truth is used for evaluating the estimated homography on real image data.
                        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)
                        # Debug test splits
                        #test_synthetic_dataloader(data_loader, True, I1_splits[i], I2_splits[i], I_splits[i], I_prime_splits[i], pts1_splits[i], gt_splits[i], patch_indices_splits[i])

                        reuse_variables = True
                        rec_loss = model.rec_loss
                        ssim_loss = model.ssim_loss
                        l1_loss = model.l1_loss
                        l1_smooth_loss = model.l1_smooth_loss
                        if i == 0:
                            self.pred_I2 = model.pred_I2
                            self.I = model.I
                            self.I_prime = model.I_prime

                            self.I1_aug = model.I1_aug
                            self.I2_aug = model.I2_aug
                            self.pred_h4p = model.pred_h4p
                            self.gt_corr = model.gt
                            self.pts1 = model.pts_1
                        else:
                            self.pred_I2 = tf.concat(
                                [self.pred_I2, model.pred_I2], axis=0)
                            self.I = tf.concat([self.I, model.I], axis=0)
                            self.I_prime = tf.concat(
                                [self.I_prime, model.I_prime], axis=0)
                            self.I1_aug = tf.concat(
                                [self.I1_aug, model.I1_aug], axis=0)
                            self.I2_aug = tf.concat(
                                [self.I2_aug, model.I2_aug], axis=0)
                            self.pred_h4p = tf.concat(
                                [self.pred_h4p, model.pred_h4p], axis=0)
                            self.gt_corr = tf.concat([self.gt_corr, model.gt],
                                                     axis=0)
                            self.pts1 = tf.concat([self.pts1, model.pts_1],
                                                  axis=0)
                        rec_losses.append(rec_loss)
                        ssim_losses.append(ssim_loss)
                        l1_losses.append(l1_loss)
                        l1_smooth_losses.append(l1_smooth_loss)
            self.total_rec_loss = tf.reduce_mean(rec_losses)
            self.total_ssim_loss = tf.reduce_mean(ssim_losses)
            self.total_l1_loss = tf.reduce_mean(l1_losses)
            self.total_l1_smooth_loss = tf.reduce_mean(l1_smooth_losses)
            self.full_I = full_I_batch
            self.full_I_prime = full_I_prime_batch

            with tf.name_scope('Losses'):
                tf.summary.scalar('Total_rec_loss', self.total_rec_loss)
                tf.summary.scalar('Total_ssim_loss', self.total_ssim_loss)
                tf.summary.scalar('Total_l1_loss', self.total_l1_loss)
                tf.summary.scalar('Total_l1_smooth_loss',
                                  self.total_l1_smooth_loss)
            self.summary_opt = tf.summary.merge_all()
Пример #6
0
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)