PCGao/MatchAnything

1import pickle2import random3 4import numpy as np5import pycolmap6from matplotlib import cm7 8from .utils.io import read_image9from .utils.viz import (10    add_text,11    cm_RdGn,12    plot_images,13    plot_keypoints,14    plot_matches,15)16 17 18def visualize_sfm_2d(19    reconstruction,20    image_dir,21    color_by="visibility",22    selected=[],23    n=1,24    seed=0,25    dpi=75,26):27    assert image_dir.exists()28    if not isinstance(reconstruction, pycolmap.Reconstruction):29        reconstruction = pycolmap.Reconstruction(reconstruction)30 31    if not selected:32        image_ids = reconstruction.reg_image_ids()33        selected = random.Random(seed).sample(image_ids, min(n, len(image_ids)))34 35    for i in selected:36        image = reconstruction.images[i]37        keypoints = np.array([p.xy for p in image.points2D])38        visible = np.array([p.has_point3D() for p in image.points2D])39 40        if color_by == "visibility":41            color = [(0, 0, 1) if v else (1, 0, 0) for v in visible]42            text = f"visible: {np.count_nonzero(visible)}/{len(visible)}"43        elif color_by == "track_length":44            tl = np.array(45                [46                    (47                        reconstruction.points3D[p.point3D_id].track.length()48                        if p.has_point3D()49                        else 150                    )51                    for p in image.points2D52                ]53            )54            max_, med_ = np.max(tl), np.median(tl[tl > 1])55            tl = np.log(tl)56            color = cm.jet(tl / tl.max()).tolist()57            text = f"max/median track length: {max_}/{med_}"58        elif color_by == "depth":59            p3ids = [p.point3D_id for p in image.points2D if p.has_point3D()]60            z = np.array(61                [62                    (image.cam_from_world * reconstruction.points3D[j].xyz)[-1]63                    for j in p3ids64                ]65            )66            z -= z.min()67            color = cm.jet(z / np.percentile(z, 99.9))68            text = f"visible: {np.count_nonzero(visible)}/{len(visible)}"69            keypoints = keypoints[visible]70        else:71            raise NotImplementedError(f"Coloring not implemented: {color_by}.")72 73        name = image.name74        fig = plot_images([read_image(image_dir / name)], dpi=dpi)75        plot_keypoints([keypoints], colors=[color], ps=4)76        add_text(0, text)77        add_text(0, name, pos=(0.01, 0.01), fs=5, lcolor=None, va="bottom")78    return fig79 80 81def visualize_loc(82    results,83    image_dir,84    reconstruction=None,85    db_image_dir=None,86    selected=[],87    n=1,88    seed=0,89    prefix=None,90    **kwargs,91):92    assert image_dir.exists()93 94    with open(str(results) + "_logs.pkl", "rb") as f:95        logs = pickle.load(f)96 97    if not selected:98        queries = list(logs["loc"].keys())99        if prefix:100            queries = [q for q in queries if q.startswith(prefix)]101        selected = random.Random(seed).sample(queries, min(n, len(queries)))102 103    if reconstruction is not None:104        if not isinstance(reconstruction, pycolmap.Reconstruction):105            reconstruction = pycolmap.Reconstruction(reconstruction)106 107    for qname in selected:108        loc = logs["loc"][qname]109        visualize_loc_from_log(110            image_dir, qname, loc, reconstruction, db_image_dir, **kwargs111        )112 113 114def visualize_loc_from_log(115    image_dir,116    query_name,117    loc,118    reconstruction=None,119    db_image_dir=None,120    top_k_db=2,121    dpi=75,122):123    q_image = read_image(image_dir / query_name)124    if loc.get("covisibility_clustering", False):125        # select the first, largest cluster if the localization failed126        loc = loc["log_clusters"][loc["best_cluster"] or 0]127 128    inliers = np.array(loc["PnP_ret"]["inliers"])129    mkp_q = loc["keypoints_query"]130    n = len(loc["db"])131    if reconstruction is not None:132        # for each pair of query keypoint and its matched 3D point,133        # we need to find its corresponding keypoint in each database image134        # that observes it. We also count the number of inliers in each.135        kp_idxs, kp_to_3D_to_db = loc["keypoint_index_to_db"]136        counts = np.zeros(n)137        dbs_kp_q_db = [[] for _ in range(n)]138        inliers_dbs = [[] for _ in range(n)]139        for i, (inl, (p3D_id, db_idxs)) in enumerate(zip(inliers, kp_to_3D_to_db)):140            track = reconstruction.points3D[p3D_id].track141            track = {el.image_id: el.point2D_idx for el in track.elements}142            for db_idx in db_idxs:143                counts[db_idx] += inl144                kp_db = track[loc["db"][db_idx]]145                dbs_kp_q_db[db_idx].append((i, kp_db))146                inliers_dbs[db_idx].append(inl)147    else:148        # for inloc the database keypoints are already in the logs149        assert "keypoints_db" in loc150        assert "indices_db" in loc151        counts = np.array([np.sum(loc["indices_db"][inliers] == i) for i in range(n)])152 153    # display the database images with the most inlier matches154    db_sort = np.argsort(-counts)155    for db_idx in db_sort[:top_k_db]:156        if reconstruction is not None:157            db = reconstruction.images[loc["db"][db_idx]]158            db_name = db.name159            db_kp_q_db = np.array(dbs_kp_q_db[db_idx])160            kp_q = mkp_q[db_kp_q_db[:, 0]]161            kp_db = np.array([db.points2D[i].xy for i in db_kp_q_db[:, 1]])162            inliers_db = inliers_dbs[db_idx]163        else:164            db_name = loc["db"][db_idx]165            kp_q = mkp_q[loc["indices_db"] == db_idx]166            kp_db = loc["keypoints_db"][loc["indices_db"] == db_idx]167            inliers_db = inliers[loc["indices_db"] == db_idx]168 169        db_image = read_image((db_image_dir or image_dir) / db_name)170        color = cm_RdGn(inliers_db).tolist()171        text = f"inliers: {sum(inliers_db)}/{len(inliers_db)}"172 173        plot_images([q_image, db_image], dpi=dpi)174        plot_matches(kp_q, kp_db, color, a=0.1)175        add_text(0, text)176        opts = dict(pos=(0.01, 0.01), fs=5, lcolor=None, va="bottom")177        add_text(0, query_name, **opts)178        add_text(1, db_name, **opts)179