hsequences_bench.py

import os, pickle
import numpy as np
from tqdm import tqdm
from COKENET.config_hpatches import get_eval_config
import HSequences_bench.tools.aux_tools as aux
import HSequences_bench.tools.geometry_tools as geo_tools
import HSequences_bench.tools.repeatability_tools as rep_tools
import HSequences_bench.tools.matching_tools as match_tools
from HSequences_bench.tools.HSequences_reader import HSequences_dataset
from HSequences_bench.tools.opencv_matcher import OpencvBruteForceMatcher

def hsequences_metrics():
    args = get_eval_config()

    print(args.split)
    aux.check_directory(args.results_bench_dir)

    # create the dataloader
    data_loader = HSequences_dataset(args.data_dir, args.split, args.split_path)
    results = aux.create_overlapping_results(args.detector_name, args.overlap)

    # matching method
    matcher = OpencvBruteForceMatcher('l2')
    count_seq = 0

    # load data and compute the keypoints
    iterate = tqdm(enumerate(data_loader.extract_hsequences()), total=len(data_loader.sequences), desc="HPatches Eval")
    print("HPatches evaluation using Keypoints and Descriptors.")
    for sample_id, sample_data in iterate:
        sequence = sample_data['sequence_name']

        count_seq += 1
        image_src = sample_data['im_src']
        images_dst = sample_data['images_dst']
        h_src_2_dst = sample_data['h_src_2_dst']
        h_dst_2_src = sample_data['h_dst_2_src']

        for idx_im in range(len(images_dst)):
            
            # create the mask to filter out the points outside of the common areas
            mask_src, mask_dst = geo_tools.create_common_region_masks(h_dst_2_src[idx_im], image_src.shape, images_dst[idx_im].shape)

            # compute the files paths
            src_pts_filename = os.path.join(args.results_dir,args.detector_name,
                            'hpatches-sequences-release', '{}/1.ppm.kpt.npy'.format(sample_data['sequence_name']))
            src_dsc_filename = os.path.join(args.results_dir,args.detector_name,
                            'hpatches-sequences-release', '{}/1.ppm.dsc.npy'.format(sample_data['sequence_name']))
            dst_pts_filename = os.path.join(args.results_dir,args.detector_name,
                            'hpatches-sequences-release', '{}/{}.ppm.kpt.npy'.format(sample_data['sequence_name'], idx_im+2))
            dst_dsc_filename = os.path.join(args.results_dir,args.detector_name,
                            'hpatches-sequences-release', '{}/{}.ppm.dsc.npy'.format(sample_data['sequence_name'], idx_im+2))

            if not os.path.isfile(src_pts_filename):
                print("Could not find the file: " + src_pts_filename)
                return False

            if not os.path.isfile(src_dsc_filename):
                print("Could not find the file: " + src_dsc_filename)
                return False

            if not os.path.isfile(dst_pts_filename):
                print("Could not find the file: " + dst_pts_filename)
                return False

            if not os.path.isfile(dst_dsc_filename):
                print("Could not find the file: " + dst_dsc_filename)
                return False

            # load the points
            src_pts = np.load(src_pts_filename)
            src_dsc = np.load(src_dsc_filename)

            dst_pts = np.load(dst_pts_filename)
            dst_dsc = np.load(dst_dsc_filename)

            if args.order_coord == 'xysr':
                src_pts = np.asarray(list(map(lambda x: [x[1], x[0], x[2], x[3]], src_pts)))
                dst_pts = np.asarray(list(map(lambda x: [x[1], x[0], x[2], x[3]], dst_pts)))

            # Check Common Points
            try:
                src_idx = rep_tools.check_common_points(src_pts, mask_src)
                # print(src_idx)

                src_pts = src_pts[src_idx]
                src_dsc = src_dsc[src_idx]

                dst_idx = rep_tools.check_common_points(dst_pts, mask_dst)
                dst_pts = dst_pts[dst_idx]
                dst_dsc = dst_dsc[dst_idx]
            except:
                continue

            # Select top K points
            if args.top_k_points:
                src_idx = rep_tools.select_top_k(src_pts, args.top_k_points)
                src_pts = src_pts[src_idx]
                src_dsc = src_dsc[src_idx]

                dst_idx = rep_tools.select_top_k(dst_pts, args.top_k_points)
                dst_pts = dst_pts[dst_idx]
                dst_dsc = dst_dsc[dst_idx]

            src_pts = np.asarray(list(map(lambda x: [x[1], x[0], x[2], x[3]], src_pts)))
            dst_pts = np.asarray(list(map(lambda x: [x[1], x[0], x[2], x[3]], dst_pts)))

            src_to_dst_pts = geo_tools.apply_homography_to_points(
                src_pts, h_src_2_dst[idx_im])

            dst_to_src_pts = geo_tools.apply_homography_to_points(
                dst_pts, h_dst_2_src[idx_im])

            if args.dst_to_src_evaluation:
                points_src = src_pts
                points_dst = dst_to_src_pts
            else:
                points_src = src_to_dst_pts
                points_dst = dst_pts

            # compute repeatability
            repeatability_results = rep_tools.compute_repeatability(points_src, points_dst, overlap_err=1-args.overlap,
                                                                    dist_match_thresh=args.pixel_threshold)

            # match descriptors
            matches = matcher.match(src_dsc, dst_dsc)
            matches_np = aux.convert_opencv_matches_to_numpy(matches)

            matches_inv = matcher.match(dst_dsc, src_dsc)
            matches_inv_np = aux.convert_opencv_matches_to_numpy(matches_inv)

            mask = matches_np[:, 0] == matches_inv_np[matches_np[:, 1], 1]
            matches_np = matches_np[mask]

            match_score, match_score_corr, num_matches = {}, {}, {}

            # compute matching based on pixel distance
            for th_i in range(1, 11):
                match_score_i, match_score_corr_i, num_matches_i = match_tools.compute_matching_based_distance(points_src, points_dst, matches_np,
                                                                       repeatability_results['total_num_points'],
                                                                       pixel_threshold=th_i,
                                                                       possible_matches=repeatability_results['possible_matches'])
                match_score[str(th_i)] = match_score_i
                match_score_corr[str(th_i)] = match_score_corr_i
                num_matches[str(th_i)] = num_matches_i

            mma = np.mean([match_score[str(idx)] for idx in match_score])

            results['rep_single_scale'].append(
                repeatability_results['rep_single_scale'])
            results['rep_multi_scale'].append(
                repeatability_results['rep_multi_scale'])
            results['num_points_single_scale'].append(
                repeatability_results['num_points_single_scale'])
            results['num_points_multi_scale'].append(
                repeatability_results['num_points_multi_scale'])
            results['error_overlap_single_scale'].append(
                repeatability_results['error_overlap_single_scale'])
            results['error_overlap_multi_scale'].append(
                repeatability_results['error_overlap_multi_scale'])

            results['mma'].append(match_score[str(args.pixel_threshold)])
            results['mma_corr'].append(match_score_corr[str(args.pixel_threshold)])
            results['num_matches'].append(num_matches[str(args.pixel_threshold)])
            results['num_mutual_corresp'].append(len(matches_np))
            results['avg_mma'].append(mma)
            results['num_features'].append(repeatability_results['total_num_points'])

            ## logging
            iterate.set_description("{}  {} / {} - {} rep_s {:.2f} , rep_m {:.2f}, p_s {:d} , p_m {:d}, eps_s {:.2f}, eps_m {:.2f} "
                .format(sequence, count_seq, len(data_loader.sequences), idx_im,
                    repeatability_results['rep_single_scale'], repeatability_results['rep_multi_scale'], repeatability_results['num_points_single_scale'], 
                    repeatability_results['num_points_multi_scale'], repeatability_results['error_overlap_single_scale'], repeatability_results['error_overlap_multi_scale']
            ) )


    # average the results
    rep_single = np.array(results['rep_single_scale']).mean()
    rep_multi = np.array(results['rep_multi_scale']).mean()
    error_overlap_s = np.array(results['error_overlap_single_scale']).mean()
    error_overlap_m = np.array(results['error_overlap_multi_scale']).mean()
    mma = np.array(results['mma']).mean()
    mma_corr = np.array(results['mma_corr']).mean()
    num_matches = np.array(results['num_matches']).mean()
    num_mutual_corresp = np.array(results['num_mutual_corresp']).mean()
    avg_mma = np.array(results['avg_mma']).mean()
    num_features = np.array(results['num_features']).mean()

    # Matching Score: Matching Score taking into account all features that have been
    # detected in any of the two images.
    # Matching Score (possible matches): Matching Score only taking into account those features that have been
    # detected in both images.
    # MMA Score is computed based on the Matching Score (all detected features)

    print('\n## Overlap @{0}:\n \
           #### Rep. Multi: {1:.4f}\n \
           #### Rep. Single: {2:.4f}\n \
           #### Overlap Multi: {3:.4f}\n \
           #### Overlap Single: {4:.4f}\n \
           #### MMA: {5:.4f}\n \
           #### MMA (possible matches): {6:.4f}\n \
           #### Num matches: {7:.4f}\n \
           #### Num Mutual Correspondences: {8:.4f}\n \
           #### Avg. over Threshold MMA: {9:.4f}\n \
           #### Num Feats: {10:.4f}'.format(
        args.overlap, rep_multi, rep_single, error_overlap_s, error_overlap_m, mma,
        mma_corr, num_matches, num_mutual_corresp, avg_mma, num_features))

    # Store data (serialize)
    output_file_path = os.path.join(args.results_bench_dir, '{0}_{1}.pickle'
        .format(args.detector_name, args.split))
    with open(output_file_path, 'wb') as handle:
            pickle.dump(results, handle, protocol=pickle.HIGHEST_PROTOCOL)


if __name__ == '__main__':

    hsequences_metrics()