{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","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"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":71885,"databundleVersionId":8143495,"sourceType":"competition"},{"sourceId":3414836,"sourceType":"datasetVersion","datasetId":2058261},{"sourceId":5373920,"sourceType":"datasetVersion","datasetId":3117886},{"sourceId":7884485,"sourceType":"datasetVersion","datasetId":4628051},{"sourceId":170475544,"sourceType":"kernelVersion"},{"sourceId":170565695,"sourceType":"kernelVersion"},{"sourceId":174129945,"sourceType":"kernelVersion"},{"sourceId":175679956,"sourceType":"kernelVersion"},{"sourceId":175684111,"sourceType":"kernelVersion"},{"sourceId":176463227,"sourceType":"kernelVersion"},{"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}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Updated information:\n- Version5:\n - Image matching methods\n - fixed: Superpoint + LightGlue ... bug fixed\n- Version4:\n - Image matching methods\n - added: MatchFormer\n - added: SIFT + LightGlue\n - added: DISK + LightGlue\n - modified: Aliked + LightGlue ... speeding up by cuda cache of keypoints/descriptors\n - modified: Superpoint + LightGlue ... speeding up by cuda cache of keypoints/descriptors\n - modified: DogHardNet + LightGlue ... speeding up by cuda cache of keypoints/descriptors\n - modified: Superpoint + SuperGlue ... added `torch.no_grad()` and speeding up by cuda cache of keypoints/descriptors\n - Configuration\n - added: CAMERA_MODEL = \"simple-radial\" or \"simple-pinhole\"\n - added: ROTATION_CORRECTION ... `check_orientation` (LightGlue series only are supported. Others image matching methods are under construction.)\n - https://github.com/ternaus/check_orientation\n - added: DRY_RUN ... to run pipeline with only 10 images\n - Pipeline\n - Parallel execution of image matching and COLMAP processing","metadata":{}},{"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-05-23T06:29:38.863442Z","iopub.execute_input":"2024-05-23T06:29:38.863874Z","iopub.status.idle":"2024-05-23T06:31:53.064634Z","shell.execute_reply.started":"2024-05-23T06:29:38.863841Z","shell.execute_reply":"2024-05-23T06:31:53.063306Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%matplotlib inline","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2024-05-23T06:31:53.0666Z","iopub.execute_input":"2024-05-23T06:31:53.066923Z","iopub.status.idle":"2024-05-23T06:31:53.073313Z","shell.execute_reply.started":"2024-05-23T06:31:53.066895Z","shell.execute_reply":"2024-05-23T06:31:53.07238Z"},"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 numpy as np\nimport pandas as pd\nimport h5py\nfrom IPython.display import clear_output\nfrom collections import defaultdict\nfrom copy import deepcopy\nimport concurrent.futures\nfrom collections import Counter\n\n# CV/ML\nimport cv2\nimport torch\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nimport kornia as K\nimport kornia.feature as KF\nfrom PIL import Image\nimport timm\nfrom timm.data import resolve_data_config\nfrom timm.data.transforms_factory import create_transform\n\nimport torchvision\n\n# 3D reconstruction\nimport pycolmap\n\nimport glob\nimport matplotlib\nfrom matplotlib import pyplot as plt\n\n# dkm\nimport sys\nsys.path.append('/kaggle/input/dkm-dependencies/DKM/')\nfrom dkm.utils.utils import tensor_to_pil, get_tuple_transform_ops\nfrom dkm import DKMv3_outdoor\n\n# LoFTR\nfrom kornia.feature import LoFTR\n\n# LightGlue\nfrom lightglue import match_pair\nfrom lightglue import ALIKED, SuperPoint, DoGHardNet, LightGlue, DISK, SIFT\nfrom lightglue.utils import load_image, rbd","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2024-05-23T06:31:53.074637Z","iopub.execute_input":"2024-05-23T06:31:53.074932Z","iopub.status.idle":"2024-05-23T06:32:00.102598Z","shell.execute_reply.started":"2024-05-23T06:31:53.07491Z","shell.execute_reply":"2024-05-23T06:32:00.101779Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print('Kornia version', K.__version__)\nprint('Pycolmap version', pycolmap.__version__)","metadata":{"execution":{"iopub.status.busy":"2024-05-23T06:32:00.104762Z","iopub.execute_input":"2024-05-23T06:32:00.105046Z","iopub.status.idle":"2024-05-23T06:32:00.109655Z","shell.execute_reply.started":"2024-05-23T06:32:00.105022Z","shell.execute_reply":"2024-05-23T06:32:00.108775Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Configurations","metadata":{}},{"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 = False\n \n # Keypoints handling\n MERGE_PARAMS = {\n \"min_matches\" : 15,\n \n # When merging keypoints, it is enable to filtering matches with cv2.findFundamentalMatrix.\n \"filter_FundamentalMatrix\" : False,\n \"filter_iterations\" : 10,\n \"filter_threshold\" : 8,\n }\n \n # Keypoints Extraction\n use_aliked_lightglue = True\n use_doghardnet_lightglue = False\n use_superpoint_lightglue = False\n use_disk_lightglue = False\n use_sift_lightglue = False\n use_loftr = False\n use_dkm = False\n use_superglue = False\n use_matchformer = False\n \n # Keypoints Extraction Parameters\n params_aliked_lightglue = {\n \"num_features\" : 8192,\n \"detection_threshold\" : 0.0000078,\n \"min_matches\" : 15,\n \"resize_to\" : 1024,\n }\n \n params_doghardnet_lightglue = {\n \"num_features\" : 8192,\n \"detection_threshold\" : 0.001,\n \"min_matches\" : 15,\n \"resize_to\" : 1024,\n }\n \n params_superpoint_lightglue = {\n \"num_features\" : 4096,\n \"detection_threshold\" : 0.005,\n \"min_matches\" : 15,\n \"resize_to\" : 1024,\n }\n \n params_disk_lightglue = {\n \"num_features\" : 8192,\n \"detection_threshold\" : 0.001,\n \"min_matches\" : 15,\n \"resize_to\" : 1024,\n }\n\n params_sift_lightglue = {\n \"num_features\" : 8192,\n \"detection_threshold\" : 0.001,\n \"min_matches\" : 15,\n \"resize_to\" : 1024,\n }\n\n params_loftr = {\n \"resize_small_edge_to\" : 750,\n \"min_matches\" : 15,\n }\n \n params_dkm = {\n \"num_features\" : 2048,\n \"detection_threshold\" : 0.4,\n \"min_matches\" : 15,\n \"resize_to\" : (540, 720), \n }\n \n # superpoint + superglue ... https://www.kaggle.com/competitions/image-matching-challenge-2023/discussion/416873\n params_sg1 = {\n \"sg_config\" : \n {\n \"superpoint\": {\n \"nms_radius\": 4, \n \"keypoint_threshold\": 0.005,\n \"max_keypoints\": -1,\n },\n \"superglue\": {\n \"weights\": \"outdoor\",\n \"sinkhorn_iterations\": 20,\n \"match_threshold\": 0.2,\n },\n },\n \"resize_to\": 1088,\n \"min_matches\": 15,\n }\n params_sg2 = {\n \"sg_config\" : \n {\n \"superpoint\": {\n \"nms_radius\": 4, \n \"keypoint_threshold\": 0.005,\n \"max_keypoints\": -1,\n },\n \"superglue\": {\n \"weights\": \"outdoor\",\n \"sinkhorn_iterations\": 20,\n \"match_threshold\": 0.2,\n },\n },\n \"resize_to\": 1280,\n \"min_matches\": 15,\n }\n params_sg3 = {\n \"sg_config\" : \n {\n \"superpoint\": {\n \"nms_radius\": 4, \n \"keypoint_threshold\": 0.005,\n \"max_keypoints\": -1,\n },\n \"superglue\": {\n \"weights\": \"outdoor\",\n \"sinkhorn_iterations\": 20,\n \"match_threshold\": 0.2,\n },\n },\n \"resize_to\": 1376,\n \"min_matches\": 15,\n }\n params_sgs = [params_sg1, params_sg2, params_sg3]\n \n params_matchformer = {\n \"detection_threshold\" : 0.15,\n \"resize_to\" : (560, 750),\n \"num_features\" : 2000,\n \"min_matches\" : 15, \n }","metadata":{"execution":{"iopub.status.busy":"2024-05-23T06:32:00.111073Z","iopub.execute_input":"2024-05-23T06:32:00.1115Z","iopub.status.idle":"2024-05-23T06:32:00.126919Z","shell.execute_reply.started":"2024-05-23T06:32:00.111471Z","shell.execute_reply":"2024-05-23T06:32:00.125984Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device=torch.device('cuda')","metadata":{"execution":{"iopub.status.busy":"2024-05-23T06:32:00.127887Z","iopub.execute_input":"2024-05-23T06:32:00.128125Z","iopub.status.idle":"2024-05-23T06:32:00.139901Z","shell.execute_reply.started":"2024-05-23T06:32:00.128105Z","shell.execute_reply":"2024-05-23T06:32:00.139139Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# COLMAP utilities","metadata":{}},{"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(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])\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\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)))","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-05-23T06:32:00.140928Z","iopub.execute_input":"2024-05-23T06:32:00.141194Z","iopub.status.idle":"2024-05-23T06:32:00.170293Z","shell.execute_reply.started":"2024-05-23T06:32:00.141173Z","shell.execute_reply":"2024-05-23T06:32:00.169416Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# h5 to colmap db","metadata":{}},{"cell_type":"code","source":"# Code to interface DISK with Colmap.\n# Forked 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 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","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-05-23T06:32:00.171347Z","iopub.execute_input":"2024-05-23T06:32:00.171614Z","iopub.status.idle":"2024-05-23T06:32:00.194977Z","shell.execute_reply.started":"2024-05-23T06:32:00.171593Z","shell.execute_reply":"2024-05-23T06:32:00.194082Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Rotation detection","metadata":{}},{"cell_type":"code","source":"from torchvision.io import read_image as T_read_image\nfrom torchvision.io import ImageReadMode\nfrom torchvision import transforms as T\nfrom check_orientation.pre_trained_models import create_model\n\ndef convert_rot_k(index):\n if index == 0:\n return 0\n elif index == 1:\n return 3\n elif index == 2:\n return 2\n else:\n return 1\n\nclass CheckRotationDataset(Dataset):\n def __init__(self, files, transform=None):\n self.transform = transform\n self.files = files\n\n def __len__(self):\n return len(self.files)\n\n def __getitem__(self, idx):\n imgPath = self.files[idx]\n image = T_read_image(imgPath, mode=ImageReadMode.RGB)\n if self.transform:\n image = self.transform(image)\n return image\n\ndef get_CheckRotation_dataloader(images, batch_size=1):\n transform = T.Compose([\n T.Resize((224, 224)),\n T.ConvertImageDtype(torch.float),\n T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))\n ])\n\n dataset = CheckRotationDataset(images, transform=transform)\n dataloader = DataLoader(\n dataset=dataset,\n shuffle=False,\n batch_size=batch_size,\n pin_memory=True,\n num_workers=2,\n drop_last=False\n )\n return dataloader\n\ndef exec_rotation_detection(img_files, device):\n model = create_model(\"swsl_resnext50_32x4d\")\n model.eval().to(device);\n \n dataloader = get_CheckRotation_dataloader(img_files)\n \n rots = []\n for idx, image in enumerate(dataloader):\n image = image.to(torch.float32).to(device)\n with torch.no_grad():\n prediction = model(image).detach().cpu().numpy()\n detected_rot = prediction[0].argmax()\n rot_k = convert_rot_k(detected_rot)\n rots.append(rot_k)\n print(f\"{os.path.basename(img_files[idx])} > rot_k={rot_k}\")\n return rots","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2024-05-23T06:32:00.196408Z","iopub.execute_input":"2024-05-23T06:32:00.196719Z","iopub.status.idle":"2024-05-23T06:32:00.872973Z","shell.execute_reply.started":"2024-05-23T06:32:00.196697Z","shell.execute_reply":"2024-05-23T06:32:00.872171Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Image Pairs","metadata":{}},{"cell_type":"code","source":"# We will use ViT global descriptor to get matching shortlists.\ndef get_global_desc(fnames, model,\n device = torch.device('cpu')):\n model = model.eval()\n model= model.to(device)\n config = resolve_data_config({}, model=model)\n transform = create_transform(**config)\n global_descs_convnext=[]\n for i, img_fname_full in tqdm(enumerate(fnames),total= len(fnames)):\n key = os.path.splitext(os.path.basename(img_fname_full))[0]\n img = Image.open(img_fname_full).convert('RGB')\n timg = transform(img).unsqueeze(0).to(device)\n with torch.no_grad(), torch.cuda.amp.autocast():\n desc = model.forward_features(timg.to(device)).mean(dim=(-1,2))#\n #print (desc.shape)\n desc = desc.view(1, -1)\n desc_norm = F.normalize(desc, dim=1, p=2)\n #print (desc_norm)\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.to(torch.float32)\n\ndef convert_1d_to_2d(idx, num_images):\n idx1 = idx // num_images\n idx2 = idx % num_images\n return (idx1, idx2)\n\ndef get_pairs_from_distancematrix(mat):\n pairs = [ convert_1d_to_2d(idx, mat.shape[0]) for idx in np.argsort(mat.flatten())]\n pairs = [ pair for pair in pairs if pair[0] < pair[1] ]\n return pairs\n\ndef get_img_pairs_exhaustive(img_fnames, model, device):\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 descs = get_global_desc(img_fnames, model, device=device)\n dm = torch.cdist(descs, descs, p=2).detach().cpu().numpy()\n matching_list = get_pairs_from_distancematrix(dm)\n return matching_list\n\n\ndef get_image_pairs_shortlist(fnames,\n sim_th = 0.6, # should be strict\n min_pairs = 20,\n exhaustive_if_less = 20,\n device=torch.device('cpu')):\n num_imgs = len(fnames)\n\n model = timm.create_model('tf_efficientnet_b7',\n checkpoint_path='/kaggle/input/tf-efficientnet/pytorch/tf-efficientnet-b7/1/tf_efficientnet_b7_ra-6c08e654.pth')\n model.eval()\n descs = get_global_desc(fnames, model, device=device)\n\n if num_imgs <= exhaustive_if_less:\n return get_img_pairs_exhaustive(fnames, model, device)\n \n dm = torch.cdist(descs, descs, p=2).detach().cpu().numpy()\n # removing half\n mask = dm <= sim_th\n total = 0\n matching_list = []\n ar = np.arange(num_imgs)\n already_there_set = []\n for st_idx in range(num_imgs-1):\n mask_idx = mask[st_idx]\n to_match = ar[mask_idx]\n if len(to_match) < min_pairs:\n to_match = np.argsort(dm[st_idx])[:min_pairs] \n for idx in to_match:\n if st_idx == idx:\n continue\n if dm[st_idx, idx] < 1000:\n matching_list.append(tuple(sorted((st_idx, idx.item()))))\n total+=1\n matching_list = sorted(list(set(matching_list)))\n return matching_list","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-05-23T06:32:00.876258Z","iopub.execute_input":"2024-05-23T06:32:00.876912Z","iopub.status.idle":"2024-05-23T06:32:00.896037Z","shell.execute_reply.started":"2024-05-23T06:32:00.876882Z","shell.execute_reply":"2024-05-23T06:32:00.894947Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Keypoints: LightGlue series","metadata":{}},{"cell_type":"code","source":"def 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\n #####################################################\n # Extract keypoints and descriptions\n #####################################################\n dict_model = {\n \"aliked\" : ALIKED,\n \"superpoint\" : SuperPoint,\n \"doghardnet\" : DoGHardNet,\n \"disk\" : DISK,\n \"sift\" : SIFT,\n }\n extractor_class = dict_model[model_name]\n \n dtype = torch.float32 # ALIKED has issues with float16\n extractor = extractor_class(\n max_num_keypoints=num_features, detection_threshold=detection_threshold, resize=resize_to\n ).eval().to(device, dtype)\n \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) # auto-resize the image, disable with resize=None\n kpts = feats0['keypoints'].reshape(-1, 2).detach()\n descs = feats0['descriptors'].reshape(len(kpts), -1).detach()\n kpts = convert_coord(kpts, w, h, rot_k)\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 kp1 = dict_kpts_cuda[key1]\n kp2 = dict_kpts_cuda[key2]\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()\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() -t \n print(f'Features matched in {t:.4f} sec ({model_name}+LightGlue)')\n return t\n","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-05-23T06:32:00.897627Z","iopub.execute_input":"2024-05-23T06:32:00.898335Z","iopub.status.idle":"2024-05-23T06:32:00.924563Z","shell.execute_reply.started":"2024-05-23T06:32:00.8983Z","shell.execute_reply":"2024-05-23T06:32:00.923633Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Keypoints: SuperGlue","metadata":{}},{"cell_type":"code","source":"import sys\nsys.path.append(\"../input/super-glue-pretrained-network\")\nfrom models.matching import Matching\nfrom models.superpoint import SuperPoint as SG_SuperPoint\nfrom models.superglue import SuperGlue\nfrom models.utils import (compute_pose_error, compute_epipolar_error,\n estimate_pose, make_matching_plot,\n error_colormap, AverageTimer, pose_auc, read_image,\n process_resize, frame2tensor,\n rotate_intrinsics, rotate_pose_inplane,\n scale_intrinsics)\n\nfrom torch.nn import functional as torchF # For resizing tensor\n\ndef sg_imread(path):\n image = cv2.imread(str(path), cv2.IMREAD_GRAYSCALE)\n return image\n\n# Preprocess\ndef sg_read_image(image, device, resize):\n w, h = image.shape[1], image.shape[0]\n w_new, h_new = process_resize(w, h, [resize,])\n \n unit_shape = 8\n w_new = w_new // unit_shape * unit_shape\n h_new = h_new // unit_shape * unit_shape\n \n scales = (float(w) / float(w_new), float(h) / float(h_new))\n image = cv2.resize(image.astype('float32'), (w_new, h_new))\n\n inp = frame2tensor(image, \"cpu\")\n return image, inp, scales, (h, w)\n\nclass SGDataset(Dataset):\n def __init__(self, img_fnames, resize_to, device):\n self.img_fnames = img_fnames\n self.resize_to = resize_to\n self.device = device\n \n def __len__(self):\n return len(self.img_fnames)\n \n def __getitem__(self, idx):\n fname = self.img_fnames[idx]\n im = cv2.imread(fname, cv2.IMREAD_GRAYSCALE)\n _, image, scale, ori_shape = sg_read_image(im, self.device, self.resize_to)\n return image, torch.tensor([idx]), torch.tensor(ori_shape)\n\ndef get_superglue_dataloader(img_fnames, resize_to, device, batch_size=1):\n dataset = SGDataset(img_fnames, resize_to, device)\n dataloader = DataLoader(\n dataset=dataset,\n shuffle=False,\n batch_size=batch_size,\n pin_memory=True,\n num_workers=2,\n drop_last=False\n )\n return dataloader\n\ndef detect_superglue(\n img_fnames, index_pairs, feature_dir, device, sg_config, file_keypoints, \n resize_to=750, min_matches=15\n): \n t=time()\n\n fnames1, fnames2, idxs1, idxs2 = [], [], [], []\n for pair_idx in progress_bar(index_pairs):\n idx1, idx2 = pair_idx\n fname1, fname2 = img_fnames[idx1], img_fnames[idx2]\n fnames1.append(fname1)\n fnames2.append(fname2)\n idxs1.append(idx1)\n idxs2.append(idx2)\n \n dataloader = get_superglue_dataloader( img_fnames, resize_to, device)\n\n #####################################################\n # Extract keypoints and descriptions\n #####################################################\n superpoint = SG_SuperPoint(sg_config[\"superpoint\"]).eval().to(device)\n dict_features_cuda = {}\n dict_shapes = {}\n dict_images = {}\n for X in dataloader:\n image, idx, ori_shape = X\n image = image[0].to(device)\n fname = img_fnames[idx]\n key = fname.split('/')[-1]\n \n with torch.no_grad(), torch.autocast(device_type=\"cuda\", dtype=torch.float16):\n pred = superpoint({'image': image})\n dict_features_cuda[key] = pred\n dict_shapes[key] = ori_shape\n dict_images[key] = image.half()\n del superpoint\n gc.collect()\n \n #####################################################\n # Matching keypoints\n #####################################################\n superglue = SuperGlue(sg_config[\"superglue\"]).eval().to(device)\n weights = sg_config[\"superglue\"][\"weights\"]\n cnt_pairs = 0\n \n with h5py.File(file_keypoints, mode='w') as f_match:\n for idx, (fname1, fname2) in enumerate(zip(fnames1, fnames2)):\n key1, key2 = fname1.split('/')[-1], fname2.split('/')[-1]\n\n data = {\"image0\": dict_images[key1], \"image1\": dict_images[key2]}\n data = {**data, **{k+'0': v for k, v in dict_features_cuda[key1].items()}}\n data = {**data, **{k+'1': v for k, v in dict_features_cuda[key2].items()}}\n for k in data:\n if isinstance(data[k], (list, tuple)):\n data[k] = torch.stack(data[k])\n with torch.no_grad(), torch.autocast(device_type=\"cuda\", dtype=torch.float16):\n pred = {**data, **superglue(data)}\n pred = {k: v[0].detach().cpu().numpy().copy() for k, v in pred.items()}\n mkpts1, mkpts2 = pred[\"keypoints0\"], pred[\"keypoints1\"]\n matches, conf = pred[\"matches0\"], pred[\"matching_scores0\"]\n\n valid = matches > -1\n mkpts1 = mkpts1[valid]\n mkpts2 = mkpts2[matches[valid]]\n mconf = conf[valid]\n\n ori_shape_1 = dict_shapes[key1][0].numpy()\n ori_shape_2 = dict_shapes[key2][0].numpy()\n \n # Scaling coords\n mkpts1[:,0] = mkpts1[:,0] * ori_shape_1[1] / dict_images[key1].shape[3] # X\n mkpts1[:,1] = mkpts1[:,1] * ori_shape_1[0] / dict_images[key1].shape[2] # Y\n mkpts2[:,0] = mkpts2[:,0] * ori_shape_2[1] / dict_images[key2].shape[3] # X\n mkpts2[:,1] = mkpts2[:,1] * ori_shape_2[0] / dict_images[key2].shape[2] # Y \n \n n_matches = mconf.shape[0]\n \n group = f_match.require_group(key1)\n if n_matches >= min_matches:\n group.create_dataset(key2, data=np.concatenate([mkpts1, mkpts2], axis=1).astype(np.float32))\n cnt_pairs+=1\n print (f'{key1}-{key2}: {n_matches} matches @ {cnt_pairs}th pair(superglue/{resize_to}/{weights})') \n else:\n print (f'{key1}-{key2}: {n_matches} matches --> skipped') \n\n del superglue\n del dict_features_cuda\n del dict_images\n torch.cuda.empty_cache()\n gc.collect()\n t=time() -t \n print(f'Features matched in {t:.4f} sec')\n return t","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-05-23T06:32:00.925877Z","iopub.execute_input":"2024-05-23T06:32:00.926171Z","iopub.status.idle":"2024-05-23T06:32:00.992149Z","shell.execute_reply.started":"2024-05-23T06:32:00.926138Z","shell.execute_reply":"2024-05-23T06:32:00.991417Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Keypoints: DKM","metadata":{}},{"cell_type":"code","source":"class DKMDataset(Dataset):\n def __init__(self, fnames1, fnames2, resize_to, device):\n self.fnames1 = fnames1\n self.fnames2 = fnames2\n self.resize_to = resize_to\n self.device = device\n self.test_transform = get_tuple_transform_ops(\n resize=self.resize_to, normalize=True\n )\n\n \n def __len__(self):\n return len(self.fnames1)\n \n def __getitem__(self, idx):\n fname1 = self.fnames1[idx]\n fname2 = self.fnames2[idx]\n \n im1, im2 = Image.open(fname1), Image.open(fname2)\n ori_shape_1 = im1.size\n ori_shape_2 = im2.size\n image1, image2 = self.test_transform((im1, im2))\n return image1, image2, torch.tensor([idx]), torch.tensor(ori_shape_1), torch.tensor(ori_shape_2)\n\ndef get_dkm_dataloader(images1, images2, resize_to, device, batch_size=4):\n dataset = DKMDataset(images1, images2, resize_to, device)\n dataloader = DataLoader(\n dataset=dataset,\n shuffle=False,\n batch_size=batch_size,\n pin_memory=True,\n num_workers=2,\n drop_last=False\n )\n return dataloader\n\ndef get_dkm_mkpts(dkm_model, bimgs1, bimgs2, shapes1, shapes2, detection_threshold=0.5, num_features = 2000, min_matches=15):\n dense_matches, dense_certainty = dkm_model.match(bimgs1, bimgs2, batched=True)\n print(\"***\", dense_matches.shape, dense_certainty.shape)\n\n store_mkpts1, store_mkpts2, store_mconf = [], [], []\n # drop low confidence pairs\n for b in range(dense_matches.shape[0]):\n u_dense_matches = dense_matches[b, dense_certainty[b,...].sqrt() >= detection_threshold, :]\n u_dense_certainty = dense_certainty[b, dense_certainty[b,...].sqrt() >= detection_threshold]\n \n if u_dense_matches.shape[0] > num_features:\n u_dense_matches, u_dense_certainty = dkm_model.sample( u_dense_matches, u_dense_certainty, num=num_features)\n \n u_dense_matches = u_dense_matches.reshape((-1, 4))\n u_dense_certainty = u_dense_certainty.reshape((-1,))\n \n mkpts1 = u_dense_matches[:, :2]\n mkpts2 = u_dense_matches[:, 2:]\n \n w1, h1 = shapes1[b, :]\n w2, h2 = shapes2[b, :]\n\n mkpts1[:, 0] = ((mkpts1[:, 0] + 1)/2) * w1\n mkpts1[:, 1] = ((mkpts1[:, 1] + 1)/2) * h1\n\n mkpts2[:, 0] = ((mkpts2[:, 0] + 1)/2) * w2\n mkpts2[:, 1] = ((mkpts2[:, 1] + 1)/2) * h2\n\n mkpts1 = mkpts1.cpu().detach().numpy()\n mkpts2 = mkpts2.cpu().detach().numpy()\n mconf = u_dense_certainty.sqrt().cpu().detach().numpy()\n\n \n if mconf.shape[0] > min_matches:\n try:\n # calc Fundamental matrix from keypoints\n F, inliers = cv2.findFundamentalMat(mkpts1, mkpts2, cv2.USAC_MAGSAC, 0.200, 0.999, 2000)\n inliers = inliers > 0\n mkpts1 = mkpts1[inliers[:,0]]\n mkpts2 = mkpts2[inliers[:,0]]\n mconf = mconf[inliers[:,0]]\n #print(\"---\", mconf.shape)\n if mconf.shape[0] > 3000:\n rand_idx = np.random.choice(range(mconf.shape[0]), 3000, replace=False)\n mkpts1 = mkpts1[rand_idx, :]\n mkpts2 = mkpts2[rand_idx, :]\n mconf = mconf[rand_idx]\n except:\n mkpts1 = np.empty((0,2))\n mkpts2 = np.empty((0,2))\n mconf = np.empty((0,))\n \n store_mkpts1.append(mkpts1)\n store_mkpts2.append(mkpts2)\n store_mconf.append(mconf)\n return store_mkpts1, store_mkpts2, store_mconf\n\ndef detect_dkm(\n img_fnames, index_pairs, feature_dir, device, \n resize_to=(540, 720), \n detection_threshold=0.4, \n num_features=2000, \n min_matches=15,\n):\n t=time()\n dkm_model = DKMv3_outdoor(device=device)\n dkm_model.upsample_preds=False\n\n fnames1, fnames2 = [], []\n for pair_idx in progress_bar(index_pairs):\n idx1, idx2 = pair_idx\n fname1, fname2 = img_fnames[idx1], img_fnames[idx2]\n fnames1.append(fname1)\n fnames2.append(fname2)\n \n cnt_pairs = 0\n with h5py.File(f'{feature_dir}/matches_dkm.h5', mode='w') as f_match: \n dataloader = get_dkm_dataloader(fnames1, fnames2, resize_to, device, batch_size=4)\n for X in tqdm(dataloader):\n images1, images2, idxs, shapes1, shapes2 = X\n store_mkpts1, store_mkpts2, store_mconf = get_dkm_mkpts(\n dkm_model, images1.to(device), images2.to(device), shapes1, shapes2, \n detection_threshold=detection_threshold, num_features = num_features, min_matches=min_matches,\n )\n \n for b in range(images1.shape[0]):\n mkpts1 = store_mkpts1[b]\n mkpts2 = store_mkpts2[b]\n mconf = store_mconf[b]\n file1 = fnames1[idxs[b]]\n file2 = fnames2[idxs[b]]\n key1, key2 = file1.split('/')[-1], file2.split('/')[-1]\n \n n_matches = mconf.shape[0]\n print (f'{key1}-{key2}: {n_matches} matches @ {cnt_pairs}th pair(dkm)') \n\n group = f_match.require_group(key1)\n if n_matches >= min_matches:\n group.create_dataset(key2, data=np.concatenate([mkpts1, mkpts2], axis=1).astype(np.float32))\n cnt_pairs+=1\n gc.collect()\n t=time() -t \n print(f'Features matched in {t:.4f} sec')\n return t","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-05-23T06:32:00.993363Z","iopub.execute_input":"2024-05-23T06:32:00.993812Z","iopub.status.idle":"2024-05-23T06:32:01.022953Z","shell.execute_reply.started":"2024-05-23T06:32:00.993779Z","shell.execute_reply":"2024-05-23T06:32:01.022123Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Keypoints: LoFTR","metadata":{}},{"cell_type":"code","source":"class LoFTRDataset(Dataset):\n def __init__(self, fnames1, fnames2, idxs1, idxs2, resize_small_edge_to, device):\n self.fnames1 = fnames1\n self.fnames2 = fnames2\n self.keys1 = [ fname.split('/')[-1] for fname in fnames1 ]\n self.keys2 = [ fname.split('/')[-1] for fname in fnames2 ]\n self.idxs1 = idxs1\n self.idxs2 = idxs2\n self.resize_small_edge_to = resize_small_edge_to\n self.device = device\n self.round_unit = 16\n \n def __len__(self):\n return len(self.images1)\n\n def load_torch_image(self, fname, device):\n img = cv2.imread(fname)\n original_shape = img.shape\n ratio = self.resize_small_edge_to / min([img.shape[0], img.shape[1]])\n w = int(img.shape[1] * ratio) # int( (img.shape[1] * ratio) // self.round_unit * self.round_unit )\n h = int(img.shape[0] * ratio) # int( (img.shape[0] * ratio) // self.round_unit * self.round_unit )\n img_resized = cv2.resize(img, (w, h))\n img_resized = K.image_to_tensor(img_resized, False).float() /255.\n img_resized = K.color.bgr_to_rgb(img_resized)\n img_resized = K.color.rgb_to_grayscale(img_resized)\n return img_resized.to(device), original_shape\n \n def __getitem__(self, idx):\n fname1 = self.fnames1[idx]\n fname2 = self.fnames2[idx]\n image1, ori_shape_1 = self.load_torch_image(fname1, device)\n image2, ori_shape_2 = self.load_torch_image(fname2, device)\n\n return image1, image2, self.keys1[idx], self.keys2[idx], self.idxs1[idx], self.idxs2[idx], ori_shape_1, ori_shape_2\n\ndef get_loftr_dataloader(images1, images2, idxs1, idxs2, resize_small_edge_to, device, batch_size=1):\n dataset = LoFTRDataset(images1, images2, idxs1, idxs2, resize_small_edge_to, device)\n dataloader = DataLoader(\n dataset=dataset,\n shuffle=False,\n batch_size=batch_size,\n pin_memory=True,\n num_workers=2,\n drop_last=False\n )\n return dataset\n \ndef detect_loftr(img_fnames, index_pairs, feature_dir, device, file_keypoints, resize_small_edge_to=750, min_matches=15):\n t=time()\n\n matcher = LoFTR(pretrained=None)\n matcher.load_state_dict(torch.load(\"../input/loftr/pytorch/outdoor/1/loftr_outdoor.ckpt\")['state_dict'])\n matcher = matcher.to(device).eval()\n\n fnames1, fnames2, idxs1, idxs2 = [], [], [], []\n for pair_idx in progress_bar(index_pairs):\n idx1, idx2 = pair_idx\n fname1, fname2 = img_fnames[idx1], img_fnames[idx2]\n fnames1.append(fname1)\n fnames2.append(fname2)\n idxs1.append(idx1)\n idxs2.append(idx2)\n \n \n dataloader = get_loftr_dataloader( fnames1, fnames2, idxs1, idxs2, resize_small_edge_to, device)\n\n cnt_pairs = 0\n\n with h5py.File(file_keypoints, mode='w') as f_match: \n store_mkpts = {}\n for X in tqdm(dataloader):\n image1, image2, key1, key2, idx1, idx2, ori_shape_1, ori_shape_2 = X\n fname1, fname2 = img_fnames[idx1], img_fnames[idx2]\n\n with torch.no_grad():\n correspondences = matcher( {\"image0\": image1.to(device),\"image1\": image2.to(device)} )\n mkpts1 = correspondences['keypoints0'].cpu().numpy()\n mkpts2 = correspondences['keypoints1'].cpu().numpy()\n mconf = correspondences['confidence'].cpu().numpy()\n\n mkpts1[:,0] *= (float(ori_shape_1[1]) / float(image1.shape[3]))\n mkpts1[:,1] *= (float(ori_shape_1[0]) / float(image1.shape[2]))\n\n mkpts2[:,0] *= (float(ori_shape_2[1]) / float(image2.shape[3]))\n mkpts2[:,1] *= (float(ori_shape_2[0]) / float(image2.shape[2]))\n \n n_matches = mconf.shape[0]\n \n group = f_match.require_group(key1)\n if n_matches >= min_matches:\n group.create_dataset(key2, data=np.concatenate([mkpts1, mkpts2], axis=1).astype(np.float32))\n cnt_pairs+=1\n print (f'{key1}-{key2}: {n_matches} matches @ {cnt_pairs}th pair(loftr)')\n else:\n print (f'{key1}-{key2}: {n_matches} matches --> skipped')\n gc.collect()\n t=time() -t \n print(f'Features matched in {t:.4f} sec')\n return t","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-05-23T06:32:01.024428Z","iopub.execute_input":"2024-05-23T06:32:01.024726Z","iopub.status.idle":"2024-05-23T06:32:01.048043Z","shell.execute_reply.started":"2024-05-23T06:32:01.024704Z","shell.execute_reply":"2024-05-23T06:32:01.047157Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Keypoints: DKM","metadata":{}},{"cell_type":"code","source":"class DKMDataset(Dataset):\n def __init__(self, fnames1, fnames2, resize_to, device):\n self.fnames1 = fnames1\n self.fnames2 = fnames2\n self.resize_to = resize_to\n self.device = device\n self.test_transform = get_tuple_transform_ops(\n resize=self.resize_to, normalize=True\n )\n\n \n def __len__(self):\n return len(self.fnames1)\n \n def __getitem__(self, idx):\n fname1 = self.fnames1[idx]\n fname2 = self.fnames2[idx]\n \n im1, im2 = Image.open(fname1), Image.open(fname2)\n ori_shape_1 = im1.size\n ori_shape_2 = im2.size\n image1, image2 = self.test_transform((im1, im2))\n return image1, image2, torch.tensor([idx]), torch.tensor(ori_shape_1), torch.tensor(ori_shape_2)\n\ndef get_dkm_dataloader(images1, images2, resize_to, device, batch_size=4):\n dataset = DKMDataset(images1, images2, resize_to, device)\n dataloader = DataLoader(\n dataset=dataset,\n shuffle=False,\n batch_size=batch_size,\n pin_memory=True,\n num_workers=2,\n drop_last=False\n )\n return dataloader\n\ndef get_dkm_mkpts(dkm_model, bimgs1, bimgs2, shapes1, shapes2, detection_threshold=0.5, num_features = 2000, min_matches=15):\n dense_matches, dense_certainty = dkm_model.match(bimgs1, bimgs2, batched=True)\n\n store_mkpts1, store_mkpts2, store_mconf = [], [], []\n # drop low confidence pairs\n for b in range(dense_matches.shape[0]):\n u_dense_matches = dense_matches[b, dense_certainty[b,...].sqrt() >= detection_threshold, :]\n u_dense_certainty = dense_certainty[b, dense_certainty[b,...].sqrt() >= detection_threshold]\n \n if u_dense_matches.shape[0] > num_features:\n u_dense_matches, u_dense_certainty = dkm_model.sample( u_dense_matches, u_dense_certainty, num=num_features)\n \n u_dense_matches = u_dense_matches.reshape((-1, 4))\n u_dense_certainty = u_dense_certainty.reshape((-1,))\n \n mkpts1 = u_dense_matches[:, :2]\n mkpts2 = u_dense_matches[:, 2:]\n \n w1, h1 = shapes1[b, :]\n w2, h2 = shapes2[b, :]\n\n mkpts1[:, 0] = ((mkpts1[:, 0] + 1)/2) * w1\n mkpts1[:, 1] = ((mkpts1[:, 1] + 1)/2) * h1\n\n mkpts2[:, 0] = ((mkpts2[:, 0] + 1)/2) * w2\n mkpts2[:, 1] = ((mkpts2[:, 1] + 1)/2) * h2\n\n mkpts1 = mkpts1.cpu().detach().numpy()\n mkpts2 = mkpts2.cpu().detach().numpy()\n mconf = u_dense_certainty.sqrt().cpu().detach().numpy()\n\n if mconf.shape[0] > min_matches:\n try:\n # calc Fundamental matrix from keypoints\n F, inliers = cv2.findFundamentalMat(mkpts1, mkpts2, cv2.USAC_MAGSAC, 0.200, 0.999, 2000)\n inliers = inliers > 0\n mkpts1 = mkpts1[inliers[:,0]]\n mkpts2 = mkpts2[inliers[:,0]]\n mconf = mconf[inliers[:,0]]\n except:\n pass\n store_mkpts1.append(mkpts1)\n store_mkpts2.append(mkpts2)\n store_mconf.append(mconf)\n return store_mkpts1, store_mkpts2, store_mconf\n\ndef detect_dkm(\n img_fnames, index_pairs, feature_dir, device, file_keypoints,\n resize_to=(540, 720), \n detection_threshold=0.4, \n num_features=2000, \n min_matches=15\n):\n t=time()\n dkm_model = DKMv3_outdoor(device=device)\n dkm_model.upsample_preds=False\n\n fnames1, fnames2 = [], []\n for pair_idx in progress_bar(index_pairs):\n idx1, idx2 = pair_idx\n fname1, fname2 = img_fnames[idx1], img_fnames[idx2]\n fnames1.append(fname1)\n fnames2.append(fname2)\n \n cnt_pairs = 0\n with h5py.File(file_keypoints, mode='w') as f_match: \n dataloader = get_dkm_dataloader(fnames1, fnames2, resize_to, device, batch_size=4)\n for X in tqdm(dataloader):\n images1, images2, idxs, shapes1, shapes2 = X\n store_mkpts1, store_mkpts2, store_mconf = get_dkm_mkpts(\n dkm_model, images1.to(device), images2.to(device), shapes1, shapes2, \n detection_threshold=detection_threshold, num_features = num_features, min_matches=min_matches,\n )\n \n for b in range(images1.shape[0]):\n mkpts1 = store_mkpts1[b]\n mkpts2 = store_mkpts2[b]\n mconf = store_mconf[b]\n file1 = fnames1[idxs[b]]\n file2 = fnames2[idxs[b]]\n key1, key2 = file1.split('/')[-1], file2.split('/')[-1]\n \n n_matches = mconf.shape[0]\n\n group = f_match.require_group(key1)\n if n_matches >= min_matches:\n group.create_dataset(key2, data=np.concatenate([mkpts1, mkpts2], axis=1).astype(np.float32))\n cnt_pairs+=1\n print (f'{key1}-{key2}: {n_matches} matches @ {cnt_pairs}th pair(dkm)')\n else:\n print (f'{key1}-{key2}: {n_matches} matches --> skipped')\n\n gc.collect()\n t=time() -t \n print(f'Features matched in {t:.4f} sec')\n return t","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-05-23T06:32:01.049485Z","iopub.execute_input":"2024-05-23T06:32:01.049834Z","iopub.status.idle":"2024-05-23T06:32:01.077883Z","shell.execute_reply.started":"2024-05-23T06:32:01.049804Z","shell.execute_reply":"2024-05-23T06:32:01.076955Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Keypoints: MatchFormer","metadata":{}},{"cell_type":"code","source":"class MatchFormerDataset(Dataset):\n def __init__(self, fnames1, fnames2, idxs1, idxs2, resize_to, device):\n self.fnames1 = fnames1\n self.fnames2 = fnames2\n self.keys1 = [ fname.split('/')[-1] for fname in fnames1 ]\n self.keys2 = [ fname.split('/')[-1] for fname in fnames2 ]\n self.idxs1 = idxs1\n self.idxs2 = idxs2\n self.resize_to = resize_to\n self.device = device\n self.round_unit = 16\n \n def __len__(self):\n return len(self.images1)\n\n def load_torch_image(self, fname, device):\n img = cv2.imread(fname)\n original_shape = img.shape\n #ratio = self.resize_long_edge_to / max([img.shape[0], img.shape[1]])\n #w = int(img.shape[1] * ratio)\n #h = int(img.shape[0] * ratio)\n img_resized = cv2.resize(img, self.resize_to)\n img_resized = K.image_to_tensor(img_resized, False).float() /255.\n img_resized = K.color.bgr_to_rgb(img_resized)\n img_resized = K.color.rgb_to_grayscale(img_resized)\n return img_resized.to(device), original_shape\n \n def __getitem__(self, idx):\n fname1 = self.fnames1[idx]\n fname2 = self.fnames2[idx]\n image1, ori_shape_1 = self.load_torch_image(fname1, device)\n image2, ori_shape_2 = self.load_torch_image(fname2, device)\n\n return image1, image2, self.keys1[idx], self.keys2[idx], self.idxs1[idx], self.idxs2[idx], ori_shape_1, ori_shape_2\n\ndef get_matchformer_dataloader(images1, images2, idxs1, idxs2, resize_to, device, batch_size=1):\n dataset = MatchFormerDataset(images1, images2, idxs1, idxs2, resize_to, device)\n dataloader = DataLoader(\n dataset=dataset,\n shuffle=False,\n batch_size=batch_size,\n pin_memory=True,\n num_workers=2,\n drop_last=False\n )\n return dataset\n \ndef detect_matchformer(\n img_fnames, index_pairs, feature_dir, device, file_keypoints,\n resize_to=(560, 750), \n detection_threshold=0.4, \n num_features=2000, \n min_matches=15\n):\n t=time()\n\n sys.path.append('/kaggle/input/matchformer/MatchFormer-main')\n\n from yacs.config import CfgNode as CN\n from model.matchformer import Matchformer\n from config import defaultmf\n\n cfg = defaultmf.get_cfg_defaults()\n cfg.MATCHFORMER.BACKBONE_TYPE = 'largela'\n cfg.MATCHFORMER.SCENS = 'outdoor'\n cfg.MATCHFORMER.RESOLUTION = (8,2)\n cfg.MATCHFORMER.MATCH_COARSE.THR = detection_threshold\n\n def lower_config(yacs_cfg):\n if not isinstance(yacs_cfg, CN):\n return yacs_cfg\n return {k.lower(): lower_config(v) for k, v in yacs_cfg.items()}\n\n _cfg = lower_config(cfg)\n\n matcher_mf = Matchformer(_cfg['matchformer'])\n\n pretrained_ckpt = '/kaggle/input/matchformer/outdoor-large-LA.ckpt'\n matcher_mf.load_state_dict({k.replace('matcher.',''):v for k,v in torch.load(pretrained_ckpt, map_location='cpu').items()})\n matcher_mf = matcher_mf.to(device).eval()\n \n \n fnames1, fnames2, idxs1, idxs2 = [], [], [], []\n for pair_idx in progress_bar(index_pairs):\n idx1, idx2 = pair_idx\n fname1, fname2 = img_fnames[idx1], img_fnames[idx2]\n fnames1.append(fname1)\n fnames2.append(fname2)\n idxs1.append(idx1)\n idxs2.append(idx2)\n \n cnt_pairs = 0\n with h5py.File(file_keypoints, mode='w') as f_match: \n dataloader = get_matchformer_dataloader(fnames1, fnames2, idxs1, idxs2, resize_to, device, batch_size=1)\n for X in tqdm(dataloader):\n image1, image2, key1, key2, idx1, idx2, ori_shape_1, ori_shape_2 = X\n fname1, fname2 = img_fnames[idx1], img_fnames[idx2]\n #print(image1.shape, image2.shape)\n input_dict = {\n \"image0\": image1, \n \"image1\": image2\n }\n\n with torch.inference_mode():\n matcher_mf(input_dict)\n\n conf = input_dict['mconf'].to('cpu').numpy()\n mkpts1 = input_dict['mkpts0_f'].to('cpu').numpy()\n mkpts2 = input_dict['mkpts1_f'].to('cpu').numpy()\n\n sorted_idx = np.argsort(-conf)\n if len(conf) > num_features:\n mkpts1 = mkpts1[sorted_idx[:num_features], :]\n mkpts2 = mkpts2[sorted_idx[:num_features], :]\n\n mkpts1[:,0] = mkpts1[:,0] * ori_shape_1[1] / image1.shape[3]\n mkpts1[:,1] = mkpts1[:,1] * ori_shape_1[0] / image1.shape[2]\n\n mkpts2[:,0] = mkpts2[:,0] * ori_shape_2[1] / image2.shape[3]\n mkpts2[:,1] = mkpts2[:,1] * ori_shape_2[0] / image2.shape[2]\n \n n_matches = mkpts1.shape[0]\n\n group = f_match.require_group(key1)\n if n_matches >= min_matches:\n group.create_dataset(key2, data=np.concatenate([mkpts1, mkpts2], axis=1).astype(np.float32))\n cnt_pairs+=1\n print (f'{key1}-{key2}: {n_matches} matches @ {cnt_pairs}th pair(MatchFormer)')\n else:\n print (f'{key1}-{key2}: {n_matches} matches --> skipped')\n\n gc.collect()\n t=time() -t \n print(f'Features matched in {t:.4f} sec')\n return t","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-05-23T06:32:01.079081Z","iopub.execute_input":"2024-05-23T06:32:01.079389Z","iopub.status.idle":"2024-05-23T06:32:01.106429Z","shell.execute_reply.started":"2024-05-23T06:32:01.079365Z","shell.execute_reply":"2024-05-23T06:32:01.105512Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Keypoints merger","metadata":{}},{"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(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 #print(inliers.shape, inlier_idxs[:5])\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 #print(f\"filter_FundamentalMatrix: {len(store_inliers)} matches --> {mkpts.shape[0]} matches\")\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","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-05-23T06:32:01.107845Z","iopub.execute_input":"2024-05-23T06:32:01.108157Z","iopub.status.idle":"2024-05-23T06:32:01.160038Z","shell.execute_reply.started":"2024-05-23T06:32:01.108133Z","shell.execute_reply":"2024-05-23T06:32:01.158959Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Keypoints wrapper function","metadata":{}},{"cell_type":"code","source":"def wrapper_keypoints(\n img_fnames, index_pairs, feature_dir, device, timings, rots\n):\n #############################################################\n # get keypoints\n #############################################################\n files_keypoints = []\n \n if CONFIG.use_superglue:\n for params_sg in CONFIG.params_sgs:\n resize_to = params_sg[\"resize_to\"]\n file_keypoints = f\"{feature_dir}/matches_superglue_{resize_to}pix.h5\"\n !rm -rf {file_keypoints}\n t = detect_superglue(\n img_fnames, index_pairs, feature_dir, device, \n params_sg[\"sg_config\"], file_keypoints, \n resize_to=params_sg[\"resize_to\"], \n min_matches=params_sg[\"min_matches\"],\n )\n gc.collect()\n files_keypoints.append( file_keypoints )\n timings['feature_matching'].append(t)\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 if CONFIG.use_doghardnet_lightglue:\n model_name = \"doghardnet\"\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_doghardnet_lightglue[\"resize_to\"],\n detection_threshold=CONFIG.params_doghardnet_lightglue[\"detection_threshold\"],\n num_features=CONFIG.params_doghardnet_lightglue[\"num_features\"],\n min_matches=CONFIG.params_doghardnet_lightglue[\"min_matches\"],\n )\n gc.collect()\n files_keypoints.append(file_keypoints)\n timings['feature_matching'].append(t)\n\n if CONFIG.use_superpoint_lightglue:\n model_name = \"superpoint\"\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_superpoint_lightglue[\"resize_to\"],\n detection_threshold=CONFIG.params_superpoint_lightglue[\"detection_threshold\"],\n num_features=CONFIG.params_superpoint_lightglue[\"num_features\"],\n min_matches=CONFIG.params_superpoint_lightglue[\"min_matches\"],\n )\n gc.collect()\n files_keypoints.append(file_keypoints)\n timings['feature_matching'].append(t)\n\n if CONFIG.use_disk_lightglue:\n model_name = \"disk\"\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_disk_lightglue[\"resize_to\"],\n detection_threshold=CONFIG.params_disk_lightglue[\"detection_threshold\"],\n num_features=CONFIG.params_disk_lightglue[\"num_features\"],\n min_matches=CONFIG.params_disk_lightglue[\"min_matches\"],\n )\n gc.collect()\n files_keypoints.append(file_keypoints)\n timings['feature_matching'].append(t)\n\n if CONFIG.use_sift_lightglue:\n model_name = \"sift\"\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_sift_lightglue[\"resize_to\"],\n detection_threshold=CONFIG.params_sift_lightglue[\"detection_threshold\"],\n num_features=CONFIG.params_sift_lightglue[\"num_features\"],\n min_matches=CONFIG.params_sift_lightglue[\"min_matches\"],\n )\n gc.collect()\n files_keypoints.append(file_keypoints)\n timings['feature_matching'].append(t)\n\n if CONFIG.use_loftr:\n file_keypoints = f'{feature_dir}/matches_loftr_{CONFIG.params_loftr[\"resize_small_edge_to\"]}pix.h5'\n t = detect_loftr(\n img_fnames, index_pairs, feature_dir, device, file_keypoints,\n resize_small_edge_to=CONFIG.params_loftr[\"resize_small_edge_to\"],\n min_matches=CONFIG.params_loftr[\"min_matches\"],\n )\n gc.collect()\n files_keypoints.append( file_keypoints )\n timings['feature_matching'].append(t)\n\n if CONFIG.use_dkm:\n file_keypoints = f'{feature_dir}/matches_dkm.h5'\n t = detect_dkm(\n img_fnames, index_pairs, feature_dir, device, file_keypoints,\n resize_to=CONFIG.params_dkm[\"resize_to\"], \n detection_threshold=CONFIG.params_dkm[\"detection_threshold\"], \n num_features=CONFIG.params_dkm[\"num_features\"], \n min_matches=CONFIG.params_dkm[\"min_matches\"]\n )\n gc.collect()\n files_keypoints.append(file_keypoints)\n timings['feature_matching'].append(t)\n\n if CONFIG.use_matchformer:\n file_keypoints = f'{feature_dir}/matches_matchformer_{CONFIG.params_matchformer[\"resize_to\"]}pix.h5'\n t = detect_matchformer(\n img_fnames, index_pairs, feature_dir, device, file_keypoints,\n resize_to=CONFIG.params_matchformer[\"resize_to\"],\n num_features=CONFIG.params_matchformer[\"num_features\"], \n min_matches=CONFIG.params_matchformer[\"min_matches\"]\n )\n gc.collect()\n files_keypoints.append( file_keypoints )\n timings['feature_matching'].append(t)\n\n #############################################################\n # merge keypoints\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","metadata":{"execution":{"iopub.status.busy":"2024-05-23T06:32:01.161352Z","iopub.execute_input":"2024-05-23T06:32:01.161983Z","iopub.status.idle":"2024-05-23T06:32:01.19864Z","shell.execute_reply.started":"2024-05-23T06:32:01.161955Z","shell.execute_reply":"2024-05-23T06:32:01.197663Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Reconstruction wrapper function","metadata":{}},{"cell_type":"code","source":"def 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()\n options = pycolmap.SiftMatchingOptions()\n options.confidence = 0.9999\n options.max_num_trials = 20000\n pycolmap.match_exhaustive(database_path, sift_options=options)\n t=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()\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 print(maps)\n clear_output(wait=False)\n t=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 scene_result[key1] = {}\n scene_result[key1][\"R\"] = deepcopy(im.rotmat())\n scene_result[key1][\"t\"] = deepcopy(np.array(im.tvec))\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 return scene_result","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-05-23T06:32:01.199834Z","iopub.execute_input":"2024-05-23T06:32:01.200157Z","iopub.status.idle":"2024-05-23T06:32:01.238745Z","shell.execute_reply.started":"2024-05-23T06:32:01.20012Z","shell.execute_reply":"2024-05-23T06:32:01.237742Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Submission utilities","metadata":{}},{"cell_type":"code","source":"def 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 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(f'{image},{dataset},{scene},{arr_to_str(R)},{arr_to_str(T)}\\n')","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-05-23T06:32:01.240125Z","iopub.execute_input":"2024-05-23T06:32:01.240488Z","iopub.status.idle":"2024-05-23T06:32:01.32479Z","shell.execute_reply.started":"2024-05-23T06:32:01.240439Z","shell.execute_reply":"2024-05-23T06:32:01.323765Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Main","metadata":{}},{"cell_type":"code","source":"src = '/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')","metadata":{"execution":{"iopub.status.busy":"2024-05-23T06:32:01.326195Z","iopub.execute_input":"2024-05-23T06:32:01.326576Z","iopub.status.idle":"2024-05-23T06:32:01.45532Z","shell.execute_reply.started":"2024-05-23T06:32:01.326543Z","shell.execute_reply":"2024-05-23T06:32:01.45437Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"out_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 print(dataset)\n if dataset not in out_results:\n out_results[dataset] = {}\n for scene in data_dict[dataset]:\n print(scene)\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}/test/{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\n #############################################################\n # get image rotations\n #############################################################\n t = time()\n if CONFIG.ROTATION_CORRECTION:\n rots = exec_rotation_detection(img_fnames, device)\n else:\n rots = [ 0 for fname in img_fnames ]\n t = 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 \n #############################################################\n # get image pairs\n #############################################################\n t=time()\n index_pairs = get_image_pairs_shortlist(img_fnames,\n sim_th = 1.0, # should be strict\n min_pairs = 50, # we select at least min_pairs PER IMAGE with biggest similarity\n exhaustive_if_less = 50,\n device=device)\n t=time() -t \n timings['shortlisting'].append(t)\n print (f'{len(index_pairs)}, pairs to match, {t:.4f} sec')\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 timings['feature_matching'] = keypoints_timings['feature_matching']\n gc.collect()\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 \n create_submission(out_results, data_dict)\n gc.collect()","metadata":{"execution":{"iopub.status.busy":"2024-05-23T06:32:01.456961Z","iopub.execute_input":"2024-05-23T06:32:01.457527Z","iopub.status.idle":"2024-05-23T06:33:51.762819Z","shell.execute_reply.started":"2024-05-23T06:32:01.457493Z","shell.execute_reply":"2024-05-23T06:33:51.76173Z"},"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-05-23T06:33:51.764138Z","iopub.execute_input":"2024-05-23T06:33:51.764441Z","iopub.status.idle":"2024-05-23T06:33:52.739541Z","shell.execute_reply.started":"2024-05-23T06:33:51.764416Z","shell.execute_reply":"2024-05-23T06:33:52.738335Z"},"trusted":true},"execution_count":null,"outputs":[]}]}