vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/keypt2subpx/submodules/glue_factory/gluefactory/eval/megadepth1500.py

@@ -0,0 +1,189 @@
+import logging
+import zipfile
+from collections import defaultdict
+from collections.abc import Iterable
+from pathlib import Path
+from pprint import pprint
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from omegaconf import OmegaConf
+from tqdm import tqdm
+
+from ..datasets import get_dataset
+from ..models.cache_loader import CacheLoader
+from ..settings import DATA_PATH, EVAL_PATH
+from ..utils.export_predictions import export_predictions
+from ..visualization.viz2d import plot_cumulative
+from .eval_pipeline import EvalPipeline
+from .io import get_eval_parser, load_model, parse_eval_args
+from .utils import eval_matches_epipolar, eval_poses, eval_relative_pose_robust
+
+logger = logging.getLogger(__name__)
+
+
+class MegaDepth1500Pipeline(EvalPipeline):
+    default_conf = {
+        "data": {
+            "name": "image_pairs",
+            "pairs": "megadepth1500/pairs_calibrated.txt",
+            "root": "megadepth1500/images/",
+            "extra_data": "relative_pose",
+            "preprocessing": {
+                "side": "long",
+            },
+        },
+        "model": {
+            "ground_truth": {
+                "name": None,  # remove gt matches
+            }
+        },
+        "eval": {
+            "estimator": "poselib",
+            "ransac_th": 1.0,  # -1 runs a bunch of thresholds and selects the best
+        },
+    }
+
+    export_keys = [
+        "keypoints0",
+        "keypoints1",
+        "keypoint_scores0",
+        "keypoint_scores1",
+        "matches0",
+        "matches1",
+        "matching_scores0",
+        "matching_scores1",
+    ]
+    optional_export_keys = []
+
+    def _init(self, conf):
+        if not (DATA_PATH / "megadepth1500").exists():
+            logger.info("Downloading the MegaDepth-1500 dataset.")
+            url = "https://cvg-data.inf.ethz.ch/megadepth/megadepth1500.zip"
+            zip_path = DATA_PATH / url.rsplit("/", 1)[-1]
+            zip_path.parent.mkdir(exist_ok=True, parents=True)
+            torch.hub.download_url_to_file(url, zip_path)
+            with zipfile.ZipFile(zip_path) as fid:
+                fid.extractall(DATA_PATH)
+            zip_path.unlink()
+
+    @classmethod
+    def get_dataloader(self, data_conf=None):
+        """Returns a data loader with samples for each eval datapoint"""
+        data_conf = data_conf if data_conf else self.default_conf["data"]
+        dataset = get_dataset(data_conf["name"])(data_conf)
+        return dataset.get_data_loader("test")
+
+    def get_predictions(self, experiment_dir, model=None, overwrite=False):
+        """Export a prediction file for each eval datapoint"""
+        pred_file = experiment_dir / "predictions.h5"
+        if not pred_file.exists() or overwrite:
+            if model is None:
+                model = load_model(self.conf.model, self.conf.checkpoint)
+            export_predictions(
+                self.get_dataloader(self.conf.data),
+                model,
+                pred_file,
+                keys=self.export_keys,
+                optional_keys=self.optional_export_keys,
+            )
+        return pred_file
+
+    def run_eval(self, loader, pred_file):
+        """Run the eval on cached predictions"""
+        conf = self.conf.eval
+        results = defaultdict(list)
+        test_thresholds = (
+            ([conf.ransac_th] if conf.ransac_th > 0 else [0.5, 1.0, 1.5, 2.0, 2.5, 3.0])
+            if not isinstance(conf.ransac_th, Iterable)
+            else conf.ransac_th
+        )
+        pose_results = defaultdict(lambda: defaultdict(list))
+        cache_loader = CacheLoader({"path": str(pred_file), "collate": None}).eval()
+        for i, data in enumerate(tqdm(loader)):
+            pred = cache_loader(data)
+            # add custom evaluations here
+            results_i = eval_matches_epipolar(data, pred)
+            for th in test_thresholds:
+                pose_results_i = eval_relative_pose_robust(
+                    data,
+                    pred,
+                    {"estimator": conf.estimator, "ransac_th": th},
+                )
+                [pose_results[th][k].append(v) for k, v in pose_results_i.items()]
+
+            # we also store the names for later reference
+            results_i["names"] = data["name"][0]
+            if "scene" in data.keys():
+                results_i["scenes"] = data["scene"][0]
+
+            for k, v in results_i.items():
+                results[k].append(v)
+
+        # summarize results as a dict[str, float]
+        # you can also add your custom evaluations here
+        summaries = {}
+        for k, v in results.items():
+            arr = np.array(v)
+            if not np.issubdtype(np.array(v).dtype, np.number):
+                continue
+            summaries[f"m{k}"] = round(np.mean(arr), 3)
+
+        best_pose_results, best_th = eval_poses(
+            pose_results, auc_ths=[5, 10, 20], key="rel_pose_error"
+        )
+        results = {**results, **pose_results[best_th]}
+        summaries = {
+            **summaries,
+            **best_pose_results,
+        }
+
+        figures = {
+            "pose_recall": plot_cumulative(
+                {self.conf.eval.estimator: results["rel_pose_error"]},
+                [0, 30],
+                unit="°",
+                title="Pose ",
+            )
+        }
+
+        return summaries, figures, results
+
+
+if __name__ == "__main__":
+    from .. import logger  # overwrite the logger
+
+    dataset_name = Path(__file__).stem
+    parser = get_eval_parser()
+    args = parser.parse_intermixed_args()
+
+    default_conf = OmegaConf.create(MegaDepth1500Pipeline.default_conf)
+
+    # mingle paths
+    output_dir = Path(EVAL_PATH, dataset_name)
+    output_dir.mkdir(exist_ok=True, parents=True)
+
+    name, conf = parse_eval_args(
+        dataset_name,
+        args,
+        "configs/",
+        default_conf,
+    )
+
+    experiment_dir = output_dir / name
+    experiment_dir.mkdir(exist_ok=True)
+
+    pipeline = MegaDepth1500Pipeline(conf)
+    s, f, r = pipeline.run(
+        experiment_dir,
+        overwrite=args.overwrite,
+        overwrite_eval=args.overwrite_eval,
+    )
+
+    pprint(s)
+
+    if args.plot:
+        for name, fig in f.items():
+            fig.canvas.manager.set_window_title(name)
+        plt.show()

