imcui/third_party/rdd/benchmarks/mega_1500.py

import sys
sys.path.append(".")
import numpy as np
import torch
from PIL import Image
import tqdm
import cv2
import argparse
from RDD.RDD_helper import RDD_helper
from RDD.RDD import build
import matplotlib.pyplot as plt
import matplotlib
import os
from benchmarks.utils import pose_auc, angle_error_vec, angle_error_mat, symmetric_epipolar_distance, compute_symmetrical_epipolar_errors, compute_pose_error, compute_relative_pose, estimate_pose, dynamic_alpha

def make_matching_figure(
        img0, img1, mkpts0, mkpts1, color,
        kpts0=None, kpts1=None, text=[], dpi=75, path=None):
    # draw image pair
    assert mkpts0.shape[0] == mkpts1.shape[0], f'mkpts0: {mkpts0.shape[0]} v.s. mkpts1: {mkpts1.shape[0]}'
    fig, axes = plt.subplots(1, 2, figsize=(10, 6), dpi=dpi)
    axes[0].imshow(img0, cmap='gray')
    axes[1].imshow(img1, cmap='gray')
    for i in range(2):   # clear all frames
        axes[i].get_yaxis().set_ticks([])
        axes[i].get_xaxis().set_ticks([])
        for spine in axes[i].spines.values():
            spine.set_visible(False)
    plt.tight_layout(pad=1)
    
    if kpts0 is not None:
        assert kpts1 is not None
        axes[0].scatter(kpts0[:, 0], kpts0[:, 1], c='w', s=2)
        axes[1].scatter(kpts1[:, 0], kpts1[:, 1], c='w', s=2)

    # draw matches
    if mkpts0.shape[0] != 0 and mkpts1.shape[0] != 0:
        fig.canvas.draw()
        transFigure = fig.transFigure.inverted()
        fkpts0 = transFigure.transform(axes[0].transData.transform(mkpts0))
        fkpts1 = transFigure.transform(axes[1].transData.transform(mkpts1))
        fig.lines = [matplotlib.lines.Line2D((fkpts0[i, 0], fkpts1[i, 0]),
                                            (fkpts0[i, 1], fkpts1[i, 1]),
                                            transform=fig.transFigure, c=color[i], linewidth=1)
                                        for i in range(len(mkpts0))]
        
        axes[0].scatter(mkpts0[:, 0], mkpts0[:, 1], c=color, s=4)
        axes[1].scatter(mkpts1[:, 0], mkpts1[:, 1], c=color, s=4)

    # put txts
    txt_color = 'k' if img0[:100, :200].mean() > 200 else 'w'
    fig.text(
        0.01, 0.99, '\n'.join(text), transform=fig.axes[0].transAxes,
        fontsize=15, va='top', ha='left', color=txt_color)

    # save or return figure
    if path:
        plt.savefig(str(path), bbox_inches='tight', pad_inches=0)
        plt.close()
    else:
        return fig

def error_colormap(err, thr, alpha=1.0):
    assert alpha <= 1.0 and alpha > 0, f"Invaid alpha value: {alpha}"
    x = 1 - np.clip(err / (thr * 2), 0, 1)
    return np.clip(
        np.stack([2-x*2, x*2, np.zeros_like(x), np.ones_like(x)*alpha], -1), 0, 1)

def _make_evaluation_figure(img0, img1, kpts0, kpts1, epi_errs, e_t, e_R, alpha='dynamic', path=None):
    conf_thr = 1e-4
    
    img0 = np.array(img0)
    img1 = np.array(img1)
    
    kpts0 = kpts0
    kpts1 = kpts1
    
    epi_errs = epi_errs.cpu().numpy()
    correct_mask = epi_errs < conf_thr
    precision = np.mean(correct_mask) if len(correct_mask) > 0 else 0
    n_correct = np.sum(correct_mask)
    
    # recall might be larger than 1, since the calculation of conf_matrix_gt
    # uses groundtruth depths and camera poses, but epipolar distance is used here.

    # matching info
    if alpha == 'dynamic':
        alpha = dynamic_alpha(len(correct_mask))
    color = error_colormap(epi_errs, conf_thr, alpha=alpha)
    
    text = [
        f'#Matches {len(kpts0)}',
        f'Precision({conf_thr:.2e}) ({100 * precision:.1f}%): {n_correct}/{len(kpts0)}',
        f'e_t: {e_t:.2f} | e_R: {e_R:.2f}',
    ]
    
    # make the figure
    figure = make_matching_figure(img0, img1, kpts0, kpts1,
                                  color, text=text, path=path)
    return figure

