third_party/pram/recognition/recmap.py

# -*- coding: UTF-8 -*-
'''=================================================
@Project -> File   pram -> recmap
@IDE    PyCharm
@Author fx221@cam.ac.uk
@Date   07/02/2024 11:02
=================================================='''
import argparse
import torch
import os
import os.path as osp
import numpy as np
import cv2
import yaml
import multiprocessing as mp
from copy import deepcopy
import logging
import h5py
from tqdm import tqdm
import open3d as o3d
from sklearn.cluster import KMeans, Birch
from collections import defaultdict
from colmap_utils.read_write_model import read_model, qvec2rotmat, write_cameras_binary, write_images_binary
from colmap_utils.read_write_model import write_points3d_binary, Image, Point3D, Camera
from colmap_utils.read_write_model import write_compressed_points3d_binary, write_compressed_images_binary
from recognition.vis_seg import generate_color_dic, vis_seg_point, plot_kpts


class RecMap:
    def __init__(self):
        self.cameras = None
        self.images = None
        self.points3D = None
        self.pcd = o3d.geometry.PointCloud()
        self.seg_color_dict = generate_color_dic(n_seg=1000)

    def load_sfm_model(self, path: str, ext='.bin'):
        self.cameras, self.images, self.points3D = read_model(path, ext)
        self.name_to_id = {image.name: i for i, image in self.images.items()}
        print('Load {:d} cameras, {:d} images, {:d} points'.format(len(self.cameras), len(self.images),
                                                                   len(self.points3D)))

    def remove_statics_outlier(self, nb_neighbors: int = 20, std_ratio: float = 2.0):
        xyzs = []
        p3d_ids = []
        for p3d_id in self.points3D.keys():
            xyzs.append(self.points3D[p3d_id].xyz)
            p3d_ids.append(p3d_id)

        xyzs = np.array(xyzs)
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(xyzs)
        new_pcd, inlier_ids = pcd.remove_statistical_outlier(nb_neighbors=nb_neighbors, std_ratio=std_ratio)

        new_point3Ds = {}
        for i in inlier_ids:
            new_point3Ds[p3d_ids[i]] = self.points3D[p3d_ids[i]]
        self.points3D = new_point3Ds
        n_outlier = xyzs.shape[0] - len(inlier_ids)
        ratio = n_outlier / xyzs.shape[0]
        print('Remove {:d} - {:d} = {:d}/{:.2f}% points'.format(xyzs.shape[0], len(inlier_ids), n_outlier, ratio * 100))

    def load_segmentation(self, path: str):
        data = np.load(path, allow_pickle=True)[()]
        p3d_id = data['id']
        seg_id = data['label']
        self.p3d_seg = {p3d_id[i]: seg_id[i] for i in range(p3d_id.shape[0])}
        self.seg_p3d = {}
        for pid in self.p3d_seg.keys():
            sid = self.p3d_seg[pid]
            if sid not in self.seg_p3d.keys():
                self.seg_p3d[sid] = [pid]
            else:
                self.seg_p3d[sid].append(pid)

        if 'xyz' not in data.keys():
            all_xyz = []
            for pid in p3d_id:
                xyz = self.points3D[pid].xyz
                all_xyz.append(xyz)
            data['xyz'] = np.array(all_xyz)
            np.save(path, data)
            print('Add xyz to ', path)

    def cluster(self, k=512, mode='xyz', min_obs=3, save_fn=None, method='kmeans', **kwargs):
        if save_fn is not None:
            if osp.isfile(save_fn):
                print('{:s} exists.'.format(save_fn))
                return
        all_xyz = []
        point3D_ids = []
        for p3d in self.points3D.values():
            track_len = len(p3d.point2D_idxs)
            if track_len < min_obs:
                continue
            all_xyz.append(p3d.xyz)
            point3D_ids.append(p3d.id)

        xyz = np.array(all_xyz)
        point3D_ids = np.array(point3D_ids)

        if mode.find('x') < 0:
            xyz[:, 0] = 0
        if mode.find('y') < 0:
            xyz[:, 1] = 0
        if mode.find('z') < 0:
            xyz[:, 2] = 0

        if method == 'kmeans':
            model = KMeans(n_clusters=k, random_state=0, verbose=True).fit(xyz)
        elif method == 'birch':
            model = Birch(threshold=kwargs.get('threshold'), n_clusters=k).fit(xyz)  # 0.01 for indoor
        else:
            print('Method {:s} for clustering does not exist'.format(method))
            exit(0)
        labels = np.array(model.labels_).reshape(-1)
        if save_fn is not None:
            np.save(save_fn, {
                'id': np.array(point3D_ids),  # should be assigned to self.points3D_ids
                'label': np.array(labels),
                'xyz': np.array(all_xyz),
            })

    def assign_point3D_descriptor(self, feature_fn: str, save_fn=None, n_process=1):
        '''
        assign each 3d point a descriptor for localization
        :param feature_fn: file name of features [h5py]
        :param save_fn:
        :param n_process:
        :return:
        '''

        def run(start_id, end_id, points3D_desc):
            for pi in tqdm(range(start_id, end_id), total=end_id - start_id):
                p3d_id = all_p3d_ids[pi]
                img_list = self.points3D[p3d_id].image_ids
                kpt_ids = self.points3D[p3d_id].point2D_idxs
                all_descs = []
                for img_id, p2d_id in zip(img_list, kpt_ids):
                    if img_id not in self.images.keys():
                        continue
                    img_fn = self.images[img_id].name
                    desc = feat_file[img_fn]['descriptors'][()].transpose()[p2d_id]
                    all_descs.append(desc)

                if len(all_descs) == 1:
                    points3D_desc[p3d_id] = all_descs[0]
                else:
                    all_descs = np.array(all_descs)  # [n, d]
                    dist = all_descs @ all_descs.transpose()  # [n, n]
                    dist = 2 - 2 * dist
                    md_dist = np.median(dist, axis=-1)  # [n]
                    min_id = np.argmin(md_dist)
                    points3D_desc[p3d_id] = all_descs[min_id]

        if osp.isfile(save_fn):
            print('{:s} exists.'.format(save_fn))
            return
        p3D_desc = {}
        feat_file = h5py.File(feature_fn, 'r')
        all_p3d_ids = sorted(self.points3D.keys())

        if n_process > 1:
            if len(all_p3d_ids) <= n_process:
                run(start_id=0, end_id=len(all_p3d_ids), points3D_desc=p3D_desc)
            else:
                manager = mp.Manager()
                output = manager.dict()  # necessary otherwise empty
                n_sample_per_process = len(all_p3d_ids) // n_process
                jobs = []
                for i in range(n_process):
                    start_id = i * n_sample_per_process
                    if i == n_process - 1:
                        end_id = len(all_p3d_ids)
                    else:
                        end_id = (i + 1) * n_sample_per_process
                    p = mp.Process(
                        target=run,
                        args=(start_id, end_id, output),
                    )
                    jobs.append(p)
                    p.start()

                for p in jobs:
                    p.join()

                p3D_desc = {}
                for k in output.keys():
                    p3D_desc[k] = output[k]
        else:
            run(start_id=0, end_id=len(all_p3d_ids), points3D_desc=p3D_desc)

        if save_fn is not None:
            np.save(save_fn, p3D_desc)

    def reproject(self, img_id, xyzs):
        qvec = self.images[img_id].qvec
        Rcw = qvec2rotmat(qvec=qvec)
        tvec = self.images[img_id].tvec
        tcw = tvec.reshape(3, )
        Tcw = np.eye(4, dtype=float)
        Tcw[:3, :3] = Rcw
        Tcw[:3, 3] = tcw
        # intrinsics
        cam = self.cameras[self.images[img_id].camera_id]
        K = self.get_intrinsics_from_camera(camera=cam)

        xyzs_homo = np.hstack([xyzs, np.ones(shape=(xyzs.shape[0], 1), dtype=float)])
        kpts = K @ ((Tcw @ xyzs_homo.transpose())[:3, :])  # [3, N]
        kpts = kpts.transpose()  # [N, 3]
        kpts[:, 0] = kpts[:, 0] / kpts[:, 2]
        kpts[:, 1] = kpts[:, 1] / kpts[:, 2]

        return kpts

    def find_covisible_frame_ids(self, image_id, images, points3D):
        covis = defaultdict(int)
        p3d_ids = images[image_id].point3D_ids

        for pid in p3d_ids:
            if pid == -1:
                continue
            if pid not in points3D.keys():
                continue
            for im in points3D[pid].image_ids:
                covis[im] += 1

        covis_ids = np.array(list(covis.keys()))
        covis_num = np.array([covis[i] for i in covis_ids])
        ind_top = np.argsort(covis_num)[::-1]
        sorted_covis_ids = [covis_ids[i] for i in ind_top]
        return sorted_covis_ids

    def create_virtual_frame_3(self, save_fn=None, save_vrf_dir=None, show_time=-1, ignored_cameras=[],
                               min_cover_ratio=0.9,
                               depth_scale=1.2,
                               radius=15,
                               min_obs=120,
                               topk_imgs=500,
                               n_vrf=10,
                               covisible_frame=20,
                               **kwargs):
        def reproject(img_id, xyzs):
            qvec = self.images[img_id].qvec
            Rcw = qvec2rotmat(qvec=qvec)
            tvec = self.images[img_id].tvec
            tcw = tvec.reshape(3, )
            Tcw = np.eye(4, dtype=float)
            Tcw[:3, :3] = Rcw
            Tcw[:3, 3] = tcw
            # intrinsics
            cam = self.cameras[self.images[img_id].camera_id]
            K = self.get_intrinsics_from_camera(camera=cam)

            xyzs_homo = np.hstack([xyzs, np.ones(shape=(xyzs.shape[0], 1), dtype=float)])
            kpts = K @ ((Tcw @ xyzs_homo.transpose())[:3, :])  # [3, N]
            kpts = kpts.transpose()  # [N, 3]
            kpts[:, 0] = kpts[:, 0] / kpts[:, 2]
            kpts[:, 1] = kpts[:, 1] / kpts[:, 2]

            return kpts

        def find_best_vrf_by_covisibility(p3d_id_list):
            all_img_ids = []
            all_xyzs = []

            img_ids_full = []
            img_id_obs = {}
            for pid in p3d_id_list:
                if pid not in self.points3D.keys():
                    continue
                all_xyzs.append(self.points3D[pid].xyz)

                img_ids = self.points3D[pid].image_ids
                for iid in img_ids:
                    if iid in all_img_ids:
                        continue
                    # valid_p3ds = [v for v in self.images[iid].point3D_ids if v > 0 and v in p3d_id_list]
                    if len(ignored_cameras) > 0:
                        ignore = False
                        img_name = self.images[iid].name
                        for c in ignored_cameras:
                            if img_name.find(c) >= 0:
                                ignore = True
                                break
                        if ignore:
                            continue
                    # valid_p3ds = np.intersect1d(np.array(self.images[iid].point3D_ids), np.array(p3d_id_list)).tolist()
                    valid_p3ds = [v for v in self.images[iid].point3D_ids if v > 0]
                    img_ids_full.append(iid)
                    if len(valid_p3ds) < min_obs:
                        continue

                    all_img_ids.append(iid)
                    img_id_obs[iid] = len(valid_p3ds)
            all_xyzs = np.array(all_xyzs)

            print('Find {} 3D points and {} images'.format(len(p3d_id_list), len(img_id_obs.keys())))
            top_img_ids_by_obs = sorted(img_id_obs.items(), key=lambda item: item[1], reverse=True)  # [(key, value), ]
            all_img_ids = []
            for item in top_img_ids_by_obs:
                all_img_ids.append(item[0])
                if len(all_img_ids) >= topk_imgs:
                    break

            # all_img_ids = all_img_ids[:200]
            if len(all_img_ids) == 0:
                print('no valid img ids with obs over {:d}'.format(min_obs))
                all_img_ids = img_ids_full

            img_observations = {}
            p3d_id_array = np.array(p3d_id_list)
            for idx, img_id in enumerate(all_img_ids):
                valid_p3ds = [v for v in self.images[img_id].point3D_ids if v > 0]
                mask = np.array([False for i in range(len(p3d_id_list))])
                for pid in valid_p3ds:
                    found_idx = np.where(p3d_id_array == pid)[0]
                    if found_idx.shape[0] == 0:
                        continue
                    mask[found_idx[0]] = True

                img_observations[img_id] = mask

            unobserved_p3d_ids = np.array([True for i in range(len(p3d_id_list))])

            candidate_img_ids = []
            total_cover_ratio = 0
            while total_cover_ratio < min_cover_ratio:
                best_img_id = -1
                best_img_obs = -1
                for idx, im_id in enumerate(all_img_ids):
                    if im_id in candidate_img_ids:
                        continue
                    obs_i = np.sum(img_observations[im_id] * unobserved_p3d_ids)
                    if obs_i > best_img_obs:
                        best_img_id = im_id
                        best_img_obs = obs_i

                if best_img_id >= 0:
                    # keep the valid img_id
                    candidate_img_ids.append(best_img_id)
                    # update the unobserved mask
                    unobserved_p3d_ids[img_observations[best_img_id]] = False
                    total_cover_ratio = 1 - np.sum(unobserved_p3d_ids) / len(p3d_id_list)
                    print(len(candidate_img_ids), best_img_obs, best_img_obs / len(p3d_id_list), total_cover_ratio)

                    if best_img_obs / len(p3d_id_list) < 0.01:
                        break

                    if len(candidate_img_ids) >= n_vrf:
                        break
                else:
                    break

            return candidate_img_ids
            # return [(v, img_observations[v]) for v in candidate_img_ids]

        if save_vrf_dir is not None:
            os.makedirs(save_vrf_dir, exist_ok=True)

        seg_ref = {}
        for sid in self.seg_p3d.keys():
            if sid == -1:  # ignore invalid segment
                continue
            all_p3d_ids = self.seg_p3d[sid]
            candidate_img_ids = find_best_vrf_by_covisibility(p3d_id_list=all_p3d_ids)

            seg_ref[sid] = {}
            for can_idx, img_id in enumerate(candidate_img_ids):
                cam = self.cameras[self.images[img_id].camera_id]
                width = cam.width
                height = cam.height
                qvec = self.images[img_id].qvec
                tvec = self.images[img_id].tvec

                img_name = self.images[img_id].name
                orig_p3d_ids = [p for p in self.images[img_id].point3D_ids if p in self.points3D.keys() and p >= 0]
                orig_xyzs = []
                new_xyzs = []
                for pid in all_p3d_ids:
                    if pid in orig_p3d_ids:
                        orig_xyzs.append(self.points3D[pid].xyz)
                    else:
                        if pid in self.points3D.keys():
                            new_xyzs.append(self.points3D[pid].xyz)

                if len(orig_xyzs) == 0:
                    continue

                orig_xyzs = np.array(orig_xyzs)
                new_xyzs = np.array(new_xyzs)

                print('img: ', osp.join(kwargs.get('image_root'), img_name))
                img = cv2.imread(osp.join(kwargs.get('image_root'), img_name))
                orig_kpts = reproject(img_id=img_id, xyzs=orig_xyzs)
                max_depth = depth_scale * np.max(orig_kpts[:, 2])
                orig_kpts = orig_kpts[:, :2]
                mask_ori = (orig_kpts[:, 0] >= 0) & (orig_kpts[:, 0] < width) & (orig_kpts[:, 1] >= 0) & (
                        orig_kpts[:, 1] < height)
                orig_kpts = orig_kpts[mask_ori]

                if orig_kpts.shape[0] == 0:
                    continue

                img_kpt = plot_kpts(img=img, kpts=orig_kpts, radius=[3 for i in range(orig_kpts.shape[0])],
                                    colors=[(0, 0, 255) for i in range(orig_kpts.shape[0])], thickness=-1)
                if new_xyzs.shape[0] == 0:
                    img_all = img_kpt
                else:
                    new_kpts = reproject(img_id=img_id, xyzs=new_xyzs)
                    mask_depth = (new_kpts[:, 2] > 0) & (new_kpts[:, 2] <= max_depth)
                    mask_in_img = (new_kpts[:, 0] >= 0) & (new_kpts[:, 0] < width) & (new_kpts[:, 1] >= 0) & (
                            new_kpts[:, 1] < height)
                    dist_all_orig = torch.from_numpy(new_kpts[:, :2])[..., None] - \
                                    torch.from_numpy(orig_kpts[:, :2].transpose())[None]
                    dist_all_orig = torch.sqrt(torch.sum(dist_all_orig ** 2, dim=1))  # [N, M]
                    min_dist = torch.min(dist_all_orig, dim=1)[0].numpy()
                    mask_close_to_img = (min_dist <= radius)

                    mask_new = (mask_depth & mask_in_img & mask_close_to_img)

                    cover_ratio = np.sum(mask_ori) + np.sum(mask_new)
                    cover_ratio = cover_ratio / len(all_p3d_ids)

                    print('idx: {:d}, img: ori {:d}/{:d}/{:.2f}, new {:d}/{:d}'.format(can_idx,
                                                                                       orig_kpts.shape[0],
                                                                                       np.sum(mask_ori),
                                                                                       cover_ratio * 100,
                                                                                       new_kpts.shape[0],
                                                                                       np.sum(mask_new)))

                    new_kpts = new_kpts[mask_new]

                    # img_all = img_kpt
                    img_all = plot_kpts(img=img_kpt, kpts=new_kpts, radius=[3 for i in range(new_kpts.shape[0])],
                                        colors=[(0, 255, 0) for i in range(new_kpts.shape[0])], thickness=-1)

                cv2.namedWindow('img', cv2.WINDOW_NORMAL)
                cv2.imshow('img', img_all)

                if save_vrf_dir is not None:
                    cv2.imwrite(osp.join(save_vrf_dir,
                                         'seg-{:05d}_can-{:05d}_'.format(sid, can_idx) + img_name.replace('/', '+')),
                                img_all)

                key = cv2.waitKey(show_time)
                if key == ord('q'):
                    cv2.destroyAllWindows()
                    exit(0)

                covisile_frame_ids = self.find_covisible_frame_ids(image_id=img_id, images=self.images,
                                                                   points3D=self.points3D)
                seg_ref[sid][can_idx] = {
                    'image_name': img_name,
                    'image_id': img_id,
                    'qvec': deepcopy(qvec),
                    'tvec': deepcopy(tvec),
                    'camera': {
                        'model': cam.model,
                        'params': cam.params,
                        'width': cam.width,
                        'height': cam.height,
                    },
                    'original_points3d': np.array(
                        [v for v in self.images[img_id].point3D_ids if v >= 0 and v in self.points3D.keys()]),
                    'covisible_frame_ids': np.array(covisile_frame_ids[:covisible_frame]),
                }
        # save vrf info
        if save_fn is not None:
            print('Save {} segments with virtual reference image information to {}'.format(len(seg_ref.keys()),
                                                                                           save_fn))
            np.save(save_fn, seg_ref)

    def visualize_3Dpoints(self):
        xyz = []
        rgb = []
        for point3D in self.points3D.values():
            xyz.append(point3D.xyz)
            rgb.append(point3D.rgb / 255)

        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(xyz)
        pcd.colors = o3d.utility.Vector3dVector(rgb)
        o3d.visualization.draw_geometries([pcd])

    def visualize_segmentation(self, p3d_segs, points3D):
        p3d_ids = p3d_segs.keys()
        xyzs = []
        rgbs = []
        for pid in p3d_ids:
            xyzs.append(points3D[pid].xyz)
            seg_color = self.seg_color_dict[p3d_segs[pid]]
            rgbs.append(np.array([seg_color[2], seg_color[1], seg_color[0]]) / 255)
        xyzs = np.array(xyzs)
        rgbs = np.array(rgbs)

        self.pcd.points = o3d.utility.Vector3dVector(xyzs)
        self.pcd.colors = o3d.utility.Vector3dVector(rgbs)

        o3d.visualization.draw_geometries([self.pcd])

    def visualize_segmentation_on_image(self, p3d_segs, image_path, feat_path):
        vis_color = generate_color_dic(n_seg=1024)
        feat_file = h5py.File(feat_path, 'r')

        cv2.namedWindow('img', cv2.WINDOW_NORMAL)
        for mi in sorted(self.images.keys()):
            im = self.images[mi]
            im_name = im.name
            p3d_ids = im.point3D_ids
            p2ds = feat_file[im_name]['keypoints'][()]
            image = cv2.imread(osp.join(image_path, im_name))
            print('img_name: ', im_name)

            sems = []
            for pid in p3d_ids:
                if pid in p3d_segs.keys():
                    sems.append(p3d_segs[pid] + 1)
                else:
                    sems.append(0)
            sems = np.array(sems)

            sems = np.array(sems)
            mask = sems > 0
            img_seg = vis_seg_point(img=image, kpts=p2ds[mask], segs=sems[mask], seg_color=vis_color)

            cv2.imshow('img', img_seg)
            key = cv2.waitKey(0)
            if key == ord('q'):
                exit(0)
            elif key == ord('r'):
                # cv2.destroyAllWindows()
                return

    def extract_query_p3ds(self, log_fn, feat_fn, save_fn=None):
        if save_fn is not None:
            if osp.isfile(save_fn):
                print('{:s} exists'.format(save_fn))
                return

        loc_log = np.load(log_fn, allow_pickle=True)[()]
        fns = loc_log.keys()
        feat_file = h5py.File(feat_fn, 'r')

        out = {}
        for fn in tqdm(fns, total=len(fns)):
            matched_kpts = loc_log[fn]['keypoints_query']
            matched_p3ds = loc_log[fn]['points3D_ids']

            query_kpts = feat_file[fn]['keypoints'][()].astype(float)
            query_p3d_ids = np.zeros(shape=(query_kpts.shape[0],), dtype=int) - 1
            print('matched kpts: {}, query kpts: {}'.format(matched_kpts.shape[0], query_kpts.shape[0]))

            if matched_kpts.shape[0] > 0:
                # [M, 2, 1] - [1, 2, N] = [M, 2, N]
                dist = torch.from_numpy(matched_kpts).unsqueeze(-1) - torch.from_numpy(
                    query_kpts.transpose()).unsqueeze(0)
                dist = torch.sum(dist ** 2, dim=1)  # [M, N]
                values, idxes = torch.topk(dist, dim=1, largest=False, k=1)  # find the matches kpts with dist of 0
                values = values.numpy()
                idxes = idxes.numpy()
                for i in range(values.shape[0]):
                    if values[i, 0] < 1:
                        query_p3d_ids[idxes[i, 0]] = matched_p3ds[i]

            out[fn] = query_p3d_ids
        np.save(save_fn, out)
        feat_file.close()

    def compute_mean_scale_p3ds(self, min_obs=5, save_fn=None):
        if save_fn is not None:
            if osp.isfile(save_fn):
                with open(save_fn, 'r') as f:
                    lines = f.readlines()
                    l = lines[0].strip().split()
                    self.mean_xyz = np.array([float(v) for v in l[:3]])
                    self.scale_xyz = np.array([float(v) for v in l[3:]])
                print('{} exists'.format(save_fn))
                return

        all_xyzs = []
        for pid in self.points3D.keys():
            p3d = self.points3D[pid]
            obs = len(p3d.point2D_idxs)
            if obs < min_obs:
                continue
            all_xyzs.append(p3d.xyz)

        all_xyzs = np.array(all_xyzs)
        mean_xyz = np.ceil(np.mean(all_xyzs, axis=0))
        all_xyz_ = all_xyzs - mean_xyz

        dx = np.max(abs(all_xyz_[:, 0]))
        dy = np.max(abs(all_xyz_[:, 1]))
        dz = np.max(abs(all_xyz_[:, 2]))
        scale_xyz = np.ceil(np.array([dx, dy, dz], dtype=float).reshape(3, ))
        scale_xyz[scale_xyz < 1] = 1
        scale_xyz[scale_xyz == 0] = 1

        # self.mean_xyz = mean_xyz
        # self.scale_xyz = scale_xyz
        #
        # if save_fn is not None:
        #     with open(save_fn, 'w') as f:
        #         text = '{:.4f} {:.4f} {:.4f} {:.4f} {:.4f} {:.4f}'.format(mean_xyz[0], mean_xyz[1], mean_xyz[2],
        #                                                                   scale_xyz[0], scale_xyz[1], scale_xyz[2])
        #         f.write(text + '\n')

    def compute_statics_inlier(self, xyz, nb_neighbors=20, std_ratio=2.0):
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(xyz)

        new_pcd, inlier_ids = pcd.remove_statistical_outlier(nb_neighbors=nb_neighbors, std_ratio=std_ratio)
        return inlier_ids

    def export_features_to_directory(self, feat_fn, save_dir, with_descriptors=True):
        def print_grp_name(grp_name, object):
            try:
                n_subgroups = len(object.keys())
            except:
                n_subgroups = 0
                dataset_list.append(object.name)

        dataset_list = []
        feat_file = h5py.File(feat_fn, 'r')
        feat_file.visititems(print_grp_name)
        all_keys = []
        os.makedirs(save_dir, exist_ok=True)
        for fn in dataset_list:
            subs = fn[1:].split('/')[:-1]  # remove the first '/'
            subs = '/'.join(map(str, subs))
            if subs in all_keys:
                continue
            all_keys.append(subs)

        for fn in tqdm(all_keys, total=len(all_keys)):
            feat = feat_file[fn]
            data = {
                # 'descriptors': feat['descriptors'][()].transpose(),
                'scores': feat['scores'][()],
                'keypoints': feat['keypoints'][()],
                'image_size': feat['image_size'][()]
            }
            np.save(osp.join(save_dir, fn.replace('/', '+')), data)
        feat_file.close()

    def get_intrinsics_from_camera(self, camera):
        if camera.model in ("SIMPLE_PINHOLE", "SIMPLE_RADIAL", "RADIAL"):
            fx = fy = camera.params[0]
            cx = camera.params[1]
            cy = camera.params[2]
        elif camera.model in ("PINHOLE", "OPENCV", "OPENCV_FISHEYE", "FULL_OPENCV"):
            fx = camera.params[0]
            fy = camera.params[1]
            cx = camera.params[2]
            cy = camera.params[3]
        else:
            raise Exception("Camera model not supported")

        # intrinsics
        K = np.identity(3)
        K[0, 0] = fx
        K[1, 1] = fy
        K[0, 2] = cx
        K[1, 2] = cy
        return K

    def compress_map_by_projection_v2(self, vrf_path, point3d_desc_path, vrf_frames=1, covisible_frames=20, radius=20,
                                      nkpts=-1, save_dir=None):
        def sparsify_by_grid(h, w, uvs, scores):
            nh = np.ceil(h / radius).astype(int)
            nw = np.ceil(w / radius).astype(int)
            grid = {}
            for ip in range(uvs.shape[0]):
                p = uvs[ip]
                iw = np.rint(p[0] // radius).astype(int)
                ih = np.rint(p[1] // radius).astype(int)
                idx = ih * nw + iw
                if idx in grid.keys():
                    if scores[ip] <= grid[idx]['score']:
                        continue
                    else:
                        grid[idx]['score'] = scores[ip]
                        grid[idx]['ip'] = ip
                else:
                    grid[idx] = {
                        'score': scores[ip],
                        'ip': ip
                    }

            retained_ips = [grid[v]['ip'] for v in grid.keys()]
            retained_ips = np.array(retained_ips)
            return retained_ips

        def choose_valid_p3ds(current_frame_id, covisible_frame_ids, reserved_images):
            curr_p3d_ids = []
            curr_xyzs = []
            for pid in self.images[current_frame_id].point3D_ids:
                if pid == -1:
                    continue
                if pid not in self.points3D.keys():
                    continue
                curr_p3d_ids.append(pid)
                curr_xyzs.append(self.points3D[pid].xyz)
            curr_xyzs = np.array(curr_xyzs)  # [N, 3]
            curr_xyzs_homo = np.hstack([curr_xyzs, np.ones((curr_xyzs.shape[0], 1), dtype=curr_xyzs.dtype)])  # [N, 4]

            curr_mask = np.array([True for mi in range(curr_xyzs.shape[0])])  # keep all at first
            for iim in covisible_frame_ids:
                cam_id = self.images[iim].camera_id
                width = self.cameras[cam_id].width
                height = self.cameras[cam_id].height
                qvec = self.images[iim].qvec
                tcw = self.images[iim].tvec
                Rcw = qvec2rotmat(qvec=qvec)
                Tcw = np.eye(4, dtype=float)
                Tcw[:3, :3] = Rcw
                Tcw[:3, 3] = tcw.reshape(3, )

                uvs = reserved_images[iim]['xys']
                K = self.get_intrinsics_from_camera(camera=self.cameras[cam_id])
                proj_xys = K @ (Tcw @ curr_xyzs_homo.transpose())[:3, :]  # [3, ]
                proj_xys = proj_xys.transpose()
                depth = proj_xys[:, 2]
                proj_xys[:, 0] = proj_xys[:, 0] / depth
                proj_xys[:, 1] = proj_xys[:, 1] / depth

                mask_in_image = (proj_xys[:, 0] >= 0) * (proj_xys[:, 0] < width) * (proj_xys[:, 1] >= 0) * (
                        proj_xys[:, 1] < height)
                mask_depth = proj_xys[:, 2] > 0

                dist_proj_uv = torch.from_numpy(proj_xys[:, :2])[..., None] - \
                               torch.from_numpy(uvs[:, :2].transpose())[None]
                dist_proj_uv = torch.sqrt(torch.sum(dist_proj_uv ** 2, dim=1))  # [N, M]
                min_dist = torch.min(dist_proj_uv, dim=1)[0].numpy()
                mask_close_to_img = (min_dist <= radius)

                mask = mask_in_image * mask_depth * mask_close_to_img  # p3ds to be discarded

                curr_mask = curr_mask * (1 - mask)

            chosen_p3d_ids = []
            for mi in range(curr_mask.shape[0]):
                if curr_mask[mi]:
                    chosen_p3d_ids.append(curr_p3d_ids[mi])

            return chosen_p3d_ids

        vrf_data = np.load(vrf_path, allow_pickle=True)[()]
        p3d_ids_in_vrf = []
        image_ids_in_vrf = []
        for sid in vrf_data.keys():
            svrf = vrf_data[sid]
            svrf_keys = [vi for vi in range(vrf_frames)]
            for vi in svrf_keys:
                if vi not in svrf.keys():
                    continue
                image_id = svrf[vi]['image_id']
                if image_id in image_ids_in_vrf:
                    continue
                image_ids_in_vrf.append(image_id)
                for pid in svrf[vi]['original_points3d']:
                    if pid in p3d_ids_in_vrf:
                        continue
                    p3d_ids_in_vrf.append(pid)

        print('Find {:d} images and {:d} 3D points in vrf'.format(len(image_ids_in_vrf), len(p3d_ids_in_vrf)))

        # first_vrf_images_covis = {}
        retained_image_ids = {}
        for frame_id in image_ids_in_vrf:
            observed = self.images[frame_id].point3D_ids
            xys = self.images[frame_id].xys
            covis = defaultdict(int)
            valid_xys = []
            valid_p3d_ids = []
            for xy, pid in zip(xys, observed):
                if pid == -1:
                    continue
                if pid not in self.points3D.keys():
                    continue
                valid_xys.append(xy)
                valid_p3d_ids.append(pid)
                for img_id in self.points3D[pid].image_ids:
                    covis[img_id] += 1

            retained_image_ids[frame_id] = {
                'xys': np.array(valid_xys),
                'p3d_ids': valid_p3d_ids,
            }

            print('Find {:d} valid connected frames'.format(len(covis.keys())))

            covis_ids = np.array(list(covis.keys()))
            covis_num = np.array([covis[i] for i in covis_ids])

            if len(covis_ids) <= covisible_frames:
                sel_covis_ids = covis_ids[np.argsort(-covis_num)]
            else:
                ind_top = np.argpartition(covis_num, -covisible_frames)
                ind_top = ind_top[-covisible_frames:]  # unsorted top k
                ind_top = ind_top[np.argsort(-covis_num[ind_top])]
                sel_covis_ids = [covis_ids[i] for i in ind_top]

            covis_frame_ids = [frame_id]
            for iim in sel_covis_ids:
                if iim == frame_id:
                    continue
                if iim in retained_image_ids.keys():
                    covis_frame_ids.append(iim)
                    continue

                chosen_p3d_ids = choose_valid_p3ds(current_frame_id=iim, covisible_frame_ids=covis_frame_ids,
                                                   reserved_images=retained_image_ids)
                if len(chosen_p3d_ids) == 0:
                    continue

                xys = []
                for xy, pid in zip(self.images[iim].xys, self.images[iim].point3D_ids):
                    if pid in chosen_p3d_ids:
                        xys.append(xy)
                xys = np.array(xys)

                covis_frame_ids.append(iim)
                retained_image_ids[iim] = {
                    'xys': xys,
                    'p3d_ids': chosen_p3d_ids,
                }

        new_images = {}
        new_point3Ds = {}
        new_cameras = {}
        for iim in retained_image_ids.keys():
            p3d_ids = retained_image_ids[iim]['p3d_ids']
            ''' this step reduces the performance
            for v in self.images[iim].point3D_ids:
                if v == -1 or v not in self.points3D:
                    continue
                if v in p3d_ids:
                    continue
                p3d_ids.append(v)
            '''

            xyzs = np.array([self.points3D[pid].xyz for pid in p3d_ids])
            obs = np.array([len(self.points3D[pid].point2D_idxs) for pid in p3d_ids])
            xys = self.images[iim].xys
            cam_id = self.images[iim].camera_id
            name = self.images[iim].name
            qvec = self.images[iim].qvec
            tvec = self.images[iim].tvec

            if nkpts > 0 and len(p3d_ids) > nkpts:
                proj_uvs = self.reproject(img_id=iim, xyzs=xyzs)
                width = self.cameras[cam_id].width
                height = self.cameras[cam_id].height
                sparsified_idxs = sparsify_by_grid(h=height, w=width, uvs=proj_uvs[:, :2], scores=obs)

                print('org / new kpts: ', len(p3d_ids), sparsified_idxs.shape)

                p3d_ids = [p3d_ids[k] for k in sparsified_idxs]

            new_images[iim] = Image(id=iim, qvec=qvec, tvec=tvec,
                                    camera_id=cam_id,
                                    name=name,
                                    xys=np.array([]),
                                    point3D_ids=np.array(p3d_ids))

            if cam_id not in new_cameras.keys():
                new_cameras[cam_id] = self.cameras[cam_id]

            for pid in p3d_ids:
                if pid in new_point3Ds.keys():
                    new_point3Ds[pid]['image_ids'].append(iim)
                else:
                    xyz = self.points3D[pid].xyz
                    rgb = self.points3D[pid].rgb
                    error = self.points3D[pid].error

                    new_point3Ds[pid] = {
                        'image_ids': [iim],
                        'rgb': rgb,
                        'xyz': xyz,
                        'error': error
                    }

        new_point3Ds_to_save = {}
        for pid in new_point3Ds.keys():
            image_ids = new_point3Ds[pid]['image_ids']
            if len(image_ids) == 0:
                continue
            xyz = new_point3Ds[pid]['xyz']
            rgb = new_point3Ds[pid]['rgb']
            error = new_point3Ds[pid]['error']

            new_point3Ds_to_save[pid] = Point3D(id=pid, xyz=xyz, rgb=rgb, error=error, image_ids=np.array(image_ids),
                                                point2D_idxs=np.array([]))

        print('Retain {:d}/{:d} images and {:d}/{:d} 3D points'.format(len(new_images), len(self.images),
                                                                       len(new_point3Ds), len(self.points3D)))

        if save_dir is not None:
            os.makedirs(save_dir, exist_ok=True)
            # write_images_binary(images=new_image_ids,
            #                     path_to_model_file=osp.join(save_dir, 'images.bin'))
            # write_points3d_binary(points3D=new_point3Ds,
            #                       path_to_model_file=osp.join(save_dir, 'points3D.bin'))
            write_compressed_images_binary(images=new_images,
                                           path_to_model_file=osp.join(save_dir, 'images.bin'))
            write_cameras_binary(cameras=new_cameras,
                                 path_to_model_file=osp.join(save_dir, 'cameras.bin'))
            write_compressed_points3d_binary(points3D=new_point3Ds_to_save,
                                             path_to_model_file=osp.join(save_dir, 'points3D.bin'))

            # Save 3d descriptors
            p3d_desc = np.load(point3d_desc_path, allow_pickle=True)[()]
            comp_p3d_desc = {}
            for k in new_point3Ds_to_save.keys():
                if k not in p3d_desc.keys():
                    print(k)
                    continue
                comp_p3d_desc[k] = deepcopy(p3d_desc[k])
            np.save(osp.join(save_dir, point3d_desc_path.split('/')[-1]), comp_p3d_desc)
            print('Save data to {:s}'.format(save_dir))


def process_dataset(dataset, dataset_dir, sfm_dir, save_dir, feature='sfd2', matcher='gml'):
    # dataset_dir = '/scratches/flyer_3/fx221/dataset'
    # sfm_dir = '/scratches/flyer_2/fx221/localization/outputs'  # your sfm results (cameras, images, points3D) and features
    # save_dir = '/scratches/flyer_3/fx221/exp/localizer'
    # local_feat = 'sfd2'
    # matcher = 'gml'
    # hloc_results_dir = '/scratches/flyer_2/fx221/exp/sgd2'

    # config_path = 'configs/datasets/CUED.yaml'
    # config_path = 'configs/datasets/7Scenes.yaml'
    # config_path = 'configs/datasets/12Scenes.yaml'
    # config_path = 'configs/datasets/CambridgeLandmarks.yaml'
    # config_path = 'configs/datasets/Aachen.yaml'

    # config_path = 'configs/datasets/Aria.yaml'
    # config_path = 'configs/datasets/DarwinRGB.yaml'
    # config_path = 'configs/datasets/ACUED.yaml'
    # config_path = 'configs/datasets/JesusCollege.yaml'
    # config_path = 'configs/datasets/CUED2Kings.yaml'

    config_path = 'configs/datasets/{:s}.yaml'.format(dataset)
    with open(config_path, 'rt') as f:
        configs = yaml.load(f, Loader=yaml.Loader)
    print(configs)

    dataset = configs['dataset']
    all_scenes = configs['scenes']
    for scene in all_scenes:
        n_cluster = configs[scene]['n_cluster']
        cluster_mode = configs[scene]['cluster_mode']
        cluster_method = configs[scene]['cluster_method']
        # if scene not in ['heads']:
        #     continue

        print('scene: ', scene, cluster_mode, cluster_method)
        # hloc_path = osp.join(hloc_root, dataset, scene)
        sfm_path = osp.join(sfm_dir, scene)
        save_path = osp.join(save_dir, feature + '-' + matcher, dataset, scene)

        n_vrf = 1
        n_cov = 30
        radius = 20
        n_kpts = 0

        if dataset in ['Aachen']:
            image_path = osp.join(dataset_dir, scene, 'images/images_upright')
            min_obs = 250
            filtering_outliers = True
            threshold = 0.2
            radius = 32

        elif dataset in ['CambridgeLandmarks', ]:
            image_path = osp.join(dataset_dir, scene)
            min_obs = 250
            filtering_outliers = True
            threshold = 0.2
            radius = 64
        elif dataset in ['Aria']:
            image_path = osp.join(dataset_dir, scene)
            min_obs = 150
            filtering_outliers = False
            threshold = 0.01
            radius = 15
        elif dataset in ['DarwinRGB']:
            image_path = osp.join(dataset_dir, scene)
            min_obs = 150
            filtering_outliers = True
            threshold = 0.2
            radius = 16
        elif dataset in ['ACUED']:
            image_path = osp.join(dataset_dir, scene)
            min_obs = 250
            filtering_outliers = True
            threshold = 0.2
            radius = 32
        elif dataset in ['7Scenes', '12Scenes']:
            image_path = osp.join(dataset_dir, scene)
            min_obs = 150
            filtering_outliers = False
            threshold = 0.01
            radius = 15
        else:
            image_path = osp.join(dataset_dir, scene)
            min_obs = 250
            filtering_outliers = True
            threshold = 0.2
            radius = 32

        # comp_map_sub_path = 'comp_model_n{:d}_{:s}_{:s}_vrf{:d}_cov{:d}_r{:d}_np{:d}_projection_v2'.format(n_cluster,
        #                                                                                                    cluster_mode,
        #                                                                                                    cluster_method,
        #                                                                                                    n_vrf,
        #                                                                                                    n_cov,
        #                                                                                                    radius,
        #                                                                                                    n_kpts)
        comp_map_sub_path = 'compress_model_{:s}'.format(cluster_method)
        seg_fn = osp.join(save_path,
                          'point3D_cluster_n{:d}_{:s}_{:s}.npy'.format(n_cluster, cluster_mode, cluster_method))
        vrf_fn = osp.join(save_path,
                          'point3D_vrf_n{:d}_{:s}_{:s}.npy'.format(n_cluster, cluster_mode, cluster_method))
        vrf_img_dir = osp.join(save_path,
                               'point3D_vrf_n{:d}_{:s}_{:s}'.format(n_cluster, cluster_mode, cluster_method))
        # p3d_query_fn = osp.join(save_path,
        #                         'point3D_query_n{:d}_{:s}_{:s}.npy'.format(n_cluster, cluster_mode, cluster_method))
        comp_map_path = osp.join(save_path, comp_map_sub_path)

        os.makedirs(save_path, exist_ok=True)

        rmap = RecMap()
        rmap.load_sfm_model(path=osp.join(sfm_path, 'sfm_{:s}-{:s}'.format(feature, matcher)))
        if filtering_outliers:
            rmap.remove_statics_outlier(nb_neighbors=20, std_ratio=2.0)

        # extract keypoints to train the recognition model (descriptors are recomputed from augmented db images)
        # we do this for ddp training (reading h5py file is not supported)
        rmap.export_features_to_directory(feat_fn=osp.join(sfm_path, 'feats-{:s}.h5'.format(feature)),
                                          save_dir=osp.join(save_path, 'feats'))  # only once for training

        rmap.cluster(k=n_cluster, mode=cluster_mode, save_fn=seg_fn, method=cluster_method, threshold=threshold)
        # rmap.visualize_3Dpoints()
        rmap.load_segmentation(path=seg_fn)
        # rmap.visualize_segmentation(p3d_segs=rmap.p3d_seg, points3D=rmap.points3D)

        # Assign each 3D point a desciptor and discard all 2D images and descriptors - for localization
        rmap.assign_point3D_descriptor(
            feature_fn=osp.join(sfm_path, 'feats-{:s}.h5'.format(feature)),
            save_fn=osp.join(save_path, 'point3D_desc.npy'.format(n_cluster, cluster_mode)),
            n_process=32)  # only once

        # exit(0)
        # rmap.visualize_segmentation_on_image(p3d_segs=rmap.p3d_seg, image_path=image_path, feat_path=feat_path)

        # for query images only - for evaluation
        # rmap.extract_query_p3ds(
        #     log_fn=osp.join(hloc_path, 'hloc_feats-{:s}_{:s}_loc.npy'.format(local_feat, matcher)),
        #     feat_fn=osp.join(sfm_path, 'feats-{:s}.h5'.format(local_feat)),
        #     save_fn=p3d_query_fn)
        # continue

        # up-to-date
        rmap.create_virtual_frame_3(
            save_fn=vrf_fn,
            save_vrf_dir=vrf_img_dir,
            image_root=image_path,
            show_time=5,
            min_cover_ratio=0.9,
            radius=radius,
            depth_scale=2.5,  # 1.2 by default
            min_obs=min_obs,
            n_vrf=10,
            covisible_frame=n_cov,
            ignored_cameras=[])

        # up-to-date
        rmap.compress_map_by_projection_v2(
            vrf_frames=n_vrf,
            vrf_path=vrf_fn,
            point3d_desc_path=osp.join(save_path, 'point3D_desc.npy'),
            save_dir=comp_map_path,
            covisible_frames=n_cov,
            radius=radius,
            nkpts=n_kpts,
        )

        # exit(0)
        # soft_link_compress_path = osp.join(save_path, 'compress_model_{:s}'.format(cluster_method))
        os.chdir(save_path)
        # if osp.isdir(soft_link_compress_path):
        #     os.unlink(soft_link_compress_path)
        # os.symlink(comp_map_sub_path, 'compress_model_{:s}'.format(cluster_method))
        # create a soft link of the full model for training
        if not osp.isdir('model'):
            os.symlink(osp.join(sfm_path, 'sfm_{:s}-{:s}'.format(feature, matcher)), '3D-models')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--dataset', type=str, required=True, help='dataset name')
    parser.add_argument('--dataset_dir', type=str, required=True, help='dataset dir')
    parser.add_argument('--sfm_dir', type=str, required=True, help='sfm dir')
    parser.add_argument('--save_dir', type=str, required=True, help='dir to save the landmarks data')
    parser.add_argument('--feature', type=str, default='sfd2', help='feature name e.g., SP, SFD2')
    parser.add_argument('--matcher', type=str, default='gml', help='matcher name e.g., SG, LSG, gml')

    args = parser.parse_args()

    process_dataset(
        dataset=args.dataset,
        dataset_dir=args.dataset_dir,
        sfm_dir=args.sfm_dir,
        save_dir=args.save_dir,
        feature=args.feature,
        matcher=args.matcher)