{"metadata":{"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":71885,"databundleVersionId":8143495,"sourceType":"competition"},{"sourceId":5373920,"sourceType":"datasetVersion","datasetId":3117886},{"sourceId":5850511,"sourceType":"datasetVersion","datasetId":3364321},{"sourceId":7884485,"sourceType":"datasetVersion","datasetId":4628051},{"sourceId":8413202,"sourceType":"datasetVersion","datasetId":4980819},{"sourceId":177657330,"sourceType":"kernelVersion"}],"dockerImageVersionId":30699,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true},"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"papermill":{"default_parameters":{},"duration":897.122094,"end_time":"2024-04-25T14:22:05.177547","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2024-04-25T14:07:08.055453","version":"2.4.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"https://www.kaggle.com/code/eduardtrulls/imc-2023-submission-example","metadata":{"papermill":{"duration":0.010536,"end_time":"2024-04-25T14:07:18.584727","exception":false,"start_time":"2024-04-25T14:07:18.574191","status":"completed"},"tags":[]}},{"cell_type":"code","source":"import imc24_kornia","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# General utilities\nimport os\nfrom tqdm import tqdm\nfrom time import time\nfrom fastprogress import progress_bar\nimport gc\nimport math\nimport numpy as np\nfrom IPython.display import clear_output\nfrom collections import defaultdict\nfrom copy import deepcopy\nimport matplotlib.pyplot as plt\nimport concurrent.futures\nfrom IPython.display import clear_output\n\n# CV/ML\nimport cv2\nimport torch\nimport torch.nn.functional as F\nimport kornia as K\nimport kornia.feature as KF\nfrom lightglue import ALIKED\nfrom PIL import Image\nimport timm\nfrom timm.data import resolve_data_config\nfrom timm.data.transforms_factory import create_transform\n\n# 3D reconstruction\nimport pycolmap\n\nprint(\"Kornia version\", K.__version__)\nprint(\"Pycolmap version\", pycolmap.__version__)","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","papermill":{"duration":13.183171,"end_time":"2024-04-25T14:07:31.777821","exception":false,"start_time":"2024-04-25T14:07:18.594650","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Global Configs","metadata":{"papermill":{"duration":0.021691,"end_time":"2024-04-25T14:07:31.821834","exception":false,"start_time":"2024-04-25T14:07:31.800143","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Mode can only be train or test. This will be used to find the image directory.\n# Use \"test\" for submission \n# MODE = \"train\"\nMODE = \"test\"\n\n# Option to change path for local testing\n# is_local = True\nis_local = False\n\nif is_local:\n    NUM_CORES = 2\n    SRC = \"./kaggle/input/image-matching-challenge-2024\"\n    MODEL_DIR = \"./kaggle/input/kornia-local-feature-weights/\"\n    DISK_PATH = \"./loftr_disk.ckpt\"\n    HARDNET_PT = \"./kaggle/input/kornia-local-feature-weights/hardnet8v2.pt\"\nelse:\n    NUM_CORES = 2\n    SRC = \"/kaggle/input/image-matching-challenge-2024\"\n    MODEL_DIR = \"/kaggle/input/kornia-local-feature-weights/\"\n    DISK_PATH = \"/kaggle/input/disk/pytorch/depth-supervision/1/loftr_outdoor.ckpt\"\n    HARDNET_PT = \"/kaggle/input/hardnet8v2/hardnet8v2.pt\"\n\nLOG_MESSAGE = \"Final submission\"\nMATCHES_CAP = None\n\n# DEBUG = True\nDEBUG = False\n\n# DEBUG_SCENE = [\"cyprus\", \"kyiv-puppet-theater\"]\n# DEBUG_SCENE = [\"cyprus\"]\n# DEBUG_SCENE = [\"kyiv-puppet-theater\"]\n# DEBUG_SCENE = [\"kyiv-puppet-theater\", \"cyprus\", \"wall\", \"chairs\"]\nDEBUG_SCENE = [\"chairs\"]\n# DEBUG_SCENE = [\"wall\"]\n\n# Longer edge limit of the input image\nhardnet_res = 1200\n\nMODEL_DICT = {\n    \"Keynet\": {\"enable\": True, \"resize_long_edge_to\": hardnet_res, \"pair_only\": False},\n    \"GFTT\": {\"enable\": True, \"resize_long_edge_to\": hardnet_res},\n    \"DoG\": {\"enable\": True, \"resize_long_edge_to\": hardnet_res},\n    \"Harris\": {\"enable\": True, \"resize_long_edge_to\": hardnet_res},\n    \"disk_lightglue\": {\"enable\": False, \"resize_long_edge_to\": hardnet_res},\n    \"aliked_lightglue\": {\"enable\": True, \"resize_long_edge_to\": hardnet_res}\n}\n\n# Find fundamental matrix parameters\nFM_PARAMS = {\"ransacReprojThreshold\": 1.4, \"confidence\": 0.9999, \"maxIters\": 50000, \"removeOutliers\": True}\n\n# Remove a \"match\" if the number of matches is lower than MATCH_FILTER_RATIO*max_num_matches\n# e.g. img1 and img2 have max 10000 matches with some other images, img2 and img1 only have 99 matches. The matches btw img1 and img2 won't be selected.\nMATCH_FILTER_RATIO = 0.01\n\n# for logging\nLOG_DICT = dict()\nLOG_DICT[\"mode\"] = MODE\nLOG_DICT[\"log_message\"] = LOG_MESSAGE\nLOG_DICT[\"matches_cap\"] = MATCHES_CAP\nLOG_DICT[\"debug\"] = DEBUG\nLOG_DICT[\"debug_scene\"] = DEBUG_SCENE\n\nif MODE == \"test\":\n    DEBUG = False\ndevice = torch.device(\"cuda\")\nprint(torch.cuda.is_available())\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Get datadict from submission file","metadata":{"papermill":{"duration":0.009659,"end_time":"2024-04-25T14:07:31.922129","exception":false,"start_time":"2024-04-25T14:07:31.912470","status":"completed"},"tags":[]}},{"cell_type":"code","source":"def list_directories(path):\n    entries = os.listdir(path)\n    directories = [entry for entry in entries if os.path.isdir(os.path.join(path, entry))]\n    return directories\n\ndef list_files(directory):\n    file_paths = []\n    for entry in os.listdir(directory):\n        entry_path = os.path.join(directory, entry)\n        if os.path.isfile(entry_path):\n            file_paths.append('/'.join(entry_path.split('/')[-4:]))\n    return file_paths\n\nIMC_PATH = f'/kaggle/input/image-matching-challenge-2024/{MODE}'\ndatasets = list_directories(IMC_PATH)\ndata_dict = {}\nall_scenes = []\nscene_len = []\n\nfor dataset in datasets:\n    data_dict[dataset] = {}\n    data_dict[dataset][dataset] = list_files(f'{IMC_PATH}/{dataset}/images')\n    all_scenes.append((dataset, dataset))\n    scene_len.append(len(data_dict[dataset][dataset]))\n    print(f\"{dataset} / {dataset} -> {len(data_dict[dataset][dataset])} images\")\n\nall_scenes = [x for _, x in sorted(zip(scene_len, all_scenes), reverse=True)]\nprint(\"\\nReconstruction order: \")\nfor scene in all_scenes:\n    print(f\" --{scene[0]} / {scene[1]}\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Submission Utils","metadata":{"papermill":{"duration":0.009936,"end_time":"2024-04-25T14:07:31.983533","exception":false,"start_time":"2024-04-25T14:07:31.973597","status":"completed"},"tags":[]}},{"cell_type":"code","source":"def arr_to_str(a):\n    return \";\".join([str(x) for x in a.reshape(-1)])\n\n\n# Function to create a submission file.\ndef create_submission(out_results1, out_results2, data_dict, mode=\"test\"):\n    if mode == \"train\":\n        file_name = \"submission_train.csv\"\n    else:\n        file_name = \"submission.csv\"\n\n    with open(file_name, \"w\") as f:\n        f.write(\"image_path,dataset,scene,rotation_matrix,translation_vector\\n\")\n        out_results = {}\n        for dataset in data_dict:\n            out_results[dataset] = {}\n            out_results[dataset][dataset] = out_results1[dataset][dataset]\n#             if out_results2[dataset][dataset]['imgs_registered'] > out_results1[dataset][dataset]['imgs_registered']:\n#                 out_results[dataset][dataset] = out_results2[dataset][dataset]\n            \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                        print(image)\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(\n                        f\"{image},{dataset},{scene},{arr_to_str(R)},{arr_to_str(T)}\\n\"\n                    )","metadata":{"papermill":{"duration":0.022128,"end_time":"2024-04-25T14:07:32.015713","exception":false,"start_time":"2024-04-25T14:07:31.993585","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Image Loading and Resize","metadata":{"papermill":{"duration":0.009759,"end_time":"2024-04-25T14:07:32.035430","exception":false,"start_time":"2024-04-25T14:07:32.025671","status":"completed"},"tags":[]}},{"cell_type":"code","source":"def load_torch_image(fname, device=torch.device(\"cpu\")):\n    img = K.image_to_tensor(cv2.imread(fname), False).float() / 255.0\n    img = K.color.bgr_to_rgb(img.to(device))\n    if fname.split('/')[-3] == 'transp_obj_glass_cup':\n        img = torch.rot90(img,dims=[2,3])\n    return img\n\n\ndef resize_torch_image(\n    timg, resize_long_edge_to=None, align=None, disable_enlarge=True\n):\n    h, w = timg.shape[2:]\n    raw_size = torch.tensor(timg.shape[2:])\n    if resize_long_edge_to is None:\n        scale = 1\n    else:\n        scale = float(resize_long_edge_to) / float(max(raw_size[0], raw_size[1]))\n\n    if disable_enlarge:\n        scale = min(scale, 1)\n\n    h_resized = int(h * scale)\n    w_resized = int(w * scale)\n\n    if align is not None:\n        assert align > 0\n        h_resized = h_resized - h_resized % align\n        w_resized = w_resized - w_resized % align\n    scale_h = h_resized / h\n    scale_w = w_resized / w\n\n    timg_resized = K.geometry.resize(timg, (h_resized, w_resized), antialias = True)\n    return timg_resized, scale_h, scale_w\n\n\ndef get_roi_image(timg, roi):\n    min_h = int(roi[\"roi_min_h\"])\n    min_w = int(roi[\"roi_min_w\"])\n    max_h = int(roi[\"roi_max_h\"])\n    max_w = int(roi[\"roi_max_w\"])\n    roi_img = timg[:, :, min_h:max_h, min_w:max_w]\n    roi_w_scale = (max_w - min_w) / timg.shape[3]\n    roi_h_scale = (max_h - min_h) / timg.shape[2]\n    return roi_img, min_h, min_w","metadata":{"papermill":{"duration":0.023172,"end_time":"2024-04-25T14:07:32.068795","exception":false,"start_time":"2024-04-25T14:07:32.045623","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Visualization Utils","metadata":{"papermill":{"duration":0.009933,"end_time":"2024-04-25T14:07:32.089258","exception":false,"start_time":"2024-04-25T14:07:32.079325","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Visualzation block\ndef draw_keypoints(img, keypoints, color=(0, 255, 0)):\n    max_edge = max(img.shape[0], img.shape[1])\n    good_radius = 4\n    for kp in keypoints:\n        x, y = kp\n        cv2.circle(img, (int(x), int(y)), color=color, radius=good_radius, thickness=-1)\n\n\ndef draw_roi(img, roi, color=(0, 255, 255)):\n    x1, y1, x2, y2 = (\n        roi[\"roi_min_w\"],\n        roi[\"roi_min_h\"],\n        roi[\"roi_max_w\"],\n        roi[\"roi_max_h\"],\n    )\n    cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), color=color, thickness=2)\n\n\ndef plot_images_with_keypoints(fname1, fname2, kpts1, kpts2, matches, rois=None):\n    print(fname1, fname2)\n    # Draw keypoints on the images\n    image1 = cv2.imread(fname1)\n    image2 = cv2.imread(fname2)\n    print(image1.shape, image2.shape)\n\n    draw_keypoints(image1, kpts1)\n    draw_keypoints(image2, kpts2)\n    if rois is not None:\n        draw_roi(image1, rois[0])\n        draw_roi(image2, rois[1])\n    print(image1.shape, image2.shape)\n    print(\"Number of matches:\", len(matches))\n    print(\"Number of keypoints:\", len(kpts1), len(kpts2))\n    #print the first match\n    print(matches[0])\n    # Resize image1 and image2 to have the same smaller height\n    display_h = 840\n    h1, w1 = image1.shape[:2]\n    h2, w2 = image2.shape[:2]\n    # new_h = min(h1, h2)\n    scale1 = display_h / h1\n    scale2 = display_h / h2\n    new_w1 = int(w1 * scale1)\n    new_w2 = int(w2 * scale2)\n\n    image1 = cv2.resize(image1, (new_w1, display_h))\n    image2 = cv2.resize(image2, (new_w2, display_h))\n\n    # Create a new image by horizontally concatenating the two images\n    concatenated_img = cv2.hconcat([image1, image2])\n\n    # Draw lines between the matching keypoints\n    for match in matches:\n        img1_idx = match[0]\n        img2_idx = match[1]\n        (x1, y1) = kpts1[img1_idx] * scale1\n        (x2, y2) = kpts2[img2_idx] * scale2\n        pt1 = (int(x1), int(y1))\n        pt2 = (int(x2) + image1.shape[1], int(y2))\n        cv2.line(concatenated_img, pt1, pt2, (0, 0, 255), 2)\n\n    # Plot the concatenated image\n    plt.figure(figsize=(20, 12))\n    plt.imshow(cv2.cvtColor(concatenated_img, cv2.COLOR_BGR2RGB))\n    plt.axis(\"off\")\n    plt.show()","metadata":{"papermill":{"duration":0.028147,"end_time":"2024-04-25T14:07:32.127691","exception":false,"start_time":"2024-04-25T14:07:32.099544","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Colmap database","metadata":{"papermill":{"duration":0.009918,"end_time":"2024-04-25T14:07:32.147656","exception":false,"start_time":"2024-04-25T14:07:32.137738","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Code to manipulate a colmap database.\n# Forked from https://github.com/colmap/colmap/blob/dev/scripts/python/database.py\n\n# Copyright (c) 2018, ETH Zurich and UNC Chapel Hill.\n# All rights reserved.\n#\n# Redistribution and use in source and binary forms, with or without\n# modification, are permitted provided that the following conditions are met:\n#\n#     * Redistributions of source code must retain the above copyright\n#       notice, this list of conditions and the following disclaimer.\n#\n#     * Redistributions in binary form must reproduce the above copyright\n#       notice, this list of conditions and the following disclaimer in the\n#       documentation and/or other materials provided with the distribution.\n#\n#     * Neither the name of ETH Zurich and UNC Chapel Hill nor the names of\n#       its contributors may be used to endorse or promote products derived\n#       from this software without specific prior written permission.\n#\n# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\n# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE\n# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE\n# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR\n# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF\n# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS\n# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN\n# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)\n# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE\n# POSSIBILITY OF SUCH DAMAGE.\n#\n# Author: Johannes L. Schoenberger (jsch-at-demuc-dot-de)\n\n# This script is based on an original implementation by True Price.\n\nimport sys\nimport sqlite3\nimport numpy as np\n\n\nIS_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(\n    MAX_IMAGE_ID\n)\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 = \"CREATE UNIQUE INDEX IF NOT EXISTS index_name ON images(name)\"\n\nCREATE_ALL = \"; \".join(\n    [\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    ]\n)\n\n\ndef 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    @staticmethod\n    def connect(database_path):\n        return sqlite3.connect(database_path, factory=COLMAPDatabase)\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 = lambda: self.executescript(CREATE_CAMERAS_TABLE)\n        self.create_descriptors_table = lambda: self.executescript(\n            CREATE_DESCRIPTORS_TABLE\n        )\n        self.create_images_table = lambda: self.executescript(CREATE_IMAGES_TABLE)\n        self.create_two_view_geometries_table = lambda: self.executescript(\n            CREATE_TWO_VIEW_GEOMETRIES_TABLE\n        )\n        self.create_keypoints_table = lambda: self.executescript(CREATE_KEYPOINTS_TABLE)\n        self.create_matches_table = lambda: self.executescript(CREATE_MATCHES_TABLE)\n        self.create_name_index = lambda: self.executescript(CREATE_NAME_INDEX)\n\n    def add_camera(\n        self, model, width, height, params, prior_focal_length=False, camera_id=None\n    ):\n        params = np.asarray(params, np.float64)\n        cursor = self.execute(\n            \"INSERT INTO cameras VALUES (?, ?, ?, ?, ?, ?)\",\n            (\n                camera_id,\n                model,\n                width,\n                height,\n                array_to_blob(params),\n                prior_focal_length,\n            ),\n        )\n        return cursor.lastrowid\n\n    def add_image(\n        self, name, camera_id, prior_q=np.zeros(4), prior_t=np.zeros(3), image_id=None\n    ):\n        cursor = self.execute(\n            \"INSERT INTO images VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)\",\n            (\n                image_id,\n                name,\n                camera_id,\n                prior_q[0],\n                prior_q[1],\n                prior_q[2],\n                prior_q[3],\n                prior_t[0],\n                prior_t[1],\n                prior_t[2],\n            ),\n        )\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\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\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\n    def add_two_view_geometry(\n        self,\n        image_id1,\n        image_id2,\n        matches,\n        F=np.eye(3),\n        E=np.eye(3),\n        H=np.eye(3),\n        config=2,\n    ):\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,)\n            + matches.shape\n            + (\n                array_to_blob(matches),\n                config,\n                array_to_blob(F),\n                array_to_blob(E),\n                array_to_blob(H),\n            ),\n        )","metadata":{"_kg_hide-input":true,"papermill":{"duration":0.041158,"end_time":"2024-04-25T14:07:32.198965","exception":false,"start_time":"2024-04-25T14:07:32.157807","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# DB operation","metadata":{"papermill":{"duration":0.009737,"end_time":"2024-04-25T14:07:32.218647","exception":false,"start_time":"2024-04-25T14:07:32.208910","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Modified from https://github.com/cvlab-epfl/disk/blob/37f1f7e971cea3055bb5ccfc4cf28bfd643fa339/colmap/h5_to_db.py\n\n#  Copyright [2020] [Michał Tyszkiewicz, Pascal Fua, Eduard Trulls]\n#\n#   Licensed under the Apache License, Version 2.0 (the \"License\");\n#   you may not use this file except in compliance with the License.\n#   You may obtain a copy of the License at\n#\n#       http://www.apache.org/licenses/LICENSE-2.0\n#\n#   Unless required by applicable law or agreed to in writing, software\n#   distributed under the License is distributed on an \"AS IS\" BASIS,\n#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n#   See the License for the specific language governing permissions and\n#   limitations under the License.\n\nimport os, argparse, h5py, warnings\nimport numpy as np\nfrom tqdm import tqdm\nfrom PIL import Image, ExifTags\n\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    \n    #\n    # Modified to add exif_ifd to exif dict\n    #\n    exif_ifd = exif.get_ifd(0x8769)\n    exif.update(exif_ifd)\n\n    focal = None\n    is_from_exif = False\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                is_from_exif = True\n                break\n\n        if focal_35mm is not None:\n            focal = focal_35mm / 35.0 * 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    #\n    # Modified to return a bool indicating if the focal length is from exif \n    #\n    return focal, is_from_exif\n\n\ndef create_camera(db, image_path, camera_model):\n    image = Image.open(image_path)\n    width, height = image.size\n\n    focal, is_from_exif = 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, 0.0, 0.0])\n\n    #\n    # Modified to set prior_focal_length if the focal length is from exif\n    #\n    return db.add_camera(\n        model, width, height, param_arr, prior_focal_length=is_from_exif\n    )\n\n\ndef add_kpts_matches(db, img_dir, kpts, matches, fms = None):\n    fname_to_id = {}\n\n    # Add keypoints\n    for filename in tqdm(kpts):\n        path = os.path.join(img_dir, filename)\n        camera_model = \"simple-radial\"\n        camera_id = create_camera(db, path, camera_model)\n        image_id = db.add_image(filename, camera_id)\n        fname_to_id[filename] = image_id\n        db.add_keypoints(image_id, kpts[filename])\n\n    n_keys = len(matches)\n    n_total = (n_keys * (n_keys - 1)) // 2\n    # Add matches\n    added = set()\n    with tqdm(total=n_total) as pbar:\n        for key1 in matches:\n            for key2 in matches[key1]:\n                id_1 = fname_to_id[key1]\n                id_2 = fname_to_id[key2]\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                db.add_matches(id_1, id_2, matches[key1][key2])\n                added.add(pair_id)\n                pbar.update(1)\n                if fms is not None:\n                    F = fms[key1][key2] if key1 in fms and key2 in fms[key1] else np.eye(3)\n                    db.add_two_view_geometry(id_1, id_2, matches[key1][key2], F, np.eye(3), np.eye(3))\n    db.commit()\n","metadata":{"papermill":{"duration":0.031153,"end_time":"2024-04-25T14:07:32.259633","exception":false,"start_time":"2024-04-25T14:07:32.228480","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def 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(\n        A, dim=dim, sorted=True, return_inverse=True, return_counts=True\n    )\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","metadata":{"papermill":{"duration":0.018667,"end_time":"2024-04-25T14:07:32.288336","exception":false,"start_time":"2024-04-25T14:07:32.269669","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# AffNetHardNet Models","metadata":{"papermill":{"duration":0.010133,"end_time":"2024-04-25T14:07:32.308440","exception":false,"start_time":"2024-04-25T14:07:32.298307","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# 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 = \"keynet\"\n    ):\n        detector_options = [\"keynet\", \"GFTT\", \"Hessian\", \"Harris\", \"DoG\"]\n        if detector not in detector_options:\n            raise ValueError(\"Detector must be one of {}\".format(detector_options))\n        \n        ori_module = (\n            KF.PassLAF()\n            if upright\n            else KF.LAFOrienter(angle_detector=KF.OriNet(False)).eval()\n        )\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\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\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        affnet_weights = torch.load(os.path.join(MODEL_DIR, \"AffNet.pth\"))[\"state_dict\"]\n        detector.aff.load_state_dict(affnet_weights)\n\n        # hardnet = KF.HardNet(False).eval()\n        # hn_weights = torch.load(os.path.join(MODEL_DIR, \"HardNetLib.pth\"))[\"state_dict\"]\n        # hardnet.load_state_dict(hn_weights)\n        # descriptor = KF.LAFDescriptor(\n        #     hardnet, patch_size=32, grayscale_descriptor=True\n        # ).to(device)\n        hardnet8 = KF.HardNet8(False).eval()\n        hn8_weights = torch.load(HARDNET_PT)\n        hardnet8.load_state_dict(hn8_weights)\n        descriptor = KF.LAFDescriptor(\n            hardnet8, patch_size=32, grayscale_descriptor=True\n        ).to(device)\n        super().__init__(detector, descriptor, scale_laf)\n","metadata":{"papermill":{"duration":0.027741,"end_time":"2024-04-25T14:07:32.346100","exception":false,"start_time":"2024-04-25T14:07:32.318359","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_unique_matches(f_match_kpts):\n    kpts = defaultdict(list)\n    match_indexes = defaultdict(dict)\n    total_kpts = defaultdict(int)\n    for key1 in f_match_kpts:\n        for key2 in f_match_kpts[key1]:\n            matches = f_match_kpts[key1][key2]\n            kpts[key1].append(matches[:, :2])\n            kpts[key2].append(matches[:, 2:])\n            current_match = torch.arange(len(matches)).reshape(-1, 1).repeat(1, 2)\n            current_match[:, 0] += total_kpts[key1]\n            current_match[:, 1] += total_kpts[key2]\n            total_kpts[key1] += len(matches)\n            total_kpts[key2] += len(matches)\n            match_indexes[key1][key2] = current_match\n\n    for key in kpts:\n        kpts[key] = np.round(np.concatenate(kpts[key], axis=0))\n\n    unique_kpts = {}\n    unique_match_idxs = {}\n    out_match = defaultdict(dict)\n\n    for key in kpts.keys():\n        uniq_kps, uniq_reverse_idxs = torch.unique(\n            torch.from_numpy(kpts[key]), dim=0, return_inverse=True\n        )\n        unique_match_idxs[key] = uniq_reverse_idxs\n        unique_kpts[key] = uniq_kps.numpy()\n\n    for key1 in match_indexes:\n        for key2 in match_indexes[key1]:\n            m2 = deepcopy(match_indexes[key1][key2])\n            m2[:, 0] = unique_match_idxs[key1][m2[:, 0]]\n            m2[:, 1] = unique_match_idxs[key2][m2[:, 1]]\n            mkpts = np.concatenate(\n                [\n                    unique_kpts[key1][m2[:, 0]],\n                    unique_kpts[key2][m2[:, 1]],\n                ],\n                axis=1,\n            )\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[key1][key2] = m2_semiclean2.numpy()\n    return unique_kpts, out_match","metadata":{"papermill":{"duration":0.026249,"end_time":"2024-04-25T14:07:32.384070","exception":false,"start_time":"2024-04-25T14:07:32.357821","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Scene feature detector","metadata":{"papermill":{"duration":0.010438,"end_time":"2024-04-25T14:07:32.404685","exception":false,"start_time":"2024-04-25T14:07:32.394247","status":"completed"},"tags":[]}},{"cell_type":"code","source":"class AffNetHardNetDetector:\n    def __init__(\n        self,\n        model,\n        device=torch.device(\"cuda\"),\n        resize_long_edge_to=600,\n        model_type=\"hardnet\",\n        min_matches=15,\n        rgb_input = False\n    ):\n        self.rgb_input = rgb_input\n        print(\"Init AffNetHardNetDetector\")\n        self.model = model\n        self.model_type = model_type\n        self.device = device\n        self.resize_long_edge_to = resize_long_edge_to\n        print(\"Longer edge will be resized to\", self.resize_long_edge_to)\n\n    def detect_features(self, img_fnames):\n        f_lafs = dict()\n        f_descs = dict()\n        f_kpts = dict()\n        f_raw_size = dict()\n        f_matches = dict()\n        # Get features\n        print(\"Detecting AffNetHardNet features\")\n        for img_path in tqdm(img_fnames):\n            img_fname = img_path.split(\"/\")[-1]\n            key = img_fname\n            f_matches[key] = dict()\n            with torch.inference_mode():\n                timg = load_torch_image(img_path, device=device)\n                raw_size = torch.tensor(timg.shape[2:])\n                timg_resized, h_scale, w_scale = resize_torch_image(\n                    timg, self.resize_long_edge_to, disable_enlarge=True\n                )\n                timg_resized = timg_resized if self.rgb_input else K.color.rgb_to_grayscale(timg_resized)\n                if self.model_type == 'hardnet':\n                    lafs, resps, descs = self.model(timg_resized)\n                elif self.model_type == 'disk':\n                    features = self.model(timg_resized, 8000, pad_if_not_divisible=True)[0]\n                    kps, descs = features.keypoints, features.descriptors\n                    lafs = KF.laf_from_center_scale_ori(kps[None], torch.ones(1, len(kps), 1, 1, device=device))\n                elif self.model_type == 'aliked':\n                    features = self.model.extract(timg_resized)\n                    kps, descs = features['keypoints'].squeeze(), features['descriptors'].squeeze()\n                    lafs = KF.laf_from_center_scale_ori(kps[None], torch.ones(1, len(kps), 1, 1, device=device))\n\n                # Recover scale?\n                lafs[:, :, 0, :] *= 1 / w_scale\n                lafs[:, :, 1, :] *= 1 / h_scale\n                desc_dim = descs.shape[-1]\n                # Move keypoints to cpu for later colmap operations\n                kpts = KF.get_laf_center(lafs).reshape(-1, 2).detach().cpu().numpy()\n                descs = descs.reshape(-1, desc_dim).detach()\n                f_lafs[key] = lafs.detach()\n                f_kpts[key] = kpts\n                f_descs[key] = descs\n                f_raw_size[key] = raw_size.clone().detach()\n        gc.collect()\n        torch.cuda.empty_cache()\n        return f_lafs, f_kpts, f_descs, f_raw_size","metadata":{"papermill":{"duration":0.026793,"end_time":"2024-04-25T14:07:32.441733","exception":false,"start_time":"2024-04-25T14:07:32.414940","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Scene LAF matcher","metadata":{"papermill":{"duration":0.011063,"end_time":"2024-04-25T14:07:32.464118","exception":false,"start_time":"2024-04-25T14:07:32.453055","status":"completed"},"tags":[]}},{"cell_type":"code","source":"class LafMatcher:\n    def __init__(self, min_matches=15, threshold_matches=0.5, device=\"cuda\", matcher=\"adalam\"):\n        self.adalam_config = KF.adalam.get_adalam_default_config()\n        self.adalam_config[\"force_seed_mnn\"] = True\n        self.adalam_config[\"search_expansion\"] = 16\n        self.adalam_config[\"ransac_iters\"] = 256\n        self.adalam_config[\"device\"] = device\n        # self.adalam_config[\"orientation_difference_threshold\"] = None\n        # self.adalam_config['scale_rate_threshold'] = None\n        self.min_matches = min_matches\n        self.threshold_matches = threshold_matches\n        self.matcher = matcher\n        if self.matcher == \"disk_lightglue\":\n            self.lg_matcher = KF.LightGlueMatcher(\"disk\").eval().to(device)\n        elif self.matcher == \"aliked_lightglue\":\n            self.lg_matcher = KF.LightGlueMatcher(\"aliked\").eval().to(device)\n\n    def match(self, img_fnames, f_lafs, f_kpts, f_descs, f_raw_size, get_roi = False):\n        index_pairs = dict()\n        num_imgs = len(img_fnames)\n        print(\"Matching to get index pairs\")\n        pair_count = 0\n        f_matches = defaultdict(dict)\n        f_rois = defaultdict(dict)\n        for idx1 in tqdm(range(num_imgs - 1)):\n            index_pairs[idx1] = []\n            for idx2 in range(idx1 + 1, num_imgs):\n                fname1, fname2 = img_fnames[idx1], img_fnames[idx2]\n                key1, key2 = fname1.split(\"/\")[-1], fname2.split(\"/\")[-1]\n                lafs1 = f_lafs[key1]\n                lafs2 = f_lafs[key2]\n                desc1 = f_descs[key1]\n                desc2 = f_descs[key2]\n                if self.matcher == \"adalam\":\n                    hw1, hw2 = f_raw_size[key1], f_raw_size[key2]\n\n                    dists, idxs = KF.match_adalam(\n                        desc1,\n                        desc2,\n                        lafs1,\n                        lafs2,  # Adalam takes into account also geometric information\n                        hw1=hw1,\n                        hw2=hw2,\n                        config=self.adalam_config,\n                    )  # Adalam also benefits from knowing image size\n                elif self.matcher in [\"disk_lightglue\",\"aliked_lightglue\"]:\n                    print()\n                    hw1, hw2 = f_raw_size[key1], f_raw_size[key2]\n                    with torch.inference_mode():\n                        dists, idxs = self.lg_matcher(desc1, desc2, lafs1, lafs2, hw1=hw1, hw2=hw2)\n                else:\n                    dists, idxs = KF.match_smnn(desc1, desc2, 0.98)\n\n                if dists.mean().detach().cpu().numpy() < self.threshold_matches:\n                    first_indices = get_unique_idxs(idxs[:, 1])\n                    idxs = idxs[first_indices]\n                    dists = dists[first_indices]\n                    n_matches = len(idxs)\n                    if n_matches >= self.min_matches:\n                        pair_count += 1\n                        index_pairs[idx1].append(\n                            [idx2, dists.mean().detach().cpu().numpy().item(), n_matches]\n                        )\n                        f_matches[key1][key2] = (\n                            idxs.detach().cpu().numpy().reshape(-1, 2)\n                        )\n\n                        # Compute ROI\n                        if get_roi:\n                            mkpts1 = f_kpts[key1][idxs.cpu().numpy()[:, 0]]\n                            mkpts2 = f_kpts[key2][idxs.cpu().numpy()[:, 1]]\n                            roi_min_w_1, roi_max_w_1 = np.percentile(mkpts1[:, 0], [5, 95])\n                            roi_min_h_1, roi_max_h_1 = np.percentile(mkpts1[:, 1], [5, 95])\n                            roi_area_1 = (roi_max_w_1 - roi_min_w_1) * (\n                                roi_max_h_1 - roi_min_h_1\n                            )\n                            roi1 = {\n                                \"roi_min_w\": roi_min_w_1,\n                                \"roi_min_h\": roi_min_h_1,\n                                \"roi_max_w\": roi_max_w_1,\n                                \"roi_max_h\": roi_max_h_1,\n                                \"area\": roi_area_1,\n                            }\n                            roi_min_w_2, roi_max_w_2 = np.percentile(mkpts2[:, 0], [5, 95])\n                            roi_min_h_2, roi_max_h_2 = np.percentile(mkpts2[:, 1], [5, 95])\n                            roi_area_2 = (roi_max_w_2 - roi_min_w_2) * (\n                                roi_max_h_2 - roi_min_h_2\n                            )\n                            roi2 = {\n                                \"roi_min_w\": roi_min_w_2,\n                                \"roi_min_h\": roi_min_h_2,\n                                \"roi_max_w\": roi_max_w_2,\n                                \"roi_max_h\": roi_max_h_2,\n                                \"area\": roi_area_2,\n                            }\n                            f_rois[key1][key2] = [roi1, roi2]\n\n        print(f\" Get {pair_count} from {int(num_imgs * (num_imgs-1)/2)} possible pairs\")\n        torch.cuda.empty_cache()\n        gc.collect()\n        return index_pairs, f_kpts, f_matches, f_rois","metadata":{"papermill":{"duration":0.034196,"end_time":"2024-04-25T14:07:32.509069","exception":false,"start_time":"2024-04-25T14:07:32.474873","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Matches and pairs operation","metadata":{"papermill":{"duration":0.010603,"end_time":"2024-04-25T14:07:32.530457","exception":false,"start_time":"2024-04-25T14:07:32.519854","status":"completed"},"tags":[]}},{"cell_type":"code","source":"def merge_kpts_matches(kpts, matches, new_kpts, new_matches, cap = None):\n    # merge kpts\n    prev_len = dict()\n    for new_key in new_kpts:\n        if new_key in kpts:\n            old_len = len(kpts[new_key])\n            kpts[new_key] = np.concatenate([kpts[new_key], new_kpts[new_key]], axis=0)\n        else:\n            old_len = 0\n            kpts[new_key] = new_kpts[new_key]\n        prev_len[new_key] = old_len\n\n    for new_key1 in new_matches:\n        for new_key2 in new_matches[new_key1]:\n            old_len1 = prev_len[new_key1]\n            old_len2 = prev_len[new_key2]\n            new_match = new_matches[new_key1][new_key2] + [old_len1, old_len2]\n            if cap is not None and len(new_match) > cap:\n                keep = np.random.choice(len(new_match), cap, replace=False)\n                new_match = new_match[keep, :]\n            if new_key1 in matches and new_key2 in matches[new_key1]:\n\n                matches[new_key1][new_key2] = np.concatenate(\n                    [\n                        matches[new_key1][new_key2],\n                        new_match,\n                    ],\n                    axis=0,\n                )\n            else:\n                if new_key1 not in matches:\n                    matches[new_key1] = dict()\n                matches[new_key1][new_key2] = new_match\n    return kpts, matches\n\n\ndef keep_matches(matches, max_num=None):\n    if max_num is None:\n        return matches\n    if len(matches) > max_num:\n        # radnomly select max_num matches\n        matches = np.random.choice(matches, max_num, replace=False)\n    return matches\n\n\ndef keep_pairs(index_pairs, max_num_pairs=20):\n    new_count = 0\n    old_count = 0\n    new_idx_count = defaultdict(int)\n    new_pairs = defaultdict(list)\n    for key1 in index_pairs:\n        # sort pairs by number of pairs\n        index_pairs[key1] = sorted(index_pairs[key1], key=lambda x: x[2], reverse=True)\n        for pair in index_pairs[key1]:\n            old_count += 1\n            idx1 = key1\n            idx2 = pair[0]\n\n            if new_idx_count[key1] < max_num_pairs:\n                new_pairs[idx1].append(pair)\n                new_count += 1\n                new_idx_count[idx1] += 1\n                new_idx_count[idx2] += 1\n            else:\n                continue\n\n    if DEBUG:\n        print(f\"origin pairs: {old_count}, kept pairs: {new_count}\")\n    return index_pairs\n\ndef select_matches(matches, keep_ratio = 0.01):\n    max_matches = defaultdict(int)\n    old_matches_count = 0\n    for key1 in matches:\n        for key2 in matches[key1]:\n            max_matches[key1] = max(max_matches[key1], len(matches[key1][key2]))\n            max_matches[key2] = max(max_matches[key2], len(matches[key1][key2]))\n            old_matches_count +=1\n\n    new_matches_count = 0\n    new_matches = defaultdict(dict)\n    for key1 in matches:\n        for key2 in matches[key1]:\n            n_matches = len(matches[key1][key2])\n            if n_matches > max_matches[key1] * keep_ratio or n_matches > max_matches[key2] * keep_ratio:\n                new_matches[key1][key2] = matches[key1][key2]\n                new_matches_count+=1\n    if DEBUG:\n        print(f\"origin matches: {old_matches_count}, kept matches: {new_matches_count}\")\n    return new_matches","metadata":{"papermill":{"duration":0.030412,"end_time":"2024-04-25T14:07:32.571594","exception":false,"start_time":"2024-04-25T14:07:32.541182","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Get Fundamental matrices from matches and keypoints","metadata":{"papermill":{"duration":0.009993,"end_time":"2024-04-25T14:07:32.591711","exception":false,"start_time":"2024-04-25T14:07:32.581718","status":"completed"},"tags":[]}},{"cell_type":"code","source":"def get_fms(kpts, matches):\n    prev_len = dict()\n    fms = defaultdict(dict)\n    print(\"Get Fundamental Matrix\")\n    for key1 in tqdm(matches):\n        for key2 in matches[key1]:\n            match = matches[key1][key2]\n            mkpts1 = kpts[key1][match[:, 0]]\n            mkpts2 = kpts[key2][match[:, 1]]\n            try:    \n                Fm, inliers = cv2.findFundamentalMat(mkpts1, mkpts2, cv2.USAC_MAGSAC, FM_PARAMS[\"ransacReprojThreshold\"], FM_PARAMS[\"confidence\"], FM_PARAMS[\"maxIters\"])\n                if FM_PARAMS[\"removeOutliers\"] == True:\n                    new_match = match[inliers.ravel() == 1]\n                    matches[key1][key2] = new_match\n                fms[key1][key2] = Fm\n            except:\n                pass\n    # print(Fm.shape)\n    return kpts, matches, fms","metadata":{"papermill":{"duration":0.020353,"end_time":"2024-04-25T14:07:32.622040","exception":false,"start_time":"2024-04-25T14:07:32.601687","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model Setup","metadata":{"papermill":{"duration":0.009612,"end_time":"2024-04-25T14:07:32.641933","exception":false,"start_time":"2024-04-25T14:07:32.632321","status":"completed"},"tags":[]}},{"cell_type":"code","source":"if MODEL_DICT[\"Keynet\"][\"enable\"]:\n    keynet_model = (\n        AffNetHardNet(num_features=8000, upright=False, device=device, detector=\"keynet\")\n        .to(device)\n        .eval()\n    )\n    keynet_detector = AffNetHardNetDetector(keynet_model, resize_long_edge_to=MODEL_DICT[\"Keynet\"][\"resize_long_edge_to\"])\n    \nif MODEL_DICT[\"GFTT\"][\"enable\"]:\n    gftt_model = (\n        AffNetHardNet(num_features=8000, upright=False, device=device, detector=\"GFTT\")\n        .to(device)\n        .eval()\n    )\n    gftt_detector = AffNetHardNetDetector(gftt_model, resize_long_edge_to=MODEL_DICT[\"GFTT\"][\"resize_long_edge_to\"])\n\nif MODEL_DICT[\"DoG\"][\"enable\"]:\n    DoG_model = (\n        AffNetHardNet(num_features=8000, upright=False, device=device, detector=\"DoG\")\n        .to(device)\n        .eval()\n    )\n    DoG_detector = AffNetHardNetDetector(DoG_model, resize_long_edge_to=MODEL_DICT[\"DoG\"][\"resize_long_edge_to\"])\n\nif MODEL_DICT[\"Harris\"][\"enable\"]:\n    harris_model = (\n        AffNetHardNet(num_features=8000, upright=False, device=device, detector=\"Harris\")\n        .to(device)\n        .eval()\n    )\n    harris_detector = AffNetHardNetDetector(harris_model, resize_long_edge_to=MODEL_DICT[\"Harris\"][\"resize_long_edge_to\"])\n    \nif MODEL_DICT[\"disk_lightglue\"][\"enable\"]:\n    disk_model = KF.DISK.from_pretrained(\"depth\").to(device).eval()\n    disk_detector = AffNetHardNetDetector(disk_model, model_type='disk', rgb_input=True, resize_long_edge_to=MODEL_DICT[\"disk_lightglue\"][\"resize_long_edge_to\"])\n    \nif MODEL_DICT[\"aliked_lightglue\"][\"enable\"]:\n    aliked_model = ALIKED(detection_threshold=0.005, max_num_keypoints=8000).eval().to(device)\n    aliked_model.preprocess_conf = {'resize': MODEL_DICT[\"aliked_lightglue\"][\"resize_long_edge_to\"]}\n    aliked_detector = AffNetHardNetDetector(aliked_model, model_type='aliked', rgb_input=True, resize_long_edge_to=MODEL_DICT[\"aliked_lightglue\"][\"resize_long_edge_to\"])","metadata":{"papermill":{"duration":3.8028,"end_time":"2024-04-25T14:07:36.454606","exception":false,"start_time":"2024-04-25T14:07:32.651806","status":"completed"},"scrolled":true,"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Function to generate scene db for reconstruction","metadata":{"papermill":{"duration":0.012643,"end_time":"2024-04-25T14:07:36.483191","exception":false,"start_time":"2024-04-25T14:07:36.470548","status":"completed"},"tags":[]}},{"cell_type":"code","source":"def generate_scene_db(dataset, scene):\n    feature_det_start = time()\n    # Process a scene and write matches and keypoints to the database\n    img_dir = f\"{SRC}/{MODE}/{dataset}/images\"\n    if not os.path.exists(img_dir):\n        print(\"Image dir does not exist:\", img_dir)\n        return\n\n    img_fnames = [f\"{SRC}/{x}\" for x in data_dict[dataset][scene]]\n    print(f\"Got {len(img_fnames)} images\")\n\n    matches = dict()\n    kpts = dict()\n\n    if MODEL_DICT[\"Keynet\"][\"enable\"]:\n#         import pdb; pdb.set_trace()\n        f_lafs, f_kpts, f_descs, f_raw_size = keynet_detector.detect_features(\n            img_fnames\n        )\n        keynet_pairs, keynet_kpts, keynet_matches, keynet_rois = laf_matcher.match(\n            img_fnames, f_lafs, f_kpts, f_descs, f_raw_size\n        )\n        if not MODEL_DICT[\"Keynet\"][\"pair_only\"]:\n            kpts, matches = merge_kpts_matches(kpts, matches, keynet_kpts, keynet_matches, MATCHES_CAP)\n\n    if MODEL_DICT[\"GFTT\"][\"enable\"]:\n        gftt_lafs, gftt_kpts, gftt_descs, gftt_raw_size = gftt_detector.detect_features(\n            img_fnames\n        )\n        index_pairs, gftt_kpts, gftt_matches, gftt_rois = laf_matcher.match(\n            img_fnames, gftt_lafs, gftt_kpts, gftt_descs, gftt_raw_size\n        )\n        kpts, matches = merge_kpts_matches(kpts, matches, gftt_kpts, gftt_matches, MATCHES_CAP)\n\n    if MODEL_DICT[\"DoG\"][\"enable\"]:\n        DoG_lafs, DoG_kpts, DoG_descs, DoG_raw_size = DoG_detector.detect_features(\n            img_fnames\n        )\n        index_pairs, DoG_kpts, DoG_matches, DoG_rois = laf_matcher.match(\n            img_fnames, DoG_lafs, DoG_kpts, DoG_descs, DoG_raw_size\n        )\n        kpts, matches = merge_kpts_matches(kpts, matches, DoG_kpts, DoG_matches, MATCHES_CAP)\n        \n    if MODEL_DICT[\"Harris\"][\"enable\"]:\n        harris_lafs, harris_kpts, harris_descs, harris_raw_size = harris_detector.detect_features(\n            img_fnames\n        )\n        harris_pairs, harris_kpts, harris_matches, harris_rois = laf_matcher.match(\n            img_fnames, harris_lafs, harris_kpts, harris_descs, harris_raw_size\n        )\n        kpts, matches = merge_kpts_matches(kpts, matches, harris_kpts, harris_matches, MATCHES_CAP)\n\n    if MODEL_DICT[\"disk_lightglue\"][\"enable\"]:\n        disk_lafs, disk_kpts, disk_descs, disk_raw_size = disk_detector.detect_features(img_fnames)\n        disk_pairs, disk_kpts, disk_matches, disk_rois = disk_glue_matcher.match(\n                img_fnames, disk_lafs, disk_kpts, disk_descs, disk_raw_size,\n            )\n        kpts, matches = merge_kpts_matches(kpts, matches, disk_kpts, disk_matches, MATCHES_CAP)\n        \n    if MODEL_DICT[\"aliked_lightglue\"][\"enable\"]:\n        aliked_lafs, aliked_kpts, aliked_descs, aliked_raw_size = aliked_detector.detect_features(img_fnames)\n        aliked_pairs, aliked_kpts, aliked_matches, aliked_rois = aliked_glue_matcher.match(\n                img_fnames, aliked_lafs, aliked_kpts, aliked_descs, aliked_raw_size,\n            )\n        kpts, matches = merge_kpts_matches(kpts, matches, aliked_kpts, aliked_matches, MATCHES_CAP)\n    \n    # Get fundamental matrices\n    kpts, matches, fms = get_fms(kpts, matches)\n    \n    matches = select_matches(matches, MATCH_FILTER_RATIO)\n    \n    if DEBUG:\n        import random\n        random.seed(0)\n        for i in range(5):\n            print(matches.keys())\n           \n            key1 = random.choice(list(matches.keys()))\n            key2 = random.choice(list(matches[key1].keys()))\n            print(key1, key2)\n            fname1, fname2 = os.path.join(img_dir, key1), os.path.join(img_dir, key2)\n\n            print(\"Plot Combined matches\")\n            plot_images_with_keypoints(\n                fname1, fname2, kpts[key1], kpts[key2], matches[key1][key2]\n            )\n    # Write to database\n    feature_dir = f\"featureout/{dataset}_{scene}\"\n    if not os.path.isdir(feature_dir):\n        os.makedirs(feature_dir, exist_ok=True)\n    database_path = f\"{feature_dir}/colmap.db\"\n    if os.path.isfile(database_path):\n        os.remove(database_path)\n\n    db = COLMAPDatabase.connect(database_path)\n    db.create_tables()\n    single_camera = False\n    print(\"Add kpts and matches to database\")\n    add_kpts_matches(db, img_dir, kpts, matches, fms)\n\n    feature_det_end = time()\n    matching_time = feature_det_end - feature_det_start\n    torch.cuda.empty_cache()\n    gc.collect()\n\n    return matching_time","metadata":{"papermill":{"duration":0.039818,"end_time":"2024-04-25T14:07:36.533762","exception":false,"start_time":"2024-04-25T14:07:36.493944","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Function of reconstruction","metadata":{"papermill":{"duration":0.010441,"end_time":"2024-04-25T14:07:36.555425","exception":false,"start_time":"2024-04-25T14:07:36.544984","status":"completed"},"tags":[]}},{"cell_type":"code","source":"def reconstruct_from_db(dataset, scene):\n    scene_result = {}\n    reconst_start = time()\n    \n    img_dir = f\"{SRC}/{MODE}/{dataset}/images\"\n    if not os.path.exists(img_dir):\n        print(\"Image dir does not exist:\", img_dir)\n        return\n\n    feature_dir = f\"featureout/{dataset}_{scene}\"\n    database_path = f\"{feature_dir}/colmap.db\"\n    db = COLMAPDatabase.connect(database_path)\n    output_path = f\"{feature_dir}/colmap_rec\"\n    t = time()\n    gc.collect()\n\n#    # Skip match_exhaustive\n#     pycolmap.match_exhaustive(database_path, match_options)\n\n    t = time() - t\n    print(f\"RANSAC in  {t:.4f} sec\")\n    t = time()\n    \n    # By default colmap does not generate a reconstruction if less than 10 images are registered. Lower it to 3.\n    mapper_options = pycolmap.IncrementalPipelineOptions({'num_threads':1})\n    mapper_options.min_model_size = 3\n\n    for attribute_name in dir(mapper_options):\n        if not attribute_name.startswith(\"__\"):\n            attribute_value = getattr(mapper_options, attribute_name)\n            print(f\"{attribute_name}: {attribute_value}\")\n    os.makedirs(output_path, exist_ok=True)\n    maps = pycolmap.incremental_mapping(\n        database_path=database_path,\n        image_path=img_dir,\n        output_path=output_path,\n        options=mapper_options,\n    )\n    print(maps)\n    # clear_output(wait=False)\n    t = time() - t\n    print(f\"Reconstruction done in  {t:.4f} sec\")\n    imgs_registered = 0\n    best_idx = None\n    print(\"Looking for the best reconstruction\")\n    if isinstance(maps, dict):\n        for idx1, rec in maps.items():\n            if len(rec.images) > imgs_registered:\n                imgs_registered = len(rec.images)\n                best_idx = idx1\n    if best_idx is not None:\n        for k, im in maps[best_idx].images.items():\n            key1 = f\"{MODE}/{dataset}/images/{im.name}\"\n            scene_result[key1] = {}\n            scene_result[key1][\"R\"] = deepcopy(im.cam_from_world.rotation.matrix())\n            scene_result[key1][\"t\"] = deepcopy(np.array(im.cam_from_world.translation))\n\n    return scene_result, imgs_registered","metadata":{"papermill":{"duration":0.028894,"end_time":"2024-04-25T14:07:36.595095","exception":false,"start_time":"2024-04-25T14:07:36.566201","status":"completed"},"scrolled":true,"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Main Loop","metadata":{"papermill":{"duration":0.011817,"end_time":"2024-04-25T14:07:36.617921","exception":false,"start_time":"2024-04-25T14:07:36.606104","status":"completed"},"tags":[]}},{"cell_type":"code","source":"    laf_matcher = LafMatcher(min_matches=15, threshold_matches=0.5, device=device)    \n    disk_glue_matcher = LafMatcher(device=device, min_matches=15, threshold_matches=0.5, matcher='disk_lightglue')\n    aliked_glue_matcher = LafMatcher(device=device, min_matches=45, threshold_matches=0.55,matcher='aliked_lightglue')\n    \n    datasets = []\n    datasets = [dataset for dataset in data_dict]\n\n    out_results1 = defaultdict(dict)\n    with concurrent.futures.ProcessPoolExecutor(max_workers=NUM_CORES) as executors:\n        futures = defaultdict(dict)\n\n        for dataset, scene in all_scenes:\n            generate_scene_db(dataset, scene)\n            futures[dataset][scene] = executors.submit(reconstruct_from_db, dataset, scene)\n\n        for dataset, scene in all_scenes:\n            result, imgs_registered = futures[dataset][scene].result()\n            if result is not None:\n                out_results1[dataset][scene] = result\n                out_results1[dataset][scene]['imgs_registered'] = imgs_registered \n                \n    clear_output(wait=False)","metadata":{"scrolled":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#     laf_matcher = LafMatcher(min_matches=100, threshold_matches=0.5, device=device)    \n#     disk_glue_matcher = LafMatcher(device=device, min_matches=100, threshold_matches=0.5, matcher='disk_lightglue')\n#     aliked_glue_matcher = LafMatcher(device=device, min_matches=100, threshold_matches=0.5,matcher='aliked_lightglue')\n    \n#     datasets = []\n#     datasets = [dataset for dataset in data_dict]\n\n#     out_results2 = defaultdict(dict)\n#     with concurrent.futures.ProcessPoolExecutor(max_workers=NUM_CORES) as executors:\n#         futures = defaultdict(dict)\n\n#         for dataset, scene in all_scenes:\n#             generate_scene_db(dataset, scene)\n#             futures[dataset][scene] = executors.submit(reconstruct_from_db, dataset, scene)\n\n#         for dataset, scene in all_scenes:\n#             result, imgs_registered = futures[dataset][scene].result()\n#             if result is not None:\n#                 out_results2[dataset][scene] = result\n#                 out_results2[dataset][scene]['imgs_registered'] = imgs_registered \n                \n#     clear_output(wait=False)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"create_submission(out_results1, {}, data_dict, MODE)","metadata":{"papermill":{"duration":0.125693,"end_time":"2024-04-25T14:22:01.846385","exception":false,"start_time":"2024-04-25T14:22:01.720692","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]}]}