vismatch/third_party/keypt2subpx/submodules/glue_factory/gluefactory/eval/utils.py

@@ -0,0 +1,272 @@
+import numpy as np
+import torch
+from kornia.geometry.homography import find_homography_dlt
+
+from ..geometry.epipolar import generalized_epi_dist, relative_pose_error
+from ..geometry.gt_generation import IGNORE_FEATURE
+from ..geometry.homography import homography_corner_error, sym_homography_error
+from ..robust_estimators import load_estimator
+from ..utils.tensor import index_batch
+from ..utils.tools import AUCMetric
+
+
+def check_keys_recursive(d, pattern):
+    if isinstance(pattern, dict):
+        {check_keys_recursive(d[k], v) for k, v in pattern.items()}
+    else:
+        for k in pattern:
+            assert k in d.keys()
+
+
+def get_matches_scores(kpts0, kpts1, matches0, mscores0):
+    m0 = matches0 > -1
+    m1 = matches0[m0]
+    pts0 = kpts0[m0]
+    pts1 = kpts1[m1]
+    scores = mscores0[m0]
+    return pts0, pts1, scores
+
+
+def eval_per_batch_item(data: dict, pred: dict, eval_f, *args, **kwargs):
+    # Batched data
+    results = [
+        eval_f(data_i, pred_i, *args, **kwargs)
+        for data_i, pred_i in zip(index_batch(data), index_batch(pred))
+    ]
+    # Return a dictionary of lists with the evaluation of each item
+    return {k: [r[k] for r in results] for k in results[0].keys()}
+
+
+def eval_matches_epipolar(data: dict, pred: dict) -> dict:
+    check_keys_recursive(data, ["view0", "view1", "T_0to1"])
+    check_keys_recursive(
+        pred, ["keypoints0", "keypoints1", "matches0", "matching_scores0"]
+    )
+
+    kp0, kp1 = pred["keypoints0"], pred["keypoints1"]
+    m0, scores0 = pred["matches0"], pred["matching_scores0"]
+    pts0, pts1, scores = get_matches_scores(kp0, kp1, m0, scores0)
+
+    results = {}
+
+    # match metrics
+    n_epi_err = generalized_epi_dist(
+        pts0[None],
+        pts1[None],
+        data["view0"]["camera"],
+        data["view1"]["camera"],
+        data["T_0to1"],
+        False,
+        essential=True,
+    )[0]
+    results["epi_prec@1e-4"] = (n_epi_err < 1e-4).float().mean()
+    results["epi_prec@5e-4"] = (n_epi_err < 5e-4).float().mean()
+    results["epi_prec@1e-3"] = (n_epi_err < 1e-3).float().mean()
+
+    results["num_matches"] = pts0.shape[0]
+    results["num_keypoints"] = (kp0.shape[0] + kp1.shape[0]) / 2.0
+
+    return results
+
+
+def eval_matches_homography(data: dict, pred: dict) -> dict:
+    check_keys_recursive(data, ["H_0to1"])
+    check_keys_recursive(
+        pred, ["keypoints0", "keypoints1", "matches0", "matching_scores0"]
+    )
+
+    H_gt = data["H_0to1"]
+    if H_gt.ndim > 2:
+        return eval_per_batch_item(data, pred, eval_matches_homography)
+
+    kp0, kp1 = pred["keypoints0"], pred["keypoints1"]
+    m0, scores0 = pred["matches0"], pred["matching_scores0"]
+    pts0, pts1, scores = get_matches_scores(kp0, kp1, m0, scores0)
+    err = sym_homography_error(pts0, pts1, H_gt)
+    results = {}
+    results["prec@1px"] = (err < 1).float().mean().nan_to_num().item()
+    results["prec@3px"] = (err < 3).float().mean().nan_to_num().item()
+    results["num_matches"] = pts0.shape[0]
+    results["num_keypoints"] = (kp0.shape[0] + kp1.shape[0]) / 2.0
+    return results
+
+
+def eval_relative_pose_robust(data, pred, conf):
+    check_keys_recursive(data, ["view0", "view1", "T_0to1"])
+    check_keys_recursive(
+        pred, ["keypoints0", "keypoints1", "matches0", "matching_scores0"]
+    )
+
+    T_gt = data["T_0to1"]
+    kp0, kp1 = pred["keypoints0"], pred["keypoints1"]
+    m0, scores0 = pred["matches0"], pred["matching_scores0"]
+    pts0, pts1, scores = get_matches_scores(kp0, kp1, m0, scores0)
+
+    results = {}
+
+    estimator = load_estimator("relative_pose", conf["estimator"])(conf)
+    data_ = {
+        "m_kpts0": pts0,
+        "m_kpts1": pts1,
+        "camera0": data["view0"]["camera"][0],
+        "camera1": data["view1"]["camera"][0],
+    }
+    est = estimator(data_)
+
+    if not est["success"]:
+        results["rel_pose_error"] = float("inf")
+        results["ransac_inl"] = 0
+        results["ransac_inl%"] = 0
+    else:
+        # R, t, inl = ret
+        M = est["M_0to1"]
+        inl = est["inliers"].numpy()
+        t_error, r_error = relative_pose_error(T_gt, M.R, M.t)
+        results["rel_pose_error"] = max(r_error, t_error)
+        results["ransac_inl"] = np.sum(inl)
+        results["ransac_inl%"] = np.mean(inl)
+
+    return results
+
+
+def eval_homography_robust(data, pred, conf):
+    H_gt = data["H_0to1"]
+    if H_gt.ndim > 2:
+        return eval_per_batch_item(data, pred, eval_relative_pose_robust, conf)
+
+    estimator = load_estimator("homography", conf["estimator"])(conf)
+
+    data_ = {}
+    if "keypoints0" in pred:
+        kp0, kp1 = pred["keypoints0"], pred["keypoints1"]
+        m0, scores0 = pred["matches0"], pred["matching_scores0"]
+        pts0, pts1, _ = get_matches_scores(kp0, kp1, m0, scores0)
+        data_["m_kpts0"] = pts0
+        data_["m_kpts1"] = pts1
+    if "lines0" in pred:
+        if "orig_lines0" in pred:
+            lines0 = pred["orig_lines0"]
+            lines1 = pred["orig_lines1"]
+        else:
+            lines0 = pred["lines0"]
+            lines1 = pred["lines1"]
+        m_lines0, m_lines1, _ = get_matches_scores(
+            lines0, lines1, pred["line_matches0"], pred["line_matching_scores0"]
+        )
+        data_["m_lines0"] = m_lines0
+        data_["m_lines1"] = m_lines1
+
+    est = estimator(data_)
+    if est["success"]:
+        M = est["M_0to1"]
+        error_r = homography_corner_error(M, H_gt, data["view0"]["image_size"]).item()
+    else:
+        error_r = float("inf")
+
+    results = {}
+    results["H_error_ransac"] = error_r
+    if "inliers" in est:
+        inl = est["inliers"]
+        results["ransac_inl"] = inl.float().sum().item()
+        results["ransac_inl%"] = inl.float().sum().item() / max(len(inl), 1)
+
+    return results
+
+
+def eval_homography_dlt(data, pred):
+    H_gt = data["H_0to1"]
+    H_inf = torch.ones_like(H_gt) * float("inf")
+
+    kp0, kp1 = pred["keypoints0"], pred["keypoints1"]
+    m0, scores0 = pred["matches0"], pred["matching_scores0"]
+    pts0, pts1, scores = get_matches_scores(kp0, kp1, m0, scores0)
+    scores = scores.to(pts0)
+    results = {}
+    try:
+        if H_gt.ndim == 2:
+            pts0, pts1, scores = pts0[None], pts1[None], scores[None]
+        h_dlt = find_homography_dlt(pts0, pts1, scores)
+        if H_gt.ndim == 2:
+            h_dlt = h_dlt[0]
+    except AssertionError:
+        h_dlt = H_inf
+
+    error_dlt = homography_corner_error(h_dlt, H_gt, data["view0"]["image_size"])
+    results["H_error_dlt"] = error_dlt.item()
+    return results
+
+
+def eval_poses(pose_results, auc_ths, key, unit="°"):
+    pose_aucs = {}
+    best_th = -1
+    for th, results_i in pose_results.items():
+        pose_aucs[th] = AUCMetric(auc_ths, results_i[key]).compute()
+    mAAs = {k: np.mean(v) for k, v in pose_aucs.items()}
+    best_th = max(mAAs, key=mAAs.get)
+
+    if len(pose_aucs) > -1:
+        print("Tested ransac setup with following results:")
+        print("AUC", pose_aucs)
+        print("mAA", mAAs)
+        print("best threshold =", best_th)
+
+    summaries = {}
+
+    for i, ath in enumerate(auc_ths):
+        summaries[f"{key}@{ath}{unit}"] = pose_aucs[best_th][i]
+    summaries[f"{key}_mAA"] = mAAs[best_th]
+
+    for k, v in pose_results[best_th].items():
+        arr = np.array(v)
+        if not np.issubdtype(np.array(v).dtype, np.number):
+            continue
+        summaries[f"m{k}"] = round(np.median(arr), 3)
+    return summaries, best_th
+
+
+def get_tp_fp_pts(pred_matches, gt_matches, pred_scores):
+    """
+    Computes the True Positives (TP), False positives (FP), the score associated
+    to each match and the number of positives for a set of matches.
+    """
+    assert pred_matches.shape == pred_scores.shape
+    ignore_mask = gt_matches != IGNORE_FEATURE
+    pred_matches, gt_matches, pred_scores = (
+        pred_matches[ignore_mask],
+        gt_matches[ignore_mask],
+        pred_scores[ignore_mask],
+    )
+    num_pos = np.sum(gt_matches != -1)
+    pred_positives = pred_matches != -1
+    tp = pred_matches[pred_positives] == gt_matches[pred_positives]
+    fp = pred_matches[pred_positives] != gt_matches[pred_positives]
+    scores = pred_scores[pred_positives]
+    return tp, fp, scores, num_pos
+
+
+def AP(tp, fp):
+    recall = tp
+    precision = tp / np.maximum(tp + fp, 1e-9)
+    recall = np.concatenate(([0.0], recall, [1.0]))
+    precision = np.concatenate(([0.0], precision, [0.0]))
+    for i in range(precision.size - 1, 0, -1):
+        precision[i - 1] = max(precision[i - 1], precision[i])
+    i = np.where(recall[1:] != recall[:-1])[0]
+    ap = np.sum((recall[i + 1] - recall[i]) * precision[i + 1])
+    return ap
+
+
+def aggregate_pr_results(results, suffix=""):
+    tp_list = np.concatenate(results["tp" + suffix], axis=0)
+    fp_list = np.concatenate(results["fp" + suffix], axis=0)
+    scores_list = np.concatenate(results["scores" + suffix], axis=0)
+    n_gt = max(results["num_pos" + suffix], 1)
+
+    out = {}
+    idx = np.argsort(scores_list)[::-1]
+    tp_vals = np.cumsum(tp_list[idx]) / n_gt
+    fp_vals = np.cumsum(fp_list[idx]) / n_gt
+    out["curve_recall" + suffix] = tp_vals
+    out["curve_precision" + suffix] = tp_vals / np.maximum(tp_vals + fp_vals, 1e-9)
+    out["AP" + suffix] = AP(tp_vals, fp_vals) * 100
+    return out

