src/mast3r_src/mast3r/colmap/database.py

# Copyright (C) 2024-present Naver Corporation. All rights reserved.
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
#
# --------------------------------------------------------
# MASt3R to colmap export functions
# --------------------------------------------------------
import os
import torch
import copy
import numpy as np
import torchvision
import numpy as np
from tqdm import tqdm
from scipy.cluster.hierarchy import DisjointSet
from scipy.spatial.transform import Rotation as R

from mast3r.utils.misc import hash_md5

from mast3r.fast_nn import extract_correspondences_nonsym, bruteforce_reciprocal_nns

import mast3r.utils.path_to_dust3r  # noqa
from dust3r.utils.geometry import find_reciprocal_matches, xy_grid  # noqa


def convert_im_matches_pairs(img0, img1, image_to_colmap, im_keypoints, matches_im0, matches_im1, viz):
    if viz:
        from matplotlib import pyplot as pl

        image_mean = torch.as_tensor(
            [0.5, 0.5, 0.5], device='cpu').reshape(1, 3, 1, 1)
        image_std = torch.as_tensor(
            [0.5, 0.5, 0.5], device='cpu').reshape(1, 3, 1, 1)
        rgb0 = img0['img'] * image_std + image_mean
        rgb0 = torchvision.transforms.functional.to_pil_image(rgb0[0])
        rgb0 = np.array(rgb0)

        rgb1 = img1['img'] * image_std + image_mean
        rgb1 = torchvision.transforms.functional.to_pil_image(rgb1[0])
        rgb1 = np.array(rgb1)

        imgs = [rgb0, rgb1]
        # visualize a few matches
        n_viz = 100
        num_matches = matches_im0.shape[0]
        match_idx_to_viz = np.round(np.linspace(
            0, num_matches - 1, n_viz)).astype(int)
        viz_matches_im0, viz_matches_im1 = matches_im0[match_idx_to_viz], matches_im1[match_idx_to_viz]

        H0, W0, H1, W1 = *imgs[0].shape[:2], *imgs[1].shape[:2]
        rgb0 = np.pad(imgs[0], ((0, max(H1 - H0, 0)),
                                (0, 0), (0, 0)), 'constant', constant_values=0)
        rgb1 = np.pad(imgs[1], ((0, max(H0 - H1, 0)),
                                (0, 0), (0, 0)), 'constant', constant_values=0)
        img = np.concatenate((rgb0, rgb1), axis=1)
        pl.figure()
        pl.imshow(img)
        cmap = pl.get_cmap('jet')
        for ii in range(n_viz):
            (x0, y0), (x1,
                       y1) = viz_matches_im0[ii].T, viz_matches_im1[ii].T
            pl.plot([x0, x1 + W0], [y0, y1], '-+', color=cmap(ii /
                    (n_viz - 1)), scalex=False, scaley=False)
        pl.show(block=True)

    matches = [matches_im0.astype(np.float64), matches_im1.astype(np.float64)]
    imgs = [img0, img1]
    imidx0 = img0['idx']
    imidx1 = img1['idx']
    ravel_matches = []
    for j in range(2):
        H, W = imgs[j]['true_shape'][0]
        with np.errstate(invalid='ignore'):
            qx, qy = matches[j].round().astype(np.int32).T
        ravel_matches_j = qx.clip(min=0, max=W - 1, out=qx) + W * qy.clip(min=0, max=H - 1, out=qy)
        ravel_matches.append(ravel_matches_j)
        imidxj = imgs[j]['idx']
        for m in ravel_matches_j:
            if m not in im_keypoints[imidxj]:
                im_keypoints[imidxj][m] = 0
            im_keypoints[imidxj][m] += 1
    imid0 = copy.deepcopy(image_to_colmap[imidx0]['colmap_imid'])
    imid1 = copy.deepcopy(image_to_colmap[imidx1]['colmap_imid'])
    if imid0 > imid1:
        colmap_matches = np.stack([ravel_matches[1], ravel_matches[0]], axis=-1)
        imid0, imid1 = imid1, imid0
        imidx0, imidx1 = imidx1, imidx0
    else:
        colmap_matches = np.stack([ravel_matches[0], ravel_matches[1]], axis=-1)
    colmap_matches = np.unique(colmap_matches, axis=0)
    return imidx0, imidx1, colmap_matches


