imcui/third_party/XoFTR/test_relative_pose.py

from collections import defaultdict, OrderedDict
import os
import os.path as osp
import numpy as np
from tqdm import tqdm
import argparse
import cv2
from pathlib import Path
import warnings
import json
import time

from src.utils.metrics import estimate_pose, relative_pose_error, error_auc, symmetric_epipolar_distance_numpy
from src.utils.plotting import dynamic_alpha, error_colormap, make_matching_figure


# Loading functions for methods
####################################################################
def load_xoftr(args):
    from src.xoftr import XoFTR
    from src.config.default import get_cfg_defaults
    from src.utils.data_io import DataIOWrapper, lower_config
    config = get_cfg_defaults(inference=True)
    config = lower_config(config)
    config["xoftr"]["match_coarse"]["thr"] = args.match_threshold
    config["xoftr"]["fine"]["thr"] = args.fine_threshold
    ckpt = args.ckpt
    matcher = XoFTR(config=config["xoftr"])
    matcher = DataIOWrapper(matcher, config=config["test"], ckpt=ckpt)
    return matcher.from_paths

####################################################################

def load_vis_tir_pairs_npz(npz_root, npz_list):
    """Load information for scene and image pairs from npz files.
    Args:
        npz_root: Directory path for npz files
        npz_list: File containing the names of the npz files to be used
    """
    with open(npz_list, 'r') as f:
        npz_names = [name.split()[0] for name in f.readlines()]
    print(f"Parse {len(npz_names)} npz from {npz_list}.")

    total_pairs = 0
    scene_pairs = {}
    
    for name in npz_names:
        print(f"Loading {name}")
        scene_info = np.load(f"{npz_root}/{name}", allow_pickle=True)
        pairs = []

        # Collect pairs
        for pair_info in scene_info['pair_infos']:
            total_pairs += 1
            (id0, id1) = pair_info
            im0 = scene_info['image_paths'][id0][0]
            im1 = scene_info['image_paths'][id1][1]
            K0 = scene_info['intrinsics'][id0][0].astype(np.float32)
            K1 = scene_info['intrinsics'][id1][1].astype(np.float32)

            dist0 =  np.array(scene_info['distortion_coefs'][id0][0], dtype=float)
            dist1 = np.array(scene_info['distortion_coefs'][id1][1], dtype=float)
            # Compute relative pose
            T0 = scene_info['poses'][id0]
            T1 = scene_info['poses'][id1]
            T_0to1 = np.matmul(T1, np.linalg.inv(T0))
            pairs.append({'im0':im0, 'im1':im1, 'dist0':dist0, 'dist1':dist1,
                          'K0':K0, 'K1':K1, 'T_0to1':T_0to1})
        scene_pairs[name] = pairs

    print(f"Loaded {total_pairs} pairs.")
    return scene_pairs



def save_matching_figure(path, img0, img1, mkpts0, mkpts1, inlier_mask, T_0to1, K0, K1,  t_err=None, R_err=None, name=None, conf_thr = 5e-4):
    """ Make and save matching figures
    """
    Tx = np.cross(np.eye(3), T_0to1[:3, 3])
    E_mat = Tx @ T_0to1[:3, :3]
    mkpts0_inliers = mkpts0[inlier_mask]
    mkpts1_inliers = mkpts1[inlier_mask] 
    if inlier_mask is not None and len(inlier_mask) != 0:
        epi_errs = symmetric_epipolar_distance_numpy(mkpts0_inliers, mkpts1_inliers, E_mat, K0, K1)

        correct_mask = epi_errs < conf_thr
        precision = np.mean(correct_mask) if len(correct_mask) > 0 else 0
        n_correct = np.sum(correct_mask)

        # matching info
        alpha = dynamic_alpha(len(correct_mask))
        color = error_colormap(epi_errs, conf_thr, alpha=alpha)
        text_precision =[
        f'Precision({conf_thr:.2e}) ({100 * precision:.1f}%): {n_correct}/{len(mkpts0_inliers)}']
    else:
        text_precision =[
        f'No inliers after ransac']

    if name is not None:
        text=[name]
    else:
        text = []
    
    if t_err is not None and R_err is not None:
        error_text = [f"err_t: {t_err:.2f} °", f"err_R: {R_err:.2f} °"]   
        text +=error_text
    
    text += text_precision
    
    # make the figure
    figure = make_matching_figure(img0, img1, mkpts0_inliers, mkpts1_inliers,
                                  color, text=text, path=path, dpi=150)


