{"metadata":{"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":71885,"databundleVersionId":8143495,"sourceType":"competition"},{"sourceId":3414836,"sourceType":"datasetVersion","datasetId":2058261},{"sourceId":5373920,"sourceType":"datasetVersion","datasetId":3117886},{"sourceId":5850511,"sourceType":"datasetVersion","datasetId":3364321},{"sourceId":7884485,"sourceType":"datasetVersion","datasetId":4628051},{"sourceId":7884725,"sourceType":"datasetVersion","datasetId":4628331},{"sourceId":170475544,"sourceType":"kernelVersion"},{"sourceId":170565695,"sourceType":"kernelVersion"},{"sourceId":172469456,"sourceType":"kernelVersion"},{"sourceId":173217852,"sourceType":"kernelVersion"},{"sourceId":174129945,"sourceType":"kernelVersion"},{"sourceId":175679956,"sourceType":"kernelVersion"},{"sourceId":175684111,"sourceType":"kernelVersion"},{"sourceId":176463227,"sourceType":"kernelVersion"},{"sourceId":3734,"sourceType":"modelInstanceVersion","modelInstanceId":2661},{"sourceId":3735,"sourceType":"modelInstanceVersion","modelInstanceId":2662},{"sourceId":3736,"sourceType":"modelInstanceVersion","modelInstanceId":2663},{"sourceId":3840,"sourceType":"modelInstanceVersion","modelInstanceId":2742},{"sourceId":3846,"sourceType":"modelInstanceVersion","modelInstanceId":2747},{"sourceId":4534,"sourceType":"modelInstanceVersion","modelInstanceId":3326},{"sourceId":17191,"sourceType":"modelInstanceVersion","modelInstanceId":14317},{"sourceId":17555,"sourceType":"modelInstanceVersion","modelInstanceId":14611}],"dockerImageVersionId":30733,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true},"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.10.13"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Dependencies","metadata":{}},{"cell_type":"code","source":"!python -m pip install --no-deps /kaggle/input/dependencies-imc/pycolmap/pycolmap-0.4.0-cp310-cp310-manylinux2014_x86_64.whl\n!python -m pip install --no-deps /kaggle/input/dependencies-imc/safetensors/safetensors-0.4.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl\n!python -m pip install --no-index --find-links=/kaggle/input/dependencies-imc/transformers/ transformers > /dev/null\n!python -m pip install  --no-deps /kaggle/input/imc2024-packages-lightglue-rerun-kornia/lightglue-0.0-py3-none-any.whl\n\n# dkm\n!python -m pip install --no-index --find-links=/kaggle/input/dkm-dependencies/packages einops > /dev/null\n\n# match former\n!python -m pip install --no-index --find-links=/kaggle/input/matchformer-dependencies yacs > /dev/null\n\n# lightglue models\n!mkdir -p /root/.cache/torch/hub/checkpoints\n!cp /kaggle/input/aliked/pytorch/aliked-n16/1/* /root/.cache/torch/hub/checkpoints/\n!cp /kaggle/input/lightglue/pytorch/aliked/1/* /root/.cache/torch/hub/checkpoints/\n!cp /kaggle/input/lightglue/pytorch/aliked/1/aliked_lightglue.pth /root/.cache/torch/hub/checkpoints/aliked_lightglue_v0-1_arxiv-pth\n!cp /kaggle/input/pytorch-lightglue-models/* /root/.cache/torch/hub/checkpoints/\n\n# dkm model\n!mkdir -p /root/.cache/torch/hub/checkpoints\n!cp /kaggle/input/dkm-dependencies/DKMv3_outdoor.pth /root/.cache/torch/hub/checkpoints/\n\n# check rotation\n!python -m pip install --no-index --find-links=/kaggle/input/pkg-check-orientation/ check_orientation==0.0.5 > /dev/null\n!cp /kaggle/input/pkg-check-orientation/2020-11-16_resnext50_32x4d.zip /root/.cache/torch/hub/checkpoints/","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2024-06-26T15:45:17.13176Z","iopub.execute_input":"2024-06-26T15:45:17.132147Z","iopub.status.idle":"2024-06-26T15:47:40.249782Z","shell.execute_reply.started":"2024-06-26T15:45:17.132116Z","shell.execute_reply":"2024-06-26T15:47:40.248334Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import concurrent.futures\nimport gc\nimport os\nimport sqlite3\nimport sys\nimport time\nimport warnings\n\nimport math\nimport numpy as np\n\nfrom IPython.display import clear_output\nfrom tqdm import tqdm\nfrom fastprogress import progress_bar\n\nimport torch\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\n\nimport cv2\nimport torchvision.transforms as T\n\nfrom torchvision.io import read_image as T_read_image, ImageReadMode\n\nimport kornia as K\nimport kornia.feature as KF\nfrom PIL import Image\nimport pycolmap\n\nimport timm\nfrom timm.data import resolve_data_config\nfrom dataclasses import dataclass\n\nfrom lightglue import ALIKED,  DoGHardNet, DISK, SIFT\nimport sys\nsys.path.append(\"../input/super-glue-pretrained-network\")\n\nfrom collections import defaultdict\nfrom copy import deepcopy\nimport h5py\nfrom h5py import *\nfrom torchvision import transforms","metadata":{"execution":{"iopub.status.busy":"2024-06-26T15:47:40.252264Z","iopub.execute_input":"2024-06-26T15:47:40.25262Z","iopub.status.idle":"2024-06-26T15:47:48.896509Z","shell.execute_reply.started":"2024-06-26T15:47:40.252587Z","shell.execute_reply":"2024-06-26T15:47:48.895416Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CONFIG:\n    # DEBUG Settings\n    DRY_RUN = False\n    DRY_RUN_MAX_IMAGES = 10\n\n    # Pipeline settings\n    NUM_CORES = 2\n    \n    # COLMAP Reconstruction\n    CAMERA_MODEL = \"simple-radial\"\n    \n    # Rotation correction\n    ROTATION_CORRECTION = 0\n    \n    # Keypoints handling\n    MERGE_PARAMS = {\n        \"min_matches\" : 15,\n        \"filter_FundamentalMatrix\" : False,\n        \"filter_iterations\" : 10,\n        \"filter_threshold\" : 8,\n    }\n    \n    use_aliked_lightglue = True\n        \n    params_aliked_lightglue = {\n        \"num_features\" : 8192,\n        \"detection_threshold\" : 0.001,\n        \"min_matches\" : 50,\n        \"resize_to\" : 1024,\n    }\n\ndevice=torch.device('cuda')\nMODEL_DIR = \"/kaggle/input/kornia-local-feature-weights/\"\nHARDNET_PT = \"/kaggle/input/hardnet8v2/hardnet8v2.pt\"\nIMG_MODEL_0 = \"/kaggle/input/tf-efficientnet/pytorch/tf-efficientnet-b5/1/tf_efficientnet_b5_ra-9a3e5369.pth\"\nIMG_MODEL_1 = '/kaggle/input/tf-efficientnet/pytorch/tf-efficientnet-b6/1/tf_efficientnet_b6_aa-80ba17e4.pth'\nIMG_MODEL_2 = '/kaggle/input/tf-efficientnet/pytorch/tf-efficientnet-b7/1/tf_efficientnet_b7_ra-6c08e654.pth'","metadata":{"execution":{"iopub.status.busy":"2024-06-26T15:47:48.897843Z","iopub.execute_input":"2024-06-26T15:47:48.898136Z","iopub.status.idle":"2024-06-26T15:47:48.905696Z","shell.execute_reply.started":"2024-06-26T15:47:48.898111Z","shell.execute_reply":"2024-06-26T15:47:48.90468Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"IS_PYTHON3 = sys.version_info[0] >= 3\n\nMAX_IMAGE_ID = 2**31 - 1\n\nCREATE_CAMERAS_TABLE = \"\"\"CREATE TABLE IF NOT EXISTS cameras (\n    camera_id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,\n    model INTEGER NOT NULL,\n    width INTEGER NOT NULL,\n    height INTEGER NOT NULL,\n    params BLOB,\n    prior_focal_length INTEGER NOT NULL)\"\"\"\n\nCREATE_DESCRIPTORS_TABLE = \"\"\"CREATE TABLE IF NOT EXISTS descriptors (\n    image_id INTEGER PRIMARY KEY NOT NULL,\n    rows INTEGER NOT NULL,\n    cols INTEGER NOT NULL,\n    data BLOB,\n    FOREIGN KEY(image_id) REFERENCES images(image_id) ON DELETE CASCADE)\"\"\"\n\nCREATE_IMAGES_TABLE = \"\"\"CREATE TABLE IF NOT EXISTS images (\n    image_id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,\n    name TEXT NOT NULL UNIQUE,\n    camera_id INTEGER NOT NULL,\n    prior_qw REAL,\n    prior_qx REAL,\n    prior_qy REAL,\n    prior_qz REAL,\n    prior_tx REAL,\n    prior_ty REAL,\n    prior_tz REAL,\n    CONSTRAINT image_id_check CHECK(image_id >= 0 and image_id < {}),\n    FOREIGN KEY(camera_id) REFERENCES cameras(camera_id))\n\"\"\".format(MAX_IMAGE_ID)\n\nCREATE_TWO_VIEW_GEOMETRIES_TABLE = \"\"\"\nCREATE TABLE IF NOT EXISTS two_view_geometries (\n    pair_id INTEGER PRIMARY KEY NOT NULL,\n    rows INTEGER NOT NULL,\n    cols INTEGER NOT NULL,\n    data BLOB,\n    config INTEGER NOT NULL,\n    F BLOB,\n    E BLOB,\n    H BLOB)\n\"\"\"\n\nCREATE_KEYPOINTS_TABLE = \"\"\"CREATE TABLE IF NOT EXISTS keypoints (\n    image_id INTEGER PRIMARY KEY NOT NULL,\n    rows INTEGER NOT NULL,\n    cols INTEGER NOT NULL,\n    data BLOB,\n    FOREIGN KEY(image_id) REFERENCES images(image_id) ON DELETE CASCADE)\n\"\"\"\n\nCREATE_MATCHES_TABLE = \"\"\"CREATE TABLE IF NOT EXISTS matches (\n    pair_id INTEGER PRIMARY KEY NOT NULL,\n    rows INTEGER NOT NULL,\n    cols INTEGER NOT NULL,\n    data BLOB)\"\"\"\n\nCREATE_NAME_INDEX = \\\n    \"CREATE UNIQUE INDEX IF NOT EXISTS index_name ON images(name)\"\n\nCREATE_ALL = \"; \".join([\n    CREATE_CAMERAS_TABLE,\n    CREATE_IMAGES_TABLE,\n    CREATE_KEYPOINTS_TABLE,\n    CREATE_DESCRIPTORS_TABLE,\n    CREATE_MATCHES_TABLE,\n    CREATE_TWO_VIEW_GEOMETRIES_TABLE,\n    CREATE_NAME_INDEX\n])","metadata":{"execution":{"iopub.status.busy":"2024-06-26T15:47:48.909355Z","iopub.execute_input":"2024-06-26T15:47:48.909667Z","iopub.status.idle":"2024-06-26T15:47:48.925383Z","shell.execute_reply.started":"2024-06-26T15:47:48.90964Z","shell.execute_reply":"2024-06-26T15:47:48.92441Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def image_ids_to_pair_id(image_id1, image_id2):\n    if image_id1 > image_id2:\n        image_id1, image_id2 = image_id2, image_id1\n    return image_id1 * MAX_IMAGE_ID + image_id2\n\n\ndef pair_id_to_image_ids(pair_id):\n    image_id2 = pair_id % MAX_IMAGE_ID\n    image_id1 = (pair_id - image_id2) / MAX_IMAGE_ID\n    return image_id1, image_id2\n\n\ndef array_to_blob(array):\n    if IS_PYTHON3:\n        return array.tostring()\n    else:\n        return np.getbuffer(array)\n\n\ndef blob_to_array(blob, dtype, shape=(-1,)):\n    if IS_PYTHON3:\n        return np.fromstring(blob, dtype=dtype).reshape(*shape)\n    else:\n        return np.frombuffer(blob, dtype=dtype).reshape(*shape)\n\n\nclass COLMAPDatabase(sqlite3.Connection):\n\n    @staticmethod\n    def connect(database_path):\n        return sqlite3.connect(database_path, factory=COLMAPDatabase)\n\n\n    def __init__(self, *args, **kwargs):\n        super(COLMAPDatabase, self).__init__(*args, **kwargs)\n\n        self.create_tables = lambda: self.executescript(CREATE_ALL)\n        self.create_cameras_table = \\\n            lambda: self.executescript(CREATE_CAMERAS_TABLE)\n        self.create_descriptors_table = \\\n            lambda: self.executescript(CREATE_DESCRIPTORS_TABLE)\n        self.create_images_table = \\\n            lambda: self.executescript(CREATE_IMAGES_TABLE)\n        self.create_two_view_geometries_table = \\\n            lambda: self.executescript(CREATE_TWO_VIEW_GEOMETRIES_TABLE)\n        self.create_keypoints_table = \\\n            lambda: self.executescript(CREATE_KEYPOINTS_TABLE)\n        self.create_matches_table = \\\n            lambda: self.executescript(CREATE_MATCHES_TABLE)\n        self.create_name_index = lambda: self.executescript(CREATE_NAME_INDEX)\n\n    def add_camera(self, model, width, height, params,\n                prior_focal_length=False, camera_id=None):\n        params = np.asarray(params, np.float64)\n        cursor = self.execute(\n            \"INSERT INTO cameras VALUES (?, ?, ?, ?, ?, ?)\",\n            (camera_id, model, width, height, array_to_blob(params),\n            prior_focal_length))\n        return cursor.lastrowid\n\n    def add_image(self, name, camera_id,\n                prior_q=np.zeros(4), prior_t=np.zeros(3), image_id=None):\n        cursor = self.execute(\n            \"INSERT INTO images VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)\",\n            (image_id, name, camera_id, prior_q[0], prior_q[1], prior_q[2],\n            prior_q[3], prior_t[0], prior_t[1], prior_t[2]))\n        return cursor.lastrowid\n\n    def add_keypoints(self, image_id, keypoints):\n        assert(len(keypoints.shape) == 2)\n        assert(keypoints.shape[1] in [2, 4, 6])\n\n        keypoints = np.asarray(keypoints, np.float32)\n        self.execute(\n            \"INSERT INTO keypoints VALUES (?, ?, ?, ?)\",\n            (image_id,) + keypoints.shape + (array_to_blob(keypoints),))\n\n    def add_descriptors(self, image_id, descriptors):\n        descriptors = np.ascontiguousarray(descriptors, np.uint8)\n        self.execute(\n            \"INSERT INTO descriptors VALUES (?, ?, ?, ?)\",\n            (image_id,) + descriptors.shape + (array_to_blob(descriptors),))\n\n    def add_matches(self, image_id1, image_id2, matches):\n        assert(len(matches.shape) == 2)\n        assert(matches.shape[1] == 2)\n\n        if image_id1 > image_id2:\n            matches = matches[:,::-1]\n\n        pair_id = image_ids_to_pair_id(image_id1, image_id2)\n        matches = np.asarray(matches, np.uint32)\n        self.execute(\n            \"INSERT INTO matches VALUES (?, ?, ?, ?)\",\n            (pair_id,) + matches.shape + (array_to_blob(matches),))\n\n    def add_two_view_geometry(self, image_id1, image_id2, matches,\n                            F=np.eye(3), E=np.eye(3), H=np.eye(3), config=2):\n        assert(len(matches.shape) == 2)\n        assert(matches.shape[1] == 2)\n\n        if image_id1 > image_id2:\n            matches = matches[:,::-1]\n\n        pair_id = image_ids_to_pair_id(image_id1, image_id2)\n        matches = np.asarray(matches, np.uint32)\n        F = np.asarray(F, dtype=np.float64)\n        E = np.asarray(E, dtype=np.float64)\n        H = np.asarray(H, dtype=np.float64)\n        self.execute(\n            \"INSERT INTO two_view_geometries VALUES (?, ?, ?, ?, ?, ?, ?, ?)\",\n            (pair_id,) + matches.shape + (array_to_blob(matches), config,\n            array_to_blob(F), array_to_blob(E), array_to_blob(H)))\n\ndef get_focal(image_path, err_on_default=False):\n    image         = Image.open(image_path)\n    max_size      = max(image.size)\n\n    exif = image.getexif()\n    focal = None\n    if exif is not None:\n        focal_35mm = None\n        # https://github.com/colmap/colmap/blob/d3a29e203ab69e91eda938d6e56e1c7339d62a99/src/util/bitmap.cc#L299\n        for tag, value in exif.items():\n            focal_35mm = None\n            if ExifTags.TAGS.get(tag, None) == 'FocalLengthIn35mmFilm':\n                focal_35mm = float(value)\n                break\n\n        if focal_35mm is not None:\n            focal = focal_35mm / 35. * max_size\n    \n    if focal is None:\n        if err_on_default:\n            raise RuntimeError(\"Failed to find focal length\")\n\n        # failed to find it in exif, use prior\n        FOCAL_PRIOR = 1.2\n        focal = FOCAL_PRIOR * max_size\n\n    return focal\n\ndef create_camera(db, image_path, camera_model):\n    image         = Image.open(image_path)\n    width, height = image.size\n\n    focal = get_focal(image_path)\n\n    if camera_model == 'simple-pinhole':\n        model = 0 # simple pinhole\n        param_arr = np.array([focal, width / 2, height / 2])\n    if camera_model == 'pinhole':\n        model = 1 # pinhole\n        param_arr = np.array([focal, focal, width / 2, height / 2])\n    elif camera_model == 'simple-radial':\n        model = 2 # simple radial\n        param_arr = np.array([focal, width / 2, height / 2, 0.1])\n    elif camera_model == 'opencv':\n        model = 4 # opencv\n        param_arr = np.array([focal, focal, width / 2, height / 2, 0., 0., 0., 0.])\n        \n    return db.add_camera(model, width, height, param_arr)\n\n\ndef add_keypoints(db, h5_path, image_path, img_ext, camera_model, single_camera = True):\n    keypoint_f = h5py.File(os.path.join(h5_path, 'keypoints.h5'), 'r')\n\n    camera_id = None\n    fname_to_id = {}\n    for filename in tqdm(list(keypoint_f.keys())):\n        keypoints = keypoint_f[filename][()]\n\n        fname_with_ext = filename# + img_ext\n        path = os.path.join(image_path, fname_with_ext)\n        if not os.path.isfile(path):\n            raise IOError(f'Invalid image path {path}')\n\n        if camera_id is None or not single_camera:\n            camera_id = create_camera(db, path, camera_model)\n        image_id = db.add_image(fname_with_ext, camera_id)\n        fname_to_id[filename] = image_id\n\n        db.add_keypoints(image_id, keypoints)\n\n    return fname_to_id\n\ndef add_matches(db, h5_path, fname_to_id):\n    match_file = h5py.File(os.path.join(h5_path, 'matches.h5'), 'r')\n    \n    added = set()\n    n_keys = len(match_file.keys())\n    n_total = (n_keys * (n_keys - 1)) // 2\n\n    with tqdm(total=n_total) as pbar:\n        for key_1 in match_file.keys():\n            group = match_file[key_1]\n            for key_2 in group.keys():\n                id_1 = fname_to_id[key_1]\n                id_2 = fname_to_id[key_2]\n\n                pair_id = image_ids_to_pair_id(id_1, id_2)\n                if pair_id in added:\n                    warnings.warn(f'Pair {pair_id} ({id_1}, {id_2}) already added!')\n                    continue\n            \n                matches = group[key_2][()]\n                db.add_matches(id_1, id_2, matches)\n\n                added.add(pair_id)\n\n                pbar.update(1)\n                \ndef import_into_colmap(img_dir,\n                    feature_dir ='.featureout',\n                    database_path = 'colmap.db',\n                    img_ext='.jpg'):\n    db = COLMAPDatabase.connect(database_path)\n    db.create_tables()\n    single_camera = False\n    fname_to_id = add_keypoints(db, feature_dir, img_dir, img_ext, CONFIG.CAMERA_MODEL, single_camera)\n    add_matches(\n        db,\n        feature_dir,\n        fname_to_id,\n    )\n\n    db.commit()\n    return\n\ndef get_global_desc(fnames, model, model2,\n                    device =  device):\n    model = model.eval()\n    model = model.to(device)\n    model2 = model2.eval()\n    model2 = model2.to(device)\n    transform = T.Compose([\n        T.ToTensor(),\n        T.Resize(600, interpolation=T.InterpolationMode.BICUBIC),\n        T.CenterCrop(600),\n        T.Normalize(mean=[0.4850, 0.4560, 0.4060], std=[0.2290, 0.2240, 0.2250])\n    ])\n    global_descs_convnext=[]\n    for _, img_fname_full in tqdm(enumerate(fnames),total= len(fnames)):\n        img = Image.open(img_fname_full).convert('RGB')\n        timg = transform(img).unsqueeze(0).to(device)\n        with torch.no_grad():\n            desc = model.forward_features(timg.to(device)).mean(dim=(-1,2))\n            desc2 = model2.forward_features(timg.to(device)).mean(dim=(-1,2))\n            desc = desc.view(1, -1)\n            desc2 = desc2.view(1, -1)\n            desc_norm = torch.cat([desc, desc2], dim=-1)\n            desc_norm = F.normalize(desc_norm, dim=1, p=2)\n        global_descs_convnext.append(desc_norm.detach().cpu())\n    global_descs_all = torch.cat(global_descs_convnext, dim=0)\n    return global_descs_all\n\n\ndef get_img_pairs_exhaustive(img_fnames):\n    index_pairs = []\n    for i in range(len(img_fnames)):\n        for j in range(i+1, len(img_fnames)):\n            index_pairs.append((i,j))\n    return index_pairs\n\ndef dfs(i, matching_list, num_imgs, index, value, sim_th=0.2):\n    if i < num_imgs:\n        return\n    for t in index[i][value[i] > sim_th]:\n        if t != i:\n            matching_list.append(tuple(sorted((i, t))))\n    dfs(i + 1, matching_list, num_imgs, index, value, sim_th=0.2)\n\ndef get_image_pairs_shortlist(fnames,\n                            exhaustive_if_less = 20,\n                            min_pairs=40,\n                            device=torch.device('cpu'),\n                            sim_th=1.0):\n    num_imgs = len(fnames)\n    matching_list = []\n    if num_imgs <= exhaustive_if_less:\n        return get_img_pairs_exhaustive(fnames)\n\n    model = timm.create_model('tf_efficientnet_b6',\n                            checkpoint_path=IMG_MODEL_1)\n    model.eval()\n    model2 = timm.create_model('tf_efficientnet_b7',\n                            checkpoint_path=IMG_MODEL_2)                  \n    model2.eval()\n    descs = get_global_desc(fnames, model, model2, device=device)\n    del model, model2\n    torch.cuda.empty_cache()\n    import gc\n    gc.collect()\n    dm = torch.einsum('bi,ki->bk', descs, descs).detach().cpu() \n    value, index = torch.topk(dm, k=min_pairs, dim=1)\n    \n    for i in range(num_imgs-1):\n        for t in index[i][value[i]>sim_th]:\n            if t != i:\n                matching_list.append(tuple(sorted((i, t.item()))))\n    \n    dfs(0, matching_list, num_imgs, index, value)\n    matching_list = sorted(list(set(matching_list)))\n    return matching_list\n\ndef load_torch_image(fname, device=torch.device('cpu')):\n    img = K.io.load_image(fname, K.io.ImageLoadType.RGB32, device=device)[None, ...]\n    return img\n\ndef convert_coord(r, w, h, rotk):\n    if rotk == 0:\n        return r\n    elif rotk == 1:\n        rx = w-1-r[:, 1]\n        ry = r[:, 0]\n        return torch.concat([rx[None], ry[None]], dim=0).T\n    elif rotk == 2:\n        rx = w-1-r[:, 0]\n        ry = h-1-r[:, 1]\n        return torch.concat([rx[None], ry[None]], dim=0).T\n    elif rotk == 3:\n        rx = r[:, 1]\n        ry = h-1-r[:, 0]\n        return torch.concat([rx[None], ry[None]], dim=0).T\n\ndef detect_common(img_fnames,\n                model_name,\n                rots,\n                file_keypoints,\n                feature_dir = '.featureout',\n                num_features = 4096,\n                resize_to = 1024,\n                detection_threshold = 0.01,\n                device=torch.device('cpu'),\n                min_matches=15,verbose=True\n                ):\n    if not os.path.isdir(feature_dir):\n        os.makedirs(feature_dir)\n    dict_model = {\n        \"aliked\" : ALIKED,\n        \"doghardnet\" : DoGHardNet,\n        \"disk\" : DISK,\n        \"sift\" : SIFT,\n    }\n    extractor_class = dict_model[model_name]\n    \n    dtype = torch.float32\n    extractor = extractor_class(\n        max_num_keypoints=num_features, detection_threshold=detection_threshold, resize=resize_to\n    ).eval().to(device, dtype)\n    dict_kpts_cuda = {}\n    dict_descs_cuda = {}\n    for (img_path, rot_k) in zip(img_fnames, rots):\n        img_fname = img_path.split('/')[-1]\n        key = img_fname\n        with torch.inference_mode():\n            image0 = load_torch_image(img_path, device=device).to(dtype)\n            h, w = image0.shape[2], image0.shape[3]\n            image1 = torch.rot90(image0, rot_k, [2, 3])\n            feats0 = extractor.extract(image1)\n            kpts = feats0['keypoints'].reshape(-1, 2).detach()\n            descs = feats0['descriptors'].reshape(len(kpts), -1).detach()\n\n            dict_kpts_cuda[f\"{key}\"] = kpts\n            dict_descs_cuda[f\"{key}\"] = descs\n            print(f\"{model_name} > rot_k={rot_k}, kpts.shape={kpts.shape}, descs.shape={descs.shape}\")\n    del extractor\n    gc.collect()\n\n    #####################################################\n    # Matching keypoints\n    #####################################################\n    lg_matcher = KF.LightGlueMatcher(model_name, {\"width_confidence\": -1,\n                                            \"depth_confidence\": -1,\n                                            \"mp\": True if 'cuda' in str(device) else False}).eval().to(device)\n    \n    cnt_pairs = 0\n    with h5py.File(file_keypoints, mode='w') as f_match:\n        for pair_idx in tqdm(index_pairs):\n            idx1, idx2 = pair_idx\n            fname1, fname2 = img_fnames[idx1], img_fnames[idx2]\n            \n            key1, key2 = fname1.split('/')[-1], fname2.split('/')[-1]\n            \n            \n            kp1 = dict_kpts_cuda[key1]\n            kp2 = dict_kpts_cuda[key2]\n            \n            desc1 = dict_descs_cuda[key1]\n            desc2 = dict_descs_cuda[key2]\n            with torch.inference_mode():\n                dists, idxs = lg_matcher(desc1,\n                                    desc2,\n                                    KF.laf_from_center_scale_ori(kp1[None]),\n                                    KF.laf_from_center_scale_ori(kp2[None]))\n            if len(idxs)  == 0:\n                continue\n            n_matches = len(idxs)\n            kp1 = kp1[idxs[:,0], :].cpu().numpy().reshape(-1, 2).astype(np.float32)\n            kp2 = kp2[idxs[:,1], :].cpu().numpy().reshape(-1, 2).astype(np.float32)\n            group  = f_match.require_group(key1)\n            if n_matches >= min_matches:\n                group.create_dataset(key2, data=np.concatenate([kp1, kp2], axis=1))\n                cnt_pairs+=1\n                print (f'{model_name}> {key1}-{key2}: {n_matches} matches @ {cnt_pairs}th pair({model_name}+lightglue)')            \n            else:\n                print (f'{model_name}> {key1}-{key2}: {n_matches} matches --> skipped')\n    del lg_matcher\n    torch.cuda.empty_cache()\n    gc.collect()\n    return\n\ndef detect_lightglue_common(\n    img_fnames, model_name, index_pairs, feature_dir, device, file_keypoints, rots,\n    resize_to=1024,\n    detection_threshold=0.01, \n    num_features=4096, \n    min_matches=15,\n):\n    t=time.time()\n    detect_common(\n        img_fnames, model_name, rots, file_keypoints, feature_dir, \n        resize_to=resize_to,\n        num_features=num_features, \n        detection_threshold=detection_threshold, \n        device=device,\n        min_matches=min_matches,\n    )\n    gc.collect()\n    t=time.time() -t \n    print(f'Features matched in  {t:.4f} sec ({model_name}+LightGlue)')\n    return t\n\n\n# Making kornia local features loading w/o internet\nclass AffNetHardNet(KF.LocalFeature):\n    \"\"\"Convenience module, which implements KeyNet detector + AffNet + HardNet descriptor.\n\n    .. image:: _static/img/keynet_affnet.jpg\n    \"\"\"\n\n    def __init__(\n        self,\n        num_features: int = 5000,\n        upright: bool = False,\n        device = torch.device('cpu'),\n        scale_laf: float = 1.0,\n        detector='GFTT',\n    ):\n        config = {\n            # Extraction Parameters\n            \"nms_size\": 15,\n            \"pyramid_levels\": 4,\n            \"up_levels\": 1,\n            \"scale_factor_levels\": math.sqrt(2),\n            \"s_mult\": 22.0,\n        }\n        ori_module = KF.PassLAF() if upright else KF.LAFOrienter(angle_detector=KF.OriNet(False)).eval()\n        if not upright:\n            weights = torch.load(os.path.join(MODEL_DIR, \"OriNet.pth\"))[\"state_dict\"]\n            ori_module.angle_detector.load_state_dict(weights)\n        if detector == \"keynet\":\n            detector = KF.KeyNetDetector(\n            False,\n            num_features=num_features,\n            ori_module=ori_module,\n            aff_module=KF.LAFAffNetShapeEstimator(False).eval(),\n            ).to(device)\n            kn_weights = torch.load(os.path.join(MODEL_DIR, \"keynet_pytorch.pth\"))[\n            \"state_dict\"\n            ]\n            detector.model.load_state_dict(kn_weights)\n        elif detector == \"GFTT\":\n            detector = KF.MultiResolutionDetector(\n                KF.CornerGFTT(),\n                num_features=num_features,\n                config=config,\n                ori_module=ori_module,\n                aff_module=KF.LAFAffNetShapeEstimator(False).eval(),\n            ).to(device)\n        elif detector == \"Harris\":\n            detector = KF.MultiResolutionDetector(\n                KF.CornerHarris(0.04),\n                num_features=num_features,\n                config=config,\n                ori_module=ori_module,\n                aff_module=KF.LAFAffNetShapeEstimator(False).eval(),\n            ).to(device)\n        elif detector == \"DoG\":\n            detector = KF.MultiResolutionDetector(\n                KF.BlobDoGSingle(),\n                num_features=num_features,\n                config=config,\n                ori_module=ori_module,\n                aff_module=KF.LAFAffNetShapeEstimator(False).eval(),\n            ).to(device)\n       \n        affnet_weights = torch.load(os.path.join(MODEL_DIR, \"AffNet.pth\"))[\"state_dict\"]\n        detector.aff.load_state_dict(affnet_weights)\n        \n        hardnet8 = KF.HardNet8(False).eval()\n        hn8_weights = torch.load(HARDNET_PT)\n        hardnet8.load_state_dict(hn8_weights)\n        descriptor = KF.LAFDescriptor(hardnet8, patch_size=32, grayscale_descriptor=True).to(device)\n        super().__init__(detector, descriptor, scale_laf)\ndef detect_features(img_fnames,\n                    num_feats = 8000,\n                    upright = False,\n                    device=torch.device('cpu'),\n                    feature_dir = '.featureout',\n                    resize_small_edge_to = 1200,\n                local_feature='DoG'):\n\n    \n    feature = AffNetHardNet(num_feats, upright, device, detector=local_feature).to(device).eval()\n    torch.cuda.empty_cache()\n    gc.collect()\n    if not os.path.isdir(feature_dir):\n        os.makedirs(feature_dir)\n    with h5py.File(f'{feature_dir}/lafs.h5', mode='w') as f_laf, \\\n        h5py.File(f'{feature_dir}/keypoints.h5', mode='w') as f_kp, \\\n        h5py.File(f'{feature_dir}/descriptors.h5', mode='w') as f_desc:\n        for img_path in progress_bar(img_fnames):\n            img_fname = img_path.split('/')[-1]\n            key = img_fname\n            with torch.inference_mode():\n                timg = load_torch_image(img_path, device=device)\n                H, W = timg.shape[2:]\n                if resize_small_edge_to is None:\n                    timg_resized = timg\n                else:\n                    timg_resized = K.geometry.resize(timg, resize_small_edge_to, antialias=True)\n                    torch.cuda.empty_cache()\n                    gc.collect()\n                    print(f'Resized {timg.shape} to {timg_resized.shape} (resize_small_edge_to={resize_small_edge_to})')\n                h, w = timg_resized.shape[2:]\n                \n                lafs, resps, descs = feature(K.color.rgb_to_grayscale(timg_resized))\n                lafs[:,:,0,:] *= float(W) / float(w)\n                lafs[:,:,1,:] *= float(H) / float(h)\n                desc_dim = descs.shape[-1]\n                kpts = KF.get_laf_center(lafs).reshape(-1, 2).detach().cpu().numpy()\n                descs = descs.reshape(-1, desc_dim).detach().cpu().numpy()\n                f_laf[key] = lafs.detach().cpu().numpy()\n                f_kp[key] = kpts\n                f_desc[key] = descs\n                del timg,timg_resized,lafs,kpts,descs\n                torch.cuda.empty_cache()\n                gc.collect()\n            \n            torch.cuda.empty_cache()\n            gc.collect()\n    del feature\n    torch.cuda.empty_cache()\n    gc.collect()\n    return\n\ndef match_features(img_fnames,\n                index_pairs,\n                file_keypoints,\n                feature_dir = '.featureout',\n                device=torch.device('cpu'),\n                min_matches=15, \n                force_mutual = True,\n                matching_alg='adalam'\n                ):\n    assert matching_alg in ['smnn', 'adalam']\n    with h5py.File(f'{feature_dir}/lafs.h5', mode='r') as f_laf, \\\n        h5py.File(f'{feature_dir}/keypoints.h5', mode='r') as f_kp, \\\n        h5py.File(f'{feature_dir}/descriptors.h5', mode='r') as f_desc, \\\n        h5py.File(file_keypoints, mode='w') as f_match:\n\n        for pair_idx in progress_bar(index_pairs):\n                    idx1, idx2 = pair_idx\n                    fname1, fname2 = img_fnames[idx1], img_fnames[idx2]\n                    key1, key2 = fname1.split('/')[-1], fname2.split('/')[-1]\n                    lafs1 = torch.from_numpy(f_laf[key1][...]).to(device)\n                    lafs2 = torch.from_numpy(f_laf[key2][...]).to(device)\n                    desc1 = torch.from_numpy(f_desc[key1][...]).to(device)\n                    desc2 = torch.from_numpy(f_desc[key2][...]).to(device)\n                    kp1 = torch.from_numpy(f_kp[key1][...]).to(device)\n                    kp2 = torch.from_numpy(f_kp[key2][...]).to(device)\n                    if matching_alg == 'adalam':\n                        img1, img2 = cv2.imread(fname1), cv2.imread(fname2)\n                        hw1, hw2 = img1.shape[:2], img2.shape[:2]\n                        adalam_config = KF.adalam.get_adalam_default_config()\n                        adalam_config['force_seed_mnn']= True\n                        adalam_config['search_expansion'] = 16\n                        adalam_config['ransac_iters'] = 128\n                        adalam_config['device'] = device\n                        dists, idxs = KF.match_adalam(desc1, desc2,\n                                                    lafs1, lafs2, # Adalam takes into account also geometric information\n                                                    hw1=hw1, hw2=hw2,\n                                                    config=adalam_config) # Adalam also benefits from knowing image size\n                    else:\n                        dists, idxs = KF.match_smnn(desc1, desc2, 0.98)\n                    if len(idxs)  == 0:\n                        continue\n                    if force_mutual:\n                        first_indices = get_unique_idxs(idxs[:,1])\n                        idxs = idxs[first_indices]\n                        dists = dists[first_indices]\n                    n_matches = len(idxs)\n                    kp1 = kp1[idxs[:,0], :].cpu().numpy().reshape(-1, 2).astype(np.float32)\n                    kp2 = kp2[idxs[:,1], :].cpu().numpy().reshape(-1, 2).astype(np.float32)\n                    if False:\n                        print (f'{key1}-{key2}: {n_matches} matches')\n                    group  = f_match.require_group(key1)\n                    if n_matches >= min_matches:\n                        group.create_dataset(key2, data=np.concatenate([kp1, kp2], axis=1))\n    torch.cuda.empty_cache()\n    gc.collect()\n    return\n\n\ndef get_unique_idxs(A, dim=0):\n    # https://stackoverflow.com/questions/72001505/how-to-get-unique-elements-and-their-firstly-appeared-indices-of-a-pytorch-tenso\n    unique, idx, counts = torch.unique(A, dim=dim, sorted=True, return_inverse=True, return_counts=True)\n    _, ind_sorted = torch.sort(idx, stable=True)\n    cum_sum = counts.cumsum(0)\n    cum_sum = torch.cat((torch.tensor([0],device=cum_sum.device), cum_sum[:-1]))\n    first_indices = ind_sorted[cum_sum]\n    return first_indices\n\ndef get_keypoint_from_h5(fp, key1, key2):\n    rc = -1\n    try:\n        kpts = np.array(fp[key1][key2])\n        rc = 0\n        return (rc, kpts)\n    except:\n        return (rc, None)\n\ndef get_keypoint_from_multi_h5(fps, key1, key2):\n    list_mkpts = []\n    for fp in fps:\n        rc, mkpts = get_keypoint_from_h5(fp, key1, key2)\n        if rc == 0:\n            list_mkpts.append(mkpts)\n    if len(list_mkpts) > 0:\n        list_mkpts = np.concatenate(list_mkpts, axis=0)\n    else:\n        list_mkpts = None\n    return list_mkpts\n\ndef matches_merger(\n    img_fnames,\n    index_pairs,\n    files_keypoints,\n    save_file,\n    feature_dir = 'featureout',\n    filter_FundamentalMatrix = False,\n    filter_iterations = 10,\n    filter_threshold = 8,\n):\n    # open h5 files\n    fps = [ h5py.File(file, mode=\"r\") for file in files_keypoints ]\n\n    with h5py.File(save_file, mode='w') as f_match:\n        counter = 0\n        for pair_idx in progress_bar(index_pairs):\n            idx1, idx2 = pair_idx\n            fname1, fname2 = img_fnames[idx1], img_fnames[idx2]\n            key1, key2 = fname1.split('/')[-1], fname2.split('/')[-1]\n\n            # extract keypoints\n            mkpts = get_keypoint_from_multi_h5(fps, key1, key2)\n            if mkpts is None:\n                print(f\"skipped key1={key1}, key2={key2}\")\n                continue\n\n            ori_size = mkpts.shape[0]\n            if mkpts.shape[0] < CONFIG.MERGE_PARAMS[\"min_matches\"]:\n                continue\n            \n            if filter_FundamentalMatrix:\n                store_inliers = { idx:0 for idx in range(mkpts.shape[0]) }\n                idxs = np.array(range(mkpts.shape[0]))\n                for iter in range(filter_iterations):\n                    try:\n                        Fm, inliers = cv2.findFundamentalMat(\n                            mkpts[:,:2], mkpts[:,2:4], cv2.USAC_MAGSAC, 0.15, 0.9999, 20000)\n                        if Fm is not None:\n                            inliers = inliers > 0\n                            inlier_idxs = idxs[inliers[:, 0]]\n                            for idx in inlier_idxs:\n                                store_inliers[idx] += 1\n                    except:\n                        print(f\"Failed to cv2.findFundamentalMat. mkpts.shape={mkpts.shape}\")\n                inliers = np.array([ count for (idx, count) in store_inliers.items() ]) >= filter_threshold\n                mkpts = mkpts[inliers]\n                if mkpts.shape[0] < 15:\n                    print(f\"skipped key1={key1}, key2={key2}: mkpts.shape={mkpts.shape} after filtered.\")\n                    continue\n            \n            \n            print (f'{key1}-{key2}: {ori_size} --> {mkpts.shape[0]} matches')            \n            # regist tmp file\n            group  = f_match.require_group(key1)\n            group.create_dataset(key2, data=mkpts)\n            counter += 1\n    print( f\"Ensembled pairs : {counter} pairs\" )\n    for fp in fps:\n        fp.close()\n\ndef keypoints_merger(\n    img_fnames,\n    index_pairs,\n    files_keypoints,\n    feature_dir = 'featureout',\n    filter_FundamentalMatrix = False,\n    filter_iterations = 10,\n    filter_threshold = 8,\n):\n    save_file = f'{feature_dir}/merge_tmp.h5'\n    !rm -rf {save_file}\n    matches_merger(\n        img_fnames,\n        index_pairs,\n        files_keypoints,\n        save_file,\n        feature_dir = feature_dir,\n        filter_FundamentalMatrix = filter_FundamentalMatrix,\n        filter_iterations = filter_iterations,\n        filter_threshold = filter_threshold,\n    )\n        \n    # Let's find unique loftr pixels and group them together.\n    kpts = defaultdict(list)\n    match_indexes = defaultdict(dict)\n    total_kpts=defaultdict(int)\n    with h5py.File(save_file, mode='r') as f_match:\n        for k1 in f_match.keys():\n            group  = f_match[k1]\n            for k2 in group.keys():\n                matches = group[k2][...]\n                total_kpts[k1]\n                kpts[k1].append(matches[:, :2])\n                kpts[k2].append(matches[:, 2:])\n                current_match = torch.arange(len(matches)).reshape(-1, 1).repeat(1, 2)\n                current_match[:, 0]+=total_kpts[k1]\n                current_match[:, 1]+=total_kpts[k2]\n                total_kpts[k1]+=len(matches)\n                total_kpts[k2]+=len(matches)\n                match_indexes[k1][k2]=current_match\n\n    for k in kpts.keys():\n        kpts[k] = np.round(np.concatenate(kpts[k], axis=0))\n    unique_kpts = {}\n    unique_match_idxs = {}\n    out_match = defaultdict(dict)\n    for k in kpts.keys():\n        uniq_kps, uniq_reverse_idxs = torch.unique(torch.from_numpy(kpts[k]),dim=0, return_inverse=True)\n        unique_match_idxs[k] = uniq_reverse_idxs\n        unique_kpts[k] = uniq_kps.numpy()\n    for k1, group in match_indexes.items():\n        for k2, m in group.items():\n            m2 = deepcopy(m)\n            m2[:,0] = unique_match_idxs[k1][m2[:,0]]\n            m2[:,1] = unique_match_idxs[k2][m2[:,1]]\n            mkpts = np.concatenate([unique_kpts[k1][ m2[:,0]],\n                                    unique_kpts[k2][  m2[:,1]],\n                                ],\n                                axis=1)\n            unique_idxs_current = get_unique_idxs(torch.from_numpy(mkpts), dim=0)\n            m2_semiclean = m2[unique_idxs_current]\n            unique_idxs_current1 = get_unique_idxs(m2_semiclean[:, 0], dim=0)\n            m2_semiclean = m2_semiclean[unique_idxs_current1]\n            unique_idxs_current2 = get_unique_idxs(m2_semiclean[:, 1], dim=0)\n            m2_semiclean2 = m2_semiclean[unique_idxs_current2]\n            out_match[k1][k2] = m2_semiclean2.numpy()\n    with h5py.File(f'{feature_dir}/keypoints.h5', mode='w') as f_kp:\n        for k, kpts1 in unique_kpts.items():\n            f_kp[k] = kpts1\n    \n    with h5py.File(f'{feature_dir}/matches.h5', mode='w') as f_match:\n        for k1, gr in out_match.items():\n            group  = f_match.require_group(k1)\n            for k2, match in gr.items():\n                group[k2] = match\n    return\n\ndef wrapper_keypoints(\n    img_fnames, index_pairs, feature_dir, device, timings, rots\n):\n    #############################################################\n    # get keypoints\n    #############################################################\n    files_keypoints = []\n    \n    local_feature_model = 'DoG' \n    detect_features(img_fnames, \n                    16000,\n                    feature_dir=local_feature_model,\n                    upright=False,\n                    device=device,\n                    resize_small_edge_to=1000,\n                    local_feature=local_feature_model,\n                            )\n    torch.cuda.empty_cache()\n    gc.collect()\n    file_keypoints = f\"{feature_dir}/matches_{local_feature_model}.h5\"\n    match_features(img_fnames, index_pairs, file_keypoints, feature_dir=local_feature_model,device=device)\n    files_keypoints.append( file_keypoints )\n    torch.cuda.empty_cache()\n    gc.collect()\n\n    if CONFIG.use_aliked_lightglue:\n        model_name = \"aliked\"\n        file_keypoints = f'{feature_dir}/matches_lightglue_{model_name}.h5'\n        t = detect_lightglue_common(\n            img_fnames, model_name, index_pairs, feature_dir, device, file_keypoints, rots,\n            resize_to=CONFIG.params_aliked_lightglue[\"resize_to\"],\n            detection_threshold=CONFIG.params_aliked_lightglue[\"detection_threshold\"],\n            num_features=CONFIG.params_aliked_lightglue[\"num_features\"],\n            min_matches=CONFIG.params_aliked_lightglue[\"min_matches\"],\n        )\n        gc.collect()\n        files_keypoints.append(file_keypoints)\n        timings['feature_matching'].append(t)\n\n\n    keypoints_merger(\n        img_fnames,\n        index_pairs,\n        files_keypoints,\n        feature_dir = feature_dir,\n        filter_FundamentalMatrix = CONFIG.MERGE_PARAMS[\"filter_FundamentalMatrix\"],\n        filter_iterations = CONFIG.MERGE_PARAMS[\"filter_iterations\"],\n        filter_threshold = CONFIG.MERGE_PARAMS[\"filter_threshold\"],\n    )    \n    return timings\n\ndef reconstruct_from_db(dataset, scene, feature_dir, img_dir, timings, image_paths):\n    scene_result = {}\n    #############################################################\n    # regist keypoints from h5 into colmap db\n    #############################################################\n    database_path = f'{feature_dir}/colmap.db'\n    if os.path.isfile(database_path):\n        os.remove(database_path)\n    gc.collect()\n    import_into_colmap(img_dir, feature_dir=feature_dir, database_path=database_path)\n    output_path = f'{feature_dir}/colmap_rec'\n\n    #############################################################\n    # Calculate fundamental matrix with colmap api\n    #############################################################\n    t=time.time()\n    options = pycolmap.SiftMatchingOptions()\n    pycolmap.match_exhaustive(database_path, sift_options=options)\n    t=time.time() - t \n    timings['RANSAC'].append(t)\n    print(f'RANSAC in  {t:.4f} sec')\n\n    #############################################################\n    # Execute bundle adjustmnet with colmap api\n    # --> Bundle adjustment Calcs Camera matrix, R and t\n    #############################################################\n    t=time.time()\n    # By default colmap does not generate a reconstruction if less than 10 images are registered. Lower it to 3.\n    mapper_options = pycolmap.IncrementalMapperOptions()\n    mapper_options.min_model_size = 3\n    os.makedirs(output_path, exist_ok=True)\n    maps = pycolmap.incremental_mapping(database_path=database_path, image_path=img_dir, output_path=output_path, options=mapper_options)\n    torch.cuda.empty_cache()\n    gc.collect()\n    clear_output(wait=False)\n    t=time.time() - t\n    timings['Reconstruction'].append(t)\n    print(f'Reconstruction done in  {t:.4f} sec')\n\n    #############################################################\n    # Extract R,t from maps \n    #############################################################            \n    imgs_registered  = 0\n    best_idx = None\n    list_num_images = []            \n    print (\"Looking for the best reconstruction\")\n    if isinstance(maps, dict):\n        for idx1, rec in maps.items():\n            print (idx1, rec.summary())\n            list_num_images.append( len(rec.images) )\n            if len(rec.images) > imgs_registered:\n                imgs_registered = len(rec.images)\n                best_idx = idx1\n    list_num_images = np.array(list_num_images)\n    print(f\"list_num_images = {list_num_images}\")\n    if best_idx is not None:\n        print (maps[best_idx].summary())\n        for k, im in maps[best_idx].images.items():\n            #key1 = f'test/{dataset}/images/{im.name}'\n            key1 = f'train/{dataset}/images/{im.name}'\n            scene_result[key1] = {}\n            scene_result[key1][\"R\"] = deepcopy(im.rotmat())\n            scene_result[key1][\"t\"] = deepcopy(np.array(im.tvec))\n            torch.cuda.empty_cache()\n            gc.collect()\n\n    print(f'Registered: {dataset} / {scene} -> {len(scene_result)} images')\n    print(f'Total: {dataset} / {scene} -> {len(image_paths)} images')\n    print(timings)\n    torch.cuda.empty_cache()\n    gc.collect()\n    return scene_result\n\ndef arr_to_str(a):\n    return ';'.join([str(x) for x in a.reshape(-1)])\n\n# Function to create a submission file.\ndef create_submission(out_results, data_dict):\n    with open(f'submission.csv', 'w') as f:\n        f.write('image_path,dataset,scene,rotation_matrix,translation_vector\\n')\n        for dataset in data_dict:\n            if dataset in out_results:\n                res = out_results[dataset]\n            else:\n                res = {}\n            for scene in data_dict[dataset]:\n                if scene in res:\n                    scene_res = res[scene]\n                else:\n                    scene_res = {\"R\":{}, \"t\":{}}\n                for image in data_dict[dataset][scene]:\n                    if image in scene_res:\n                        R = scene_res[image]['R'].reshape(-1)\n                        T = scene_res[image]['t'].reshape(-1)\n                    else:\n                        R = np.eye(3).reshape(-1)\n                        T = np.zeros((3))\n                    f.write(f'{image},{dataset},{scene},{arr_to_str(R)},{arr_to_str(T)}\\n')\n\nsrc = '/kaggle/input/image-matching-challenge-2024'\n\n# Get data from csv.\ndata_dict = {}\nwith open(f'{src}/sample_submission.csv', 'r') as f:\n    for i, l in enumerate(f):\n        # Skip header.\n        if l and i > 0:\n            image, dataset, scene, _, _ = l.strip().split(',')\n            if dataset not in data_dict:\n                data_dict[dataset] = {}\n            if scene not in data_dict[dataset]:\n                data_dict[dataset][scene] = []\n            data_dict[dataset][scene].append(image)\n            \n            if CONFIG.DRY_RUN:\n                if len(data_dict[dataset][scene]) == CONFIG.DRY_RUN_MAX_IMAGES:\n                    break\n                    \nfor dataset in data_dict:\n    for scene in data_dict[dataset]:\n        print(f'{dataset} / {scene} -> {len(data_dict[dataset][scene])} images')\n\nout_results = {}\ntimings = {\n    \"rotation_detection\" : [],\n    \"shortlisting\":[],\n    \"feature_detection\": [],\n    \"feature_matching\":[],\n    \"RANSAC\": [],\n    \"Reconstruction\": []\n}\n\ngc.collect()\ndatasets = []\nfor dataset in data_dict:\n    datasets.append(dataset)\n\nwith concurrent.futures.ProcessPoolExecutor(max_workers=CONFIG.NUM_CORES) as executors:\n    futures = defaultdict(dict)\n    for dataset in datasets:\n        if dataset not in out_results:\n            out_results[dataset] = {}\n        for scene in data_dict[dataset]:\n            # Fail gently if the notebook has not been submitted and the test data is not populated.\n            # You may want to run this on the training data in that case?\n            img_dir = f'{src}/train/{dataset}/images'\n            if not os.path.exists(img_dir):\n                continue\n\n            out_results[dataset][scene] = {}\n            img_fnames = [f'{src}/{x}' for x in data_dict[dataset][scene]]\n            print (f\"Got {len(img_fnames)} images\")\n            feature_dir = f'featureout/{dataset}_{scene}'\n            if not os.path.isdir(feature_dir):\n                os.makedirs(feature_dir, exist_ok=True)\n            t = time.time()\n            rots = [ 0 for fname in img_fnames ]\n            t = time.time()-t\n            timings['rotation_detection'].append(t)\n            print (f'rotation_detection for {len(img_fnames)} images : {t:.4f} sec')\n            gc.collect()\n            t=time.time()\n\n            index_pairs = get_image_pairs_shortlist(img_fnames,\n                                                sim_th = 0.20,# should be strict\n                                                min_pairs = 40, # we select at least min_pairs PER IMAGE with biggest similarity\n                                                exhaustive_if_less = 40,\n                                                device=torch.device('cpu')   \n                                                )\n            t=time.time() -t \n            timings['shortlisting'].append(t)\n            print (f'{len(index_pairs)}, pairs to match, {t:.4f} sec')\n            torch.cuda.empty_cache()\n            gc.collect()\n\n            #############################################################\n            # get keypoints\n            #############################################################            \n            keypoints_timings = wrapper_keypoints(\n                img_fnames, index_pairs, feature_dir, device, timings, rots\n            )\n            torch.cuda.empty_cache()\n            gc.collect()\n            timings['feature_matching'] = keypoints_timings['feature_matching']\n\n            #############################################################\n            # kick COLMAP reconstruction\n            #############################################################            \n            futures[dataset][scene] = executors.submit(\n                reconstruct_from_db, \n                dataset, scene, feature_dir, img_dir, timings, data_dict[dataset][scene])\n                \n    #############################################################\n    # reconstruction results\n    #############################################################            \n    for dataset in datasets:\n        for scene in data_dict[dataset]:\n            # wait to complete COLMAP reconstruction\n            result = futures[dataset][scene].result()\n            if result is not None:\n                out_results[dataset][scene] = result   # get R and t from result\n            torch.cuda.empty_cache()\n            gc.collect()\n    create_submission(out_results, data_dict)\n    torch.cuda.empty_cache()\n    gc.collect()\n ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pandas as pd\ndef register_by_Horn(ev_coord, gt_coord, ransac_threshold, inl_cf, strict_cf):\n    \n    # remove invalid cameras, the index is returned\n    idx_cams = np.all(np.isfinite(ev_coord), axis=0)\n    ev_coord = ev_coord[:, idx_cams]\n    gt_coord = gt_coord[:, idx_cams]\n\n    # initialization\n    n = ev_coord.shape[1]\n    r = ransac_threshold.shape[0]\n    ransac_threshold = np.expand_dims(ransac_threshold, axis=0)\n    ransac_threshold2 = ransac_threshold**2\n    ev_coord_1 = np.vstack((ev_coord, np.ones(n)))\n\n    max_no_inl = np.zeros((1, r))\n    best_inl_err = np.full(r, np.Inf)\n    best_transf_matrix = np.zeros((r, 4, 4))\n    best_err = np.full((n, r), np.Inf)\n    strict_inl = np.full((n, r), False)\n    triplets_used = np.zeros((3, r))\n\n    # run on camera triplets\n    for ii in range(n-2):\n        for jj in range(ii+1, n-1):\n            for kk in range(jj+1, n):\n                i = [ii, jj, kk]\n                triplets_used_now = np.full((n), False)\n                triplets_used_now[i] = True\n                # if both ii, jj, kk are strict inliers for the best current model just skip\n                if np.all(strict_inl[i]):\n                    continue\n                # get transformation T by Horn on the triplet camera center correspondences\n                transf_matrix = affine_matrix_from_points(ev_coord[:, i], gt_coord[:, i], usesvd=False)\n                # apply transformation T to test camera centres\n                rotranslated = np.matmul(transf_matrix[:3], ev_coord_1)\n                # compute error and inliers\n                err = np.sum((rotranslated - gt_coord)**2, axis=0)\n                inl = np.expand_dims(err, axis=1) < ransac_threshold2\n                no_inl = np.sum(inl, axis=0)\n                # if the number of inliers is close to that of the best model so far, go for refinement\n                to_ref = np.squeeze(((no_inl > 2) & (no_inl > max_no_inl * inl_cf)), axis=0)\n                for q in np.argwhere(to_ref):                        \n                    qq = q[0]\n                    if np.any(np.all((np.expand_dims(inl[:, qq], axis=1) == inl[:, :qq]), axis=0)):\n                        # already done for this set of inliers\n                        continue\n                    # get transformation T by Horn on the inlier camera center correspondences\n                    transf_matrix = affine_matrix_from_points(ev_coord[:, inl[:, qq]], gt_coord[:, inl[:, qq]])\n                    # apply transformation T to test camera centres\n                    rotranslated = np.matmul(transf_matrix[:3], ev_coord_1)\n                    # compute error and inliers\n                    err_ref = np.sum((rotranslated - gt_coord)**2, axis=0)\n                    err_ref_sum = np.sum(err_ref, axis=0)\n                    err_ref = np.expand_dims(err_ref, axis=1)\n                    inl_ref = err_ref < ransac_threshold2\n                    no_inl_ref = np.sum(inl_ref, axis=0)\n                    # update the model if better for each threshold\n                    to_update = np.squeeze((no_inl_ref > max_no_inl) | ((no_inl_ref == max_no_inl) & (err_ref_sum < best_inl_err)), axis=0)\n                    if np.any(to_update):\n                        triplets_used[0, to_update] = ii\n                        triplets_used[1, to_update] = jj\n                        triplets_used[2, to_update] = kk\n                        max_no_inl[:, to_update] = no_inl_ref[to_update]\n                        best_err[:, to_update] = np.sqrt(err_ref)\n                        best_inl_err[to_update] = err_ref_sum\n                        strict_inl[:, to_update] = (best_err[:, to_update] < strict_cf * ransac_threshold[:, to_update])\n                        best_transf_matrix[to_update] = transf_matrix\n\n#     for i in range(r):\n#        print(f'Registered cameras {int(max_no_inl[0, i])}/{n} for threshold {ransac_threshold[0, i]}')\n\n    best_model = {\n        \"valid_cams\": idx_cams,        \n        \"no_inl\": max_no_inl,\n        \"err\": best_err,\n        \"triplets_used\": triplets_used,\n        \"transf_matrix\": best_transf_matrix}\n    return best_model\ndef mAA_on_cameras(err, thresholds, n, skip_top_thresholds, to_dec=3):\n    \n    aux = err[:, skip_top_thresholds:] < np.expand_dims(np.asarray(thresholds[skip_top_thresholds:]), axis=0)\n    return np.sum(np.maximum(np.sum(aux, axis=0) - to_dec, 0)) / (len(thresholds[skip_top_thresholds:]) * (n - to_dec))\ndef get_camera_centers_from_df(df):\n    out = {}\n    for row in df.iterrows():\n        row = row[1]\n        fname = row['image_path']\n        R = np.array([float(x) for x in (row['rotation_matrix'].split(';'))]).reshape(3,3)\n        t = np.array([float(x) for x in (row['translation_vector'].split(';'))]).reshape(3)\n        center = -R.T @ t\n        out[fname] = center\n    return out\ntranslation_thresholds_meters_dict = {\n 'multi-temporal-temple-baalshamin':  np.array([0.025,  0.05,  0.1,  0.2,  0.5,  1.0]),\n 'pond':                              np.array([0.025,  0.05,  0.1,  0.2,  0.5,  1.0]),\n 'transp_obj_glass_cylinder':         np.array([0.0025, 0.005, 0.01, 0.02, 0.05, 0.1]),\n 'transp_obj_glass_cup':              np.array([0.0025, 0.005, 0.01, 0.02, 0.05, 0.1]),\n 'church':                            np.array([0.025,  0.05,  0.1,  0.2,  0.5,  1.0]),\n 'lizard':                            np.array([0.025,  0.05,  0.1,  0.2,  0.5,  1.0]),\n 'dioscuri':                          np.array([0.025,  0.05,  0.1,  0.2,  0.5,  1.0]), \n}\ndef evaluate_rec(gt_df, user_df, inl_cf = 0.8, strict_cf=0.5, skip_top_thresholds=2, to_dec=3,\n                 thresholds=[0.005, 0.01, 0.02, 0.03, 0.04, 0.05, 0.1, 0.15, 0.2]):\n    # get camera centers\n    ucameras = get_camera_centers_from_df(user_df)\n    gcameras = get_camera_centers_from_df(gt_df)    \n\n    # the denominator for mAA ratio\n    m = gt_df.shape[0]\n    \n    # get the image list to use\n    good_cams = []\n    for image_path in gcameras.keys():\n        if image_path in ucameras.keys():\n            good_cams.append(image_path)\n        \n    # put corresponding camera centers into matrices\n    n = len(good_cams)\n    u_cameras = np.zeros((3, n))\n    g_cameras = np.zeros((3, n))\n    \n    ii = 0\n    for i in good_cams:\n        u_cameras[:, ii] = ucameras[i]\n        g_cameras[:, ii] = gcameras[i]\n        ii += 1\n        \n    # Horn camera centers registration, a different best model for each camera threshold\n    model = register_by_Horn(u_cameras, g_cameras, np.asarray(thresholds), inl_cf, strict_cf)\n    \n    # transformation matrix\n#     print(\"\\nTransformation matrix for maximum threshold\")\n    T = np.squeeze(model['transf_matrix'][-1])\n#     print(T)\n    \n    # mAA\n    mAA = mAA_on_cameras(model[\"err\"], thresholds, m, skip_top_thresholds, to_dec)\n    # print(f'mAA = {mAA * 100 : .2f}% considering {m} input cameras - {to_dec}')\n    return mAA\n\ndef score(solution: pd.DataFrame, submission: pd.DataFrame) -> float:\n    \n    scenes = list(set(solution['dataset'].tolist()))\n    results_per_dataset = []\n    for dataset in scenes:\n        print(f\"\\n*** {dataset} ***\")\n#         start = time.time()\n        gt_ds = solution[solution['dataset'] == dataset]\n        user_ds = submission[submission['dataset'] == dataset]\n        gt_ds = gt_ds.sort_values(by=['image_path'], ascending = True)\n        user_ds = user_ds.sort_values(by=['image_path'], ascending = True)\n        result = evaluate_rec(gt_ds, user_ds, inl_cf=0, strict_cf=-1, skip_top_thresholds=0, to_dec=3,\n                 thresholds=translation_thresholds_meters_dict[dataset])\n#         end = time.time()\n        print(f\"\\nmAA: {round(result,4)}\")\n#         print(\"Running time: %s\" % (end - start))        \n        results_per_dataset.append(result)\n    return float(np.array(results_per_dataset).mean())","metadata":{"execution":{"iopub.status.busy":"2024-06-26T16:26:22.432395Z","iopub.execute_input":"2024-06-26T16:26:22.432775Z","iopub.status.idle":"2024-06-26T16:26:22.472669Z","shell.execute_reply.started":"2024-06-26T16:26:22.432745Z","shell.execute_reply":"2024-06-26T16:26:22.471163Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#from imc24 import score\ndef image_path(row):\n        row['image_path'] = 'train/' + row['dataset'] + '/images/' + row['image_name']\n        return row\n\ntrain_df = pd.read_csv(f'/kaggle/input/image-matching-challenge-2024/train/train_labels.csv')\ntrain_df = train_df.apply(image_path,axis=1).drop_duplicates(subset=['image_path'])\nG = train_df.groupby(['dataset','scene'])['image_path']\nimage_paths = []\n    \nfor g in G:\n    n = 40\n    n = n if n < len(g[1]) else len(g[1])\n    g = g[0],g[1].sample(n,random_state=42).reset_index(drop=True)\n    for image_path in g[1]:\n        image_paths.append(image_path)\n\ngt_df = train_df[train_df.image_path.isin(image_paths)].reset_index(drop=True)\npred_df = gt_df[['image_path','dataset','scene','rotation_matrix','translation_vector']]\npred_df.to_csv('pred_df.csv',index=False)\npred_df = pd.read_csv('submission.csv')\nmAA = round(score(gt_df, pred_df),4)\nprint('*** Total mean Average Accuracy ***')\nprint(f\"mAA: {mAA}\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!ls","metadata":{"execution":{"iopub.status.busy":"2024-06-26T16:13:28.564918Z","iopub.status.idle":"2024-06-26T16:13:28.565288Z","shell.execute_reply.started":"2024-06-26T16:13:28.565109Z","shell.execute_reply":"2024-06-26T16:13:28.565124Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Result","metadata":{}},{"cell_type":"code","source":"!cat submission.csv","metadata":{"execution":{"iopub.status.busy":"2024-06-26T16:13:28.5678Z","iopub.status.idle":"2024-06-26T16:13:28.568402Z","shell.execute_reply.started":"2024-06-26T16:13:28.568081Z","shell.execute_reply":"2024-06-26T16:13:28.568102Z"},"trusted":true},"execution_count":null,"outputs":[]}]}