def get_im_matches(pred1, pred2, pairs, image_to_colmap, im_keypoints, conf_thr,
                   is_sparse=True, subsample=8, pixel_tol=0, viz=False, device='cuda'):
    im_matches = {}
    for i in range(len(pred1['pts3d'])):
        imidx0 = pairs[i][0]['idx']
        imidx1 = pairs[i][1]['idx']
        if 'desc' in pred1:  # mast3r
            descs = [pred1['desc'][i], pred2['desc'][i]]
            confidences = [pred1['desc_conf'][i], pred2['desc_conf'][i]]
            desc_dim = descs[0].shape[-1]

            if is_sparse:
                corres = extract_correspondences_nonsym(descs[0], descs[1], confidences[0], confidences[1],
                                                        device=device, subsample=subsample, pixel_tol=pixel_tol)
                conf = corres[2]
                mask = conf >= conf_thr
                matches_im0 = corres[0][mask].cpu().numpy()
                matches_im1 = corres[1][mask].cpu().numpy()
            else:
                confidence_masks = [confidences[0] >=
                                    conf_thr, confidences[1] >= conf_thr]
                pts2d_list, desc_list = [], []
                for j in range(2):
                    conf_j = confidence_masks[j].cpu().numpy().flatten()
                    true_shape_j = pairs[i][j]['true_shape'][0]
                    pts2d_j = xy_grid(
                        true_shape_j[1], true_shape_j[0]).reshape(-1, 2)[conf_j]
                    desc_j = descs[j].detach().cpu(
                    ).numpy().reshape(-1, desc_dim)[conf_j]
                    pts2d_list.append(pts2d_j)
                    desc_list.append(desc_j)
                if len(desc_list[0]) == 0 or len(desc_list[1]) == 0:
                    continue

                nn0, nn1 = bruteforce_reciprocal_nns(desc_list[0], desc_list[1],
                                                     device=device, dist='dot', block_size=2**13)
                reciprocal_in_P0 = (nn1[nn0] == np.arange(len(nn0)))

                matches_im1 = pts2d_list[1][nn0][reciprocal_in_P0]
                matches_im0 = pts2d_list[0][reciprocal_in_P0]
        else:
            pts3d = [pred1['pts3d'][i], pred2['pts3d_in_other_view'][i]]
            confidences = [pred1['conf'][i], pred2['conf'][i]]

            if is_sparse:
                corres = extract_correspondences_nonsym(pts3d[0], pts3d[1], confidences[0], confidences[1],
                                                        device=device, subsample=subsample, pixel_tol=pixel_tol,
                                                        ptmap_key='3d')
                conf = corres[2]
                mask = conf >= conf_thr
                matches_im0 = corres[0][mask].cpu().numpy()
                matches_im1 = corres[1][mask].cpu().numpy()
            else:
                confidence_masks = [confidences[0] >=
                                    conf_thr, confidences[1] >= conf_thr]
                # find 2D-2D matches between the two images
                pts2d_list, pts3d_list = [], []
                for j in range(2):
                    conf_j = confidence_masks[j].cpu().numpy().flatten()
                    true_shape_j = pairs[i][j]['true_shape'][0]
                    pts2d_j = xy_grid(true_shape_j[1], true_shape_j[0]).reshape(-1, 2)[conf_j]
                    pts3d_j = pts3d[j].detach().cpu().numpy().reshape(-1, 3)[conf_j]
                    pts2d_list.append(pts2d_j)
                    pts3d_list.append(pts3d_j)

                PQ, PM = pts3d_list[0], pts3d_list[1]
                if len(PQ) == 0 or len(PM) == 0:
                    continue
                reciprocal_in_PM, nnM_in_PQ, num_matches = find_reciprocal_matches(
                    PQ, PM)

                matches_im1 = pts2d_list[1][reciprocal_in_PM]
                matches_im0 = pts2d_list[0][nnM_in_PQ][reciprocal_in_PM]

        if len(matches_im0) == 0:
            continue
        imidx0, imidx1, colmap_matches = convert_im_matches_pairs(pairs[i][0], pairs[i][1],
                                                                  image_to_colmap, im_keypoints,
                                                                  matches_im0, matches_im1, viz)
        im_matches[(imidx0, imidx1)] = colmap_matches
    return im_matches