vismatch/third_party/keypt2subpx/submodules/glue_factory/gluefactory/geometry/depth.py

@@ -0,0 +1,88 @@
+import kornia
+import torch
+
+from .utils import get_image_coords
+from .wrappers import Camera
+
+
+def sample_fmap(pts, fmap):
+    h, w = fmap.shape[-2:]
+    grid_sample = torch.nn.functional.grid_sample
+    pts = (pts / pts.new_tensor([[w, h]]) * 2 - 1)[:, None]
+    # @TODO: This might still be a source of noise --> bilinear interpolation dangerous
+    interp_lin = grid_sample(fmap, pts, align_corners=False, mode="bilinear")
+    interp_nn = grid_sample(fmap, pts, align_corners=False, mode="nearest")
+    return torch.where(torch.isnan(interp_lin), interp_nn, interp_lin)[:, :, 0].permute(
+        0, 2, 1
+    )
+
+
+def sample_depth(pts, depth_):
+    depth = torch.where(depth_ > 0, depth_, depth_.new_tensor(float("nan")))
+    depth = depth[:, None]
+    interp = sample_fmap(pts, depth).squeeze(-1)
+    valid = (~torch.isnan(interp)) & (interp > 0)
+    return interp, valid
+
+
+def sample_normals_from_depth(pts, depth, K):
+    depth = depth[:, None]
+    normals = kornia.geometry.depth.depth_to_normals(depth, K)
+    normals = torch.where(depth > 0, normals, 0.0)
+    interp = sample_fmap(pts, normals)
+    valid = (~torch.isnan(interp)) & (interp > 0)
+    return interp, valid
+
+
+def project(
+    kpi,
+    di,
+    depthj,
+    camera_i,
+    camera_j,
+    T_itoj,
+    validi,
+    ccth=None,
+    sample_depth_fun=sample_depth,
+    sample_depth_kwargs=None,
+):
+    if sample_depth_kwargs is None:
+        sample_depth_kwargs = {}
+
+    kpi_3d_i = camera_i.image2cam(kpi)
+    kpi_3d_i = kpi_3d_i * di[..., None]
+    kpi_3d_j = T_itoj.transform(kpi_3d_i)
+    kpi_j, validj = camera_j.cam2image(kpi_3d_j)
+    # di_j = kpi_3d_j[..., -1]
+    validi = validi & validj
+    if depthj is None or ccth is None:
+        return kpi_j, validi & validj
+    else:
+        # circle consistency
+        dj, validj = sample_depth_fun(kpi_j, depthj, **sample_depth_kwargs)
+        kpi_j_3d_j = camera_j.image2cam(kpi_j) * dj[..., None]
+        kpi_j_i, validj_i = camera_i.cam2image(T_itoj.inv().transform(kpi_j_3d_j))
+        consistent = ((kpi - kpi_j_i) ** 2).sum(-1) < ccth
+        visible = validi & consistent & validj_i & validj
+        # visible = validi
+        return kpi_j, visible
+
+
+def dense_warp_consistency(
+    depthi: torch.Tensor,
+    depthj: torch.Tensor,
+    T_itoj: torch.Tensor,
+    camerai: Camera,
+    cameraj: Camera,
+    **kwargs,
+):
+    kpi = get_image_coords(depthi).flatten(-3, -2)
+    di = depthi.flatten(
+        -2,
+    )
+    validi = di > 0
+    kpir, validir = project(kpi, di, depthj, camerai, cameraj, T_itoj, validi, **kwargs)
+
+    return kpir.unflatten(-2, depthi.shape[-2:]), validir.unflatten(
+        -1, (depthj.shape[-2:])
+    )