class MegaDepthPoseMNNBenchmark:
    def __init__(self, data_root="./megadepth_test_1500", scene_names = None) -> None:
        if scene_names is None:
            self.scene_names = [
                "0015_0.1_0.3.npz",
                "0015_0.3_0.5.npz",
                "0022_0.1_0.3.npz",
                "0022_0.3_0.5.npz",
                "0022_0.5_0.7.npz",
            ]

        else:
            self.scene_names = scene_names
        self.scenes = [
            np.load(f"{data_root}/{scene}", allow_pickle=True)
            for scene in self.scene_names
        ]
        self.data_root = data_root

    def benchmark(self, model_helper, model_name = None, scale_intrinsics = False, calibrated = True, plot_every_iter=1, plot=False, method='sparse'):
        
        with torch.no_grad():
            data_root = self.data_root
            tot_e_t, tot_e_R, tot_e_pose = [], [], []
            thresholds = [5, 10, 20]
            for scene_ind in range(len(self.scenes)):
                import os
                scene_name = os.path.splitext(self.scene_names[scene_ind])[0]
                print(f"Processing {scene_name}")
                scene = self.scenes[scene_ind]
                pairs = scene["pair_infos"]
                intrinsics = scene["intrinsics"]
                poses = scene["poses"]
                im_paths = scene["image_paths"]
                pair_inds = range(len(pairs))
                for pairind in tqdm.tqdm(pair_inds):
                    idx0, idx1 = pairs[pairind][0]
                    K0 = intrinsics[idx0].copy()
                    T0 = poses[idx0].copy()
                    R0, t0 = T0[:3, :3], T0[:3, 3]
                    K1 = intrinsics[idx1].copy()
                    T1 = poses[idx1].copy()
                    R1, t1 = T1[:3, :3], T1[:3, 3]
                    R, t = compute_relative_pose(R0, t0, R1, t1)
                    T0_to_1 = np.concatenate((R,t[:,None]), axis=-1)
                    im_A_path = f"{data_root}/{im_paths[idx0]}"
                    im_B_path = f"{data_root}/{im_paths[idx1]}"
                    
                    im_A = cv2.imread(im_A_path)
                    im_B = cv2.imread(im_B_path)
                
                    if method == 'dense':
                        kpts0, kpts1, conf = model_helper.match_dense(im_A, im_B, thr=0.01, resize=1600)
                    elif method == 'lightglue':
                        kpts0, kpts1, conf = model_helper.match_lg(im_A, im_B, thr=0.01, resize=1600)
                    elif method == 'sparse':
                        kpts0, kpts1, conf = model_helper.match(im_A, im_B, thr=0.01, resize=1600)
                    else:
                        kpts0, kpts1, conf = model_helper.match_3rd_party(im_A, im_B, thr=0.01, resize=1600, model=method)

                    im_A = Image.open(im_A_path)
                    w0, h0 = im_A.size
                    im_B = Image.open(im_B_path)
                    w1, h1 = im_B.size
                    if scale_intrinsics:
                        scale0 = 840 / max(w0, h0)
                        scale1 = 840 / max(w1, h1)
                        w0, h0 = scale0 * w0, scale0 * h0
                        w1, h1 = scale1 * w1, scale1 * h1
                        K0, K1 = K0.copy(), K1.copy()
                        K0[:2] = K0[:2] * scale0
                        K1[:2] = K1[:2] * scale1
                    
                        
                    threshold = 0.5 
                    if calibrated:
                        norm_threshold = threshold / (np.mean(np.abs(K0[:2, :2])) + np.mean(np.abs(K1[:2, :2])))
                        ret = estimate_pose(
                            kpts0,
                            kpts1,
                            K0,
                            K1,
                            norm_threshold,
                            conf=0.99999,
                        )
                    if ret is not None:
                        R_est, t_est, mask = ret
                        T0_to_1_est = np.concatenate((R_est, t_est), axis=-1)  #
                        T0_to_1 = np.concatenate((R, t[:,None]), axis=-1)
                        e_t, e_R = compute_pose_error(T0_to_1_est, R, t)
                        
                        epi_errs = compute_symmetrical_epipolar_errors(T0_to_1, kpts0, kpts1, K0, K1)
                        if scene_ind % plot_every_iter == 0 and plot:

                            if not os.path.exists(f'outputs/mega_1500/{model_name}_{method}'):
                                os.mkdir(f'outputs/mega_1500/{model_name}_{method}')
                            name = f'outputs/mega_1500/{model_name}_{method}/{scene_name}_{pairind}.png'
                            _make_evaluation_figure(im_A, im_B, kpts0, kpts1, epi_errs, e_t, e_R, path=name)
                        e_pose = max(e_t, e_R)
                        
                        tot_e_t.append(e_t)
                        tot_e_R.append(e_R)
                        tot_e_pose.append(e_pose)
                            
            tot_e_pose = np.array(tot_e_pose)
            auc = pose_auc(tot_e_pose, thresholds)
            acc_5 = (tot_e_pose < 5).mean()
            acc_10 = (tot_e_pose < 10).mean()
            acc_15 = (tot_e_pose < 15).mean()
            acc_20 = (tot_e_pose < 20).mean()
            map_5 = acc_5
            map_10 = np.mean([acc_5, acc_10])
            map_20 = np.mean([acc_5, acc_10, acc_15, acc_20])
            print(f"{model_name} auc: {auc}")
            return {
                "auc_5": auc[0],
                "auc_10": auc[1],
                "auc_20": auc[2],
                "map_5": map_5,
                "map_10": map_10,
                "map_20": map_20,
            }
            
            
def parse_arguments():
    parser = argparse.ArgumentParser(description="Testing script.")
    
    parser.add_argument("--data_root", type=str, default="./data/megadepth_test_1500", help="Path to the MegaDepth dataset.")

    parser.add_argument("--weights", type=str, default="./weights/RDD-v2.pth", help="Path to the model checkpoint.")

    parser.add_argument("--plot", action="store_true", help="Whether to plot the results.")

    parser.add_argument("--method", type=str, default="sparse", help="Method for matching.")
    
    return parser.parse_args()

if __name__ == "__main__":
    args = parse_arguments()    
    if not os.path.exists('outputs'):
        os.mkdir('outputs')

    if not os.path.exists(f'outputs/mega_1500'):
        os.mkdir(f'outputs/mega_1500')
        
    model = build(weights=args.weights)
    benchmark = MegaDepthPoseMNNBenchmark(data_root=args.data_root)
    model.eval()
    model_helper = RDD_helper(model)
    with torch.no_grad():
        method = args.method
        out = benchmark.benchmark(model_helper, model_name='RDD', plot_every_iter=1, plot=args.plot, method=method)
        with open(f'outputs/mega_1500/RDD_{method}.txt', 'w') as f:
            f.write(str(out))