def get_im_matches_from_cache(pairs, cache_path, desc_conf, subsample,
                              image_to_colmap, im_keypoints, conf_thr,
                              viz=False, device='cuda'):
    im_matches = {}
    for i in range(len(pairs)):
        imidx0 = pairs[i][0]['idx']
        imidx1 = pairs[i][1]['idx']

        corres_idx1 = hash_md5(pairs[i][0]['instance'])
        corres_idx2 = hash_md5(pairs[i][1]['instance'])

        path_corres = cache_path + f'/corres_conf={desc_conf}_{subsample=}/{corres_idx1}-{corres_idx2}.pth'
        if os.path.isfile(path_corres):
            score, (xy1, xy2, confs) = torch.load(path_corres, map_location=device)
        else:
            path_corres = cache_path + f'/corres_conf={desc_conf}_{subsample=}/{corres_idx2}-{corres_idx1}.pth'
            score, (xy2, xy1, confs) = torch.load(path_corres, map_location=device)
        mask = confs >= conf_thr
        matches_im0 = xy1[mask].cpu().numpy()
        matches_im1 = xy2[mask].cpu().numpy()

        if len(matches_im0) == 0:
            continue
        imidx0, imidx1, colmap_matches = convert_im_matches_pairs(pairs[i][0], pairs[i][1],
                                                                  image_to_colmap, im_keypoints,
                                                                  matches_im0, matches_im1, viz)
        im_matches[(imidx0, imidx1)] = colmap_matches
    return im_matches


def export_images(db, images, image_paths, focals, ga_world_to_cam, camera_model):
    # add cameras/images to the db
    # with the output of ga as prior
    image_to_colmap = {}
    im_keypoints = {}
    for idx in range(len(image_paths)):
        im_keypoints[idx] = {}
        H, W = images[idx]["orig_shape"]
        if focals is None:
            focal_x = focal_y = 1.2 * max(W, H)
            prior_focal_length = False
            cx = W / 2.0
            cy = H / 2.0
        elif isinstance(focals[idx], np.ndarray) and len(focals[idx].shape) == 2:
            # intrinsics
            focal_x = focals[idx][0, 0]
            focal_y = focals[idx][1, 1]
            cx = focals[idx][0, 2] * images[idx]["to_orig"][0, 0]
            cy = focals[idx][1, 2] * images[idx]["to_orig"][1, 1]
            prior_focal_length = True
        else:
            focal_x = focal_y = float(focals[idx])
            prior_focal_length = True
            cx = W / 2.0
            cy = H / 2.0
        focal_x = focal_x * images[idx]["to_orig"][0, 0]
        focal_y = focal_y * images[idx]["to_orig"][1, 1]

        if camera_model == "SIMPLE_PINHOLE":
            model_id = 0
            focal = (focal_x + focal_y) / 2.0
            params = np.asarray([focal, cx, cy], np.float64)
        elif camera_model == "PINHOLE":
            model_id = 1
            params = np.asarray([focal_x, focal_y, cx, cy], np.float64)
        elif camera_model == "SIMPLE_RADIAL":
            model_id = 2
            focal = (focal_x + focal_y) / 2.0
            params = np.asarray([focal, cx, cy, 0.0], np.float64)
        elif camera_model == "OPENCV":
            model_id = 4
            params = np.asarray([focal_x, focal_y, cx, cy, 0.0, 0.0, 0.0, 0.0], np.float64)
        else:
            raise ValueError(f"invalid camera model {camera_model}")

        H, W = int(H), int(W)
        # OPENCV camera model
        camid = db.add_camera(
            model_id, W, H, params, prior_focal_length=prior_focal_length)
        if ga_world_to_cam is None:
            prior_t = np.zeros(3)
            prior_q = np.zeros(4)
        else:
            q = R.from_matrix(ga_world_to_cam[idx][:3, :3]).as_quat()
            prior_t = ga_world_to_cam[idx][:3, 3]
            prior_q = np.array([q[-1], q[0], q[1], q[2]])
        imid = db.add_image(
            image_paths[idx], camid, prior_q=prior_q, prior_t=prior_t)
        image_to_colmap[idx] = {
            'colmap_imid': imid,
            'colmap_camid': camid
        }
    return image_to_colmap, im_keypoints