vismatch/third_party/keypt2subpx/submodules/glue_factory/gluefactory/geometry/epipolar.py

@@ -0,0 +1,155 @@
+import torch
+
+from .utils import skew_symmetric, to_homogeneous
+from .wrappers import Camera, Pose
+
+
+def T_to_E(T: Pose):
+    """Convert batched poses (..., 4, 4) to batched essential matrices."""
+    return skew_symmetric(T.t) @ T.R
+
+
+def T_to_F(cam0: Camera, cam1: Camera, T_0to1: Pose):
+    return E_to_F(cam0, cam1, T_to_E(T_0to1))
+
+
+def E_to_F(cam0: Camera, cam1: Camera, E: torch.Tensor):
+    assert cam0._data.shape[-1] == 6, "only pinhole cameras supported"
+    assert cam1._data.shape[-1] == 6, "only pinhole cameras supported"
+    K0 = cam0.calibration_matrix()
+    K1 = cam1.calibration_matrix()
+    return K1.inverse().transpose(-1, -2) @ E @ K0.inverse()
+
+
+def F_to_E(cam0: Camera, cam1: Camera, F: torch.Tensor):
+    assert cam0._data.shape[-1] == 6, "only pinhole cameras supported"
+    assert cam1._data.shape[-1] == 6, "only pinhole cameras supported"
+    K0 = cam0.calibration_matrix()
+    K1 = cam1.calibration_matrix()
+    return K1.transpose(-1, -2) @ F @ K0
+
+
+def sym_epipolar_distance(p0, p1, E, squared=True):
+    """Compute batched symmetric epipolar distances.
+    Args:
+        p0, p1: batched tensors of N 2D points of size (..., N, 2).
+        E: essential matrices from camera 0 to camera 1, size (..., 3, 3).
+    Returns:
+        The symmetric epipolar distance of each point-pair: (..., N).
+    """
+    assert p0.shape[-2] == p1.shape[-2]
+    if p0.shape[-2] == 0:
+        return torch.zeros(p0.shape[:-1]).to(p0)
+    if p0.shape[-1] != 3:
+        p0 = to_homogeneous(p0)
+    if p1.shape[-1] != 3:
+        p1 = to_homogeneous(p1)
+    p1_E_p0 = torch.einsum("...ni,...ij,...nj->...n", p1, E, p0)
+    E_p0 = torch.einsum("...ij,...nj->...ni", E, p0)
+    Et_p1 = torch.einsum("...ij,...ni->...nj", E, p1)
+    d0 = (E_p0[..., 0] ** 2 + E_p0[..., 1] ** 2).clamp(min=1e-6)
+    d1 = (Et_p1[..., 0] ** 2 + Et_p1[..., 1] ** 2).clamp(min=1e-6)
+    if squared:
+        d = p1_E_p0**2 * (1 / d0 + 1 / d1)
+    else:
+        d = p1_E_p0.abs() * (1 / d0.sqrt() + 1 / d1.sqrt()) / 2
+    return d
+
+
+def sym_epipolar_distance_all(p0, p1, E, eps=1e-15):
+    if p0.shape[-1] != 3:
+        p0 = to_homogeneous(p0)
+    if p1.shape[-1] != 3:
+        p1 = to_homogeneous(p1)
+    p1_E_p0 = torch.einsum("...mi,...ij,...nj->...nm", p1, E, p0).abs()
+    E_p0 = torch.einsum("...ij,...nj->...ni", E, p0)
+    Et_p1 = torch.einsum("...ij,...mi->...mj", E, p1)
+    d0 = p1_E_p0 / (E_p0[..., None, 0] ** 2 + E_p0[..., None, 1] ** 2 + eps).sqrt()
+    d1 = (
+        p1_E_p0
+        / (Et_p1[..., None, :, 0] ** 2 + Et_p1[..., None, :, 1] ** 2 + eps).sqrt()
+    )
+    return (d0 + d1) / 2
+
+
+def generalized_epi_dist(
+    kpts0, kpts1, cam0: Camera, cam1: Camera, T_0to1: Pose, all=True, essential=True
+):
+    if essential:
+        E = T_to_E(T_0to1)
+        p0 = cam0.image2cam(kpts0)
+        p1 = cam1.image2cam(kpts1)
+        if all:
+            return sym_epipolar_distance_all(p0, p1, E, agg="max")
+        else:
+            return sym_epipolar_distance(p0, p1, E, squared=False)
+    else:
+        assert cam0._data.shape[-1] == 6
+        assert cam1._data.shape[-1] == 6
+        K0, K1 = cam0.calibration_matrix(), cam1.calibration_matrix()
+        F = K1.inverse().transpose(-1, -2) @ T_to_E(T_0to1) @ K0.inverse()
+        if all:
+            return sym_epipolar_distance_all(kpts0, kpts1, F)
+        else:
+            return sym_epipolar_distance(kpts0, kpts1, F, squared=False)
+
+
+def decompose_essential_matrix(E):
+    # decompose matrix by its singular values
+    U, _, V = torch.svd(E)
+    Vt = V.transpose(-2, -1)
+
+    mask = torch.ones_like(E)
+    mask[..., -1:] *= -1.0  # fill last column with negative values
+
+    maskt = mask.transpose(-2, -1)
+
+    # avoid singularities
+    U = torch.where((torch.det(U) < 0.0)[..., None, None], U * mask, U)
+    Vt = torch.where((torch.det(Vt) < 0.0)[..., None, None], Vt * maskt, Vt)
+
+    W = skew_symmetric(E.new_tensor([[0, 0, 1]]))
+    W[..., 2, 2] += 1.0
+
+    # reconstruct rotations and retrieve translation vector
+    U_W_Vt = U @ W @ Vt
+    U_Wt_Vt = U @ W.transpose(-2, -1) @ Vt
+
+    # return values
+    R1 = U_W_Vt
+    R2 = U_Wt_Vt
+    T = U[..., -1]
+    return R1, R2, T
+
+
+# pose errors
+# TODO: test for batched data
+def angle_error_mat(R1, R2):
+    cos = (torch.trace(torch.einsum("...ij, ...jk -> ...ik", R1.T, R2)) - 1) / 2
+    cos = torch.clip(cos, -1.0, 1.0)  # numerical errors can make it out of bounds
+    return torch.rad2deg(torch.abs(torch.arccos(cos)))
+
+
+def angle_error_vec(v1, v2, eps=1e-10):
+    n = torch.clip(v1.norm(dim=-1) * v2.norm(dim=-1), min=eps)
+    v1v2 = (v1 * v2).sum(dim=-1)  # dot product in the last dimension
+    return torch.rad2deg(torch.arccos(torch.clip(v1v2 / n, -1.0, 1.0)))
+
+
+def relative_pose_error(T_0to1, R, t, ignore_gt_t_thr=0.0, eps=1e-10):
+    if isinstance(T_0to1, torch.Tensor):
+        R_gt, t_gt = T_0to1[:3, :3], T_0to1[:3, 3]
+    else:
+        R_gt, t_gt = T_0to1.R, T_0to1.t
+    R_gt, t_gt = torch.squeeze(R_gt), torch.squeeze(t_gt)
+
+    # angle error between 2 vectors
+    t_err = angle_error_vec(t, t_gt, eps)
+    t_err = torch.minimum(t_err, 180 - t_err)  # handle E ambiguity
+    if t_gt.norm() < ignore_gt_t_thr:  # pure rotation is challenging
+        t_err = 0
+
+    # angle error between 2 rotation matrices
+    r_err = angle_error_mat(R, R_gt)
+
+    return t_err, r_err