def aggregiate_scenes(scene_pose_auc, thresholds):
    """Averages the auc results for cloudy_cloud and cloudy_sunny scenes
    """
    temp_pose_auc = {}
    for npz_name in scene_pose_auc.keys():
        scene_name = npz_name.split("_scene")[0]
        temp_pose_auc[scene_name] = [np.zeros(len(thresholds), dtype=np.float32), 0] # [sum, total_number]
    for npz_name in scene_pose_auc.keys():
        scene_name = npz_name.split("_scene")[0]
        temp_pose_auc[scene_name][0] += scene_pose_auc[npz_name]
        temp_pose_auc[scene_name][1] += 1
    
    agg_pose_auc = {}
    for scene_name in temp_pose_auc.keys():
        agg_pose_auc[scene_name] = temp_pose_auc[scene_name][0] / temp_pose_auc[scene_name][1]
    
    return agg_pose_auc

def eval_relapose(
    matcher,
    data_root,
    scene_pairs,
    ransac_thres,
    thresholds,
    save_figs,
    figures_dir=None,
    method=None,
    print_out=False,
    debug=False,
):
    scene_pose_auc = {}
    for scene_name in scene_pairs.keys():
        scene_dir =  osp.join(figures_dir, scene_name.split(".")[0])
        if save_figs and not osp.exists(scene_dir):
            os.makedirs(scene_dir)

        pairs = scene_pairs[scene_name]
        statis = defaultdict(list)
        np.set_printoptions(precision=2)

        # Eval on pairs
        print(f"\nStart evaluation on VisTir \n")
        for i, pair in tqdm(enumerate(pairs), smoothing=.1, total=len(pairs)):
            if debug and i > 10:
                break

            T_0to1 = pair['T_0to1']
            im0 = str(data_root / pair['im0'])
            im1 = str(data_root / pair['im1'])
            match_res = matcher(im0, im1, pair['K0'], pair['K1'], pair['dist0'], pair['dist1'])
            matches = match_res['matches']
            new_K0 = match_res['new_K0']
            new_K1 = match_res['new_K1']
            mkpts0 = match_res['mkpts0']
            mkpts1 = match_res['mkpts1']

            # Calculate pose errors
            ret = estimate_pose(
                mkpts0, mkpts1, new_K0, new_K1, thresh=ransac_thres
            )
            
            if ret is None:
                R, t, inliers = None, None, None
                t_err, R_err = np.inf, np.inf
                statis['failed'].append(i)
                statis['R_errs'].append(R_err)
                statis['t_errs'].append(t_err)
                statis['inliers'].append(np.array([]).astype(np.bool_))
            else:
                R, t, inliers = ret
                t_err, R_err = relative_pose_error(T_0to1, R, t)
                statis['R_errs'].append(R_err)
                statis['t_errs'].append(t_err)
                statis['inliers'].append(inliers.sum() / len(mkpts0))
                if print_out:
                    print(f"#M={len(matches)} R={R_err:.3f}, t={t_err:.3f}")
            
            if save_figs:
                img0_name = f"{'vis' if 'visible' in pair['im0'] else 'tir'}_{osp.basename(pair['im0']).split('.')[0]}"
                img1_name = f"{'vis' if 'visible' in pair['im1'] else 'tir'}_{osp.basename(pair['im1']).split('.')[0]}"
                fig_path = osp.join(scene_dir, f"{img0_name}_{img1_name}.jpg")
                save_matching_figure(path=fig_path, 
                                    img0=match_res['img0_undistorted'] if 'img0_undistorted' in match_res.keys() else match_res['img0'], 
                                    img1=match_res['img1_undistorted'] if 'img1_undistorted' in match_res.keys() else match_res['img1'], 
                                    mkpts0=mkpts0,
                                    mkpts1=mkpts1,
                                    inlier_mask=inliers,
                                    T_0to1=T_0to1,
                                    K0=new_K0,
                                    K1=new_K1,
                                    t_err=t_err,
                                    R_err=R_err,
                                    name=method
                                    ) 

        print(f"Scene: {scene_name} Total samples: {len(pairs)} Failed:{len(statis['failed'])}. \n")
        pose_errors = np.max(np.stack([statis['R_errs'], statis['t_errs']]), axis=0)
        pose_auc = error_auc(pose_errors, thresholds)  # (auc@5, auc@10, auc@20)
        scene_pose_auc[scene_name] = 100 * np.array([pose_auc[f'auc@{t}'] for t in thresholds])
        print(f"{scene_name} {pose_auc}")
    agg_pose_auc = aggregiate_scenes(scene_pose_auc, thresholds)
    return scene_pose_auc, agg_pose_auc

def test_relative_pose_vistir(
    data_root_dir,
    method="xoftr",
    exp_name = "VisTIR",
    ransac_thres=1.5,
    print_out=False,
    save_dir=None,
    save_figs=False,
    debug=False,
    args=None

):  
    if not osp.exists(osp.join(save_dir, method)):
        os.makedirs(osp.join(save_dir, method))

    counter = 0
    path = osp.join(save_dir, method, f"{exp_name}"+"_{}")
    while osp.exists(path.format(counter)):
        counter += 1
    exp_dir = path.format(counter)
    os.mkdir(exp_dir)
    results_file = osp.join(exp_dir, "results.json")
    figures_dir = osp.join(exp_dir, "match_figures")
    if save_figs:
        os.mkdir(figures_dir)

    # Init paths
    npz_root = data_root_dir / 'index/scene_info_test/'
    npz_list = data_root_dir / 'index/val_test_list/test_list.txt'
    data_root = data_root_dir 

    # Load pairs
    scene_pairs = load_vis_tir_pairs_npz(npz_root, npz_list)

    # Load method
    matcher = eval(f"load_{method}")(args)
    
    thresholds=[5, 10, 20]
    # Eval
    scene_pose_auc, agg_pose_auc = eval_relapose(
        matcher, 
        data_root,
        scene_pairs,
        ransac_thres=ransac_thres,
        thresholds=thresholds,
        save_figs=save_figs,
        figures_dir=figures_dir,
        method=method,
        print_out=print_out,
        debug=debug,
    )

    # Create result dict
    results = OrderedDict({"method": method, 
                        "exp_name": exp_name, 
                        "ransac_thres": ransac_thres,
                        "auc_thresholds": thresholds})
    results.update({key:value for key, value in vars(args).items() if key not in results})
    results.update({key:value.tolist() for key, value in agg_pose_auc.items()})
    results.update({key:value.tolist() for key, value in scene_pose_auc.items()})

    print(f"Results: {json.dumps(results, indent=4)}")
    
    # Save to json file
    with open(results_file, 'w') as outfile:
        json.dump(results, outfile, indent=4)

    print(f"Results saved to {results_file}")

if __name__ == '__main__':

    def add_common_arguments(parser):
        parser.add_argument('--gpu', '-gpu', type=str, default='0')
        parser.add_argument('--exp_name', type=str, default="VisTIR")
        parser.add_argument('--data_root_dir', type=str, default="./data/METU_VisTIR/")
        parser.add_argument('--save_dir', type=str, default="./results_relative_pose")
        parser.add_argument('--ransac_thres', type=float, default=1.5)
        parser.add_argument('--print_out', action='store_true')
        parser.add_argument('--debug', action='store_true')
        parser.add_argument('--save_figs', action='store_true')
    
    def add_xoftr_arguments(subparsers):
        subcommand = subparsers.add_parser('xoftr')
        subcommand.add_argument('--match_threshold', type=float, default=0.3)
        subcommand.add_argument('--fine_threshold', type=float, default=0.1)
        subcommand.add_argument('--ckpt', type=str, default="./weights/weights_xoftr_640.ckpt")
        add_common_arguments(subcommand)

    parser = argparse.ArgumentParser(description='Benchmark Relative Pose')
    add_common_arguments(parser)

    # Create subparsers for top-level commands
    subparsers = parser.add_subparsers(dest="method")
    add_xoftr_arguments(subparsers)
    
    args = parser.parse_args()

    os.environ['CUDA_VISIBLE_DEVICES'] = "0"
    tt = time.time()
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        test_relative_pose_vistir(
            Path(args.data_root_dir),
            args.method,
            args.exp_name,
            ransac_thres=args.ransac_thres,
            print_out=args.print_out,
            save_dir = args.save_dir,
            save_figs = args.save_figs,
            debug=args.debug,
            args=args
        )
    print(f"Elapsed time: {time.time() - tt}")