def export_matches(db, images, image_to_colmap, im_keypoints, im_matches, min_len_track, skip_geometric_verification):
    colmap_image_pairs = []
    # 2D-2D are quite dense
    # we want to remove the very small tracks
    # and export only kpt for which we have values
    # build tracks
    print("building tracks")
    keypoints_to_track_id = {}
    track_id_to_kpt_list = []
    to_merge = []
    for (imidx0, imidx1), colmap_matches in tqdm(im_matches.items()):
        if imidx0 not in keypoints_to_track_id:
            keypoints_to_track_id[imidx0] = {}
        if imidx1 not in keypoints_to_track_id:
            keypoints_to_track_id[imidx1] = {}

        for m in colmap_matches:
            if m[0] not in keypoints_to_track_id[imidx0] and m[1] not in keypoints_to_track_id[imidx1]:
                # new pair of kpts never seen before
                track_idx = len(track_id_to_kpt_list)
                keypoints_to_track_id[imidx0][m[0]] = track_idx
                keypoints_to_track_id[imidx1][m[1]] = track_idx
                track_id_to_kpt_list.append(
                    [(imidx0, m[0]), (imidx1, m[1])])
            elif m[1] not in keypoints_to_track_id[imidx1]:
                # 0 has a track, not 1
                track_idx = keypoints_to_track_id[imidx0][m[0]]
                keypoints_to_track_id[imidx1][m[1]] = track_idx
                track_id_to_kpt_list[track_idx].append((imidx1, m[1]))
            elif m[0] not in keypoints_to_track_id[imidx0]:
                # 1 has a track, not 0
                track_idx = keypoints_to_track_id[imidx1][m[1]]
                keypoints_to_track_id[imidx0][m[0]] = track_idx
                track_id_to_kpt_list[track_idx].append((imidx0, m[0]))
            else:
                # both have tracks, merge them
                track_idx0 = keypoints_to_track_id[imidx0][m[0]]
                track_idx1 = keypoints_to_track_id[imidx1][m[1]]
                if track_idx0 != track_idx1:
                    # let's deal with them later
                    to_merge.append((track_idx0, track_idx1))

    # regroup merge targets
    print("merging tracks")
    unique = np.unique(to_merge)
    tree = DisjointSet(unique)
    for track_idx0, track_idx1 in tqdm(to_merge):
        tree.merge(track_idx0, track_idx1)

    subsets = tree.subsets()
    print("applying merge")
    for setvals in tqdm(subsets):
        new_trackid = len(track_id_to_kpt_list)
        kpt_list = []
        for track_idx in setvals:
            kpt_list.extend(track_id_to_kpt_list[track_idx])
            for imidx, kpid in track_id_to_kpt_list[track_idx]:
                keypoints_to_track_id[imidx][kpid] = new_trackid
        track_id_to_kpt_list.append(kpt_list)

    # binc = np.bincount([len(v) for v in track_id_to_kpt_list])
    # nonzero = np.nonzero(binc)
    # nonzerobinc = binc[nonzero[0]]
    # print(nonzero[0].tolist())
    # print(nonzerobinc)
    num_valid_tracks = sum(
        [1 for v in track_id_to_kpt_list if len(v) >= min_len_track])

    keypoints_to_idx = {}
    print(f"squashing keypoints - {num_valid_tracks} valid tracks")
    for imidx, keypoints_imid in tqdm(im_keypoints.items()):
        imid = image_to_colmap[imidx]['colmap_imid']
        keypoints_kept = []
        keypoints_to_idx[imidx] = {}
        for kp in keypoints_imid.keys():
            if kp not in keypoints_to_track_id[imidx]:
                continue
            track_idx = keypoints_to_track_id[imidx][kp]
            track_length = len(track_id_to_kpt_list[track_idx])
            if track_length < min_len_track:
                continue
            keypoints_to_idx[imidx][kp] = len(keypoints_kept)
            keypoints_kept.append(kp)
        if len(keypoints_kept) == 0:
            continue
        keypoints_kept = np.array(keypoints_kept)
        keypoints_kept = np.unravel_index(keypoints_kept, images[imidx]['true_shape'][0])[
            0].base[:, ::-1].copy().astype(np.float32)
        # rescale coordinates
        keypoints_kept[:, 0] += 0.5
        keypoints_kept[:, 1] += 0.5
        keypoints_kept = geotrf(images[imidx]['to_orig'], keypoints_kept, norm=True)

        H, W = images[imidx]['orig_shape']
        keypoints_kept[:, 0] = keypoints_kept[:, 0].clip(min=0, max=W - 0.01)
        keypoints_kept[:, 1] = keypoints_kept[:, 1].clip(min=0, max=H - 0.01)

        db.add_keypoints(imid, keypoints_kept)

    print("exporting im_matches")
    for (imidx0, imidx1), colmap_matches in im_matches.items():
        imid0, imid1 = image_to_colmap[imidx0]['colmap_imid'], image_to_colmap[imidx1]['colmap_imid']
        assert imid0 < imid1
        final_matches = np.array([[keypoints_to_idx[imidx0][m[0]], keypoints_to_idx[imidx1][m[1]]]
                                  for m in colmap_matches
                                  if m[0] in keypoints_to_idx[imidx0] and m[1] in keypoints_to_idx[imidx1]])
        if len(final_matches) > 0:
            colmap_image_pairs.append(
                (images[imidx0]['instance'], images[imidx1]['instance']))
            db.add_matches(imid0, imid1, final_matches)
            if skip_geometric_verification:
                db.add_two_view_geometry(imid0, imid1, final_matches)
    return colmap_image_pairs