{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":71885,"databundleVersionId":8143495,"sourceType":"competition"},{"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":"# 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":{"execution":{"iopub.status.busy":"2024-05-14T16:54:27.144673Z","iopub.execute_input":"2024-05-14T16:54:27.145012Z","iopub.status.idle":"2024-05-14T16:56:41.527518Z","shell.execute_reply.started":"2024-05-14T16:54:27.144982Z","shell.execute_reply":"2024-05-14T16:56:41.526014Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%matplotlib inline","metadata":{"execution":{"iopub.status.busy":"2024-05-14T16:56:41.529889Z","iopub.execute_input":"2024-05-14T16:56:41.530843Z","iopub.status.idle":"2024-05-14T16:56:41.537191Z","shell.execute_reply.started":"2024-05-14T16:56:41.530802Z","shell.execute_reply":"2024-05-14T16:56:41.53615Z"},"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","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-05-14T16:56:41.538639Z","iopub.execute_input":"2024-05-14T16:56:41.53917Z","iopub.status.idle":"2024-05-14T16:56:49.928708Z","shell.execute_reply.started":"2024-05-14T16:56:41.539136Z","shell.execute_reply":"2024-05-14T16:56:49.927854Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from lightglue import match_pair\nfrom lightglue import ALIKED, SuperPoint, DoGHardNet, LightGlue, DISK, SIFT\nfrom lightglue.utils import load_image, rbd\nfrom kornia.feature import LoFTR","metadata":{"execution":{"iopub.status.busy":"2024-05-14T16:56:49.930831Z","iopub.execute_input":"2024-05-14T16:56:49.931136Z","iopub.status.idle":"2024-05-14T16:56:49.940435Z","shell.execute_reply.started":"2024-05-14T16:56:49.931109Z","shell.execute_reply":"2024-05-14T16:56:49.939594Z"},"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-14T16:56:49.941472Z","iopub.execute_input":"2024-05-14T16:56:49.941756Z","iopub.status.idle":"2024-05-14T16:56:49.953979Z","shell.execute_reply.started":"2024-05-14T16:56:49.941733Z","shell.execute_reply":"2024-05-14T16:56:49.953082Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Configurations","metadata":{}},{"cell_type":"code","source":"class CONFIG:\n # DEBUG Settings\n DRY_RUN = True\n DRY_RUN_MAX_IMAGES = 10\n \n # processing settings\n NUM_CORES = 2\n \n # Rotation correction\n ROTATION_CORRECTION = False\n \n # Keypoints handling\n MERGE_PARAMS = {\n \"min_matches\" : 15,\n \"filter_FundamentalMatrix\" : False,\n \"filter_iterations\" : 5,\n \"filter_threshold\" : 3,\n }\n \n # Keypoints Extraction(1st)\n use_aliked_lightglue = False # True\n use_aliked_lightglue_rots = False #True\n use_aliked_lightglue_4rots = True\n use_doghardnet_lightglue = False\n use_disk_lightglue = False\n use_sift_lightglue = False\n use_sift_lightglue_rots = False\n use_loftr = False\n use_dkm = False\n use_matchformer = False\n \n # Keypoints Extraction(2nd)\n use_aliked_lightglue_crop = True\n use_doghardnet_lightglue_crop = False\n use_superpoint_lightglue_crop = False\n use_disk_lightglue_crop = False\n use_sift_lightglue_crop = False\n \n # Keypoints Extraction Parameters\n params_aliked_lightglue_rots1 = {\n \"model_name\" : \"aliked-n16rot\",\n \"num_features\" : 2048,\n \"detection_threshold\" : 0.001,\n \"min_matches\" : 15,\n \"resize_to\" : 512,\n }\n params_aliked_lightglue_rots2 = {\n \"model_name\" : \"aliked-n16\",\n \"num_features\" : 2048,\n \"detection_threshold\" : 0.001,\n \"min_matches\" : 15,\n \"resize_to\" : 512,\n }\n list_params_aliked_lightglue_rots = [ params_aliked_lightglue_rots2]\n params_aliked_lightglue1 = {\n \"model_name\" : \"aliked-n16rot\",\n \"num_features\" : 8192,\n \"detection_threshold\" : 0.001,\n \"min_matches\" : 15,\n \"resize_to\" : 512,\n }\n params_aliked_lightglue2 = {\n \"model_name\" : \"aliked-n16\",\n \"num_features\" : 8192,\n \"detection_threshold\" : 0.001,\n \"min_matches\" : 15,\n \"resize_to\" : 512,\n }\n list_params_aliked_lightglue = [ params_aliked_lightglue2]\n params_aliked_lightglue_crop = {\n \"model_name\" : \"aliked-n16\",\n \"num_features\" : 2048,\n \"detection_threshold\" : 0.001,\n \"min_matches\" : 15,\n \"resize_to\" : 512,\n }\n \n params_doghardnet_lightglue = {\n \"model_name\" : \"doghardnet\",\n \"num_features\" : 8192,\n \"detection_threshold\" : 0.001,\n \"min_matches\" : 400,\n \"resize_to\" : 1024,\n }\n params_doghardnet_lightglue_crop = params_doghardnet_lightglue.copy()\n \n params_disk_lightglue = {\n \"model_name\" : \"disk\",\n \"num_features\" : 8192,\n \"detection_threshold\" : 0.001,\n \"min_matches\" : 400,\n \"resize_to\" : 1024,\n }\n params_disk_lightglue_crop = params_disk_lightglue.copy()\n\n params_sift_lightglue = {\n \"model_name\" : \"sift\",\n \"num_features\" : 8192,\n \"detection_threshold\" : 0.001,\n \"min_matches\" : 15,\n \"resize_to\" : 1024,\n }\n params_sift_lightglue_rots = {\n \"model_name\" : \"sift\",\n \"num_features\" : 2048,\n \"detection_threshold\" : 0.001,\n \"min_matches\" : 15,\n \"resize_to\" : 1024,\n }\n params_sift_lightglue_crop = params_sift_lightglue.copy()\n params_sift_lightglue_crop[\"num_features\"] = 8192\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 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-14T16:56:49.955149Z","iopub.execute_input":"2024-05-14T16:56:49.955408Z","iopub.status.idle":"2024-05-14T16:56:49.973891Z","shell.execute_reply.started":"2024-05-14T16:56:49.955386Z","shell.execute_reply":"2024-05-14T16:56:49.973097Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pandas as pd\ntmp = pd.read_csv(\"/kaggle/input/image-matching-challenge-2024/sample_submission.csv\")\nif tmp.shape[0] != 41:\n CONFIG.DRY_RUN = False","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device0=torch.device('cuda:0')\ndevice1=torch.device('cuda:1')","metadata":{"execution":{"iopub.status.busy":"2024-05-14T16:56:49.974951Z","iopub.execute_input":"2024-05-14T16:56:49.975206Z","iopub.status.idle":"2024-05-14T16:56:49.987942Z","shell.execute_reply.started":"2024-05-14T16:56:49.975184Z","shell.execute_reply":"2024-05-14T16:56:49.987244Z"},"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-14T16:56:49.989268Z","iopub.execute_input":"2024-05-14T16:56:49.989598Z","iopub.status.idle":"2024-05-14T16:56:50.01678Z","shell.execute_reply.started":"2024-05-14T16:56:49.989561Z","shell.execute_reply":"2024-05-14T16:56:50.015977Z"},"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 camera_model = 'simple-radial',\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, 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-14T16:56:50.017914Z","iopub.execute_input":"2024-05-14T16:56:50.018156Z","iopub.status.idle":"2024-05-14T16:56:50.039045Z","shell.execute_reply.started":"2024-05-14T16:56:50.018135Z","shell.execute_reply":"2024-05-14T16:56:50.038235Z"},"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_crop(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_crop(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":{"execution":{"iopub.status.busy":"2024-05-14T16:56:50.042563Z","iopub.execute_input":"2024-05-14T16:56:50.042837Z","iopub.status.idle":"2024-05-14T16:56:50.515551Z","shell.execute_reply.started":"2024-05-14T16:56:50.042815Z","shell.execute_reply":"2024-05-14T16:56:50.514751Z"},"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-14T16:56:50.516701Z","iopub.execute_input":"2024-05-14T16:56:50.516979Z","iopub.status.idle":"2024-05-14T16:56:50.534229Z","shell.execute_reply.started":"2024-05-14T16:56:50.516956Z","shell.execute_reply":"2024-05-14T16:56:50.533373Z"},"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 # np.array([rx, ry]).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 # np.array([rx, ry]).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 # np.array([rx, ry]).T\n\n#######\ndef detect_common_rot(img_fnames,\n model_name,\n rot,\n enough_pairs,\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-n16\" : ALIKED,\n \"aliked-n16rot\" : ALIKED,\n \"aliked-n32\" : ALIKED,\n \"superpoint\" : SuperPoint,\n \"doghardnet\" : DoGHardNet,\n \"disk\" : DISK,\n \"sift\" : SIFT,\n }\n extractor_class = dict_model[model_name]\n dtype = torch.float32\n if not \"aliked\" in model_name:\n extractor = extractor_class(\n max_num_keypoints=num_features, detection_threshold=detection_threshold, resize=resize_to\n ).eval().to(device, dtype)\n else:\n extractor = extractor_class(\n max_num_keypoints=num_features, detection_threshold=detection_threshold, resize=resize_to,\n model_name = model_name,\n ).eval().to(device, dtype)\n \n dict_kpts_cuda = {}\n dict_descs_cuda = {}\n for img_path in img_fnames:\n img_fname = img_path.split('/')[-1]\n key = img_fname\n\n with torch.inference_mode():\n rot_k = 0\n image0 = load_torch_image(img_path, device=device).to(dtype)\n h, w = image0.shape[2], image0.shape[3]\n feats0 = extractor.extract(image0) # 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 dict_kpts_cuda[f\"{key}_{rot_k}\"] = kpts\n dict_descs_cuda[f\"{key}_{rot_k}\"] = descs\n print(f\"{model_name} > rot_k={rot_k}, kpts.shape={kpts.shape}, descs.shape={descs.shape}\")\n\n if rot != 0:\n with torch.inference_mode():\n rot_k = rot\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, [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}_{rot_k}\"] = kpts\n dict_descs_cuda[f\"{key}_{rot_k}\"] = 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.split(\"-\")[0], {\"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 is_enough = False\n for ep in enough_pairs:\n if (key1 == ep[0]) and (key2 == ep[1]):\n is_enough = True\n if is_enough:\n continue\n\n kp1 = dict_kpts_cuda[f\"{key1}_0\"]\n desc1 = dict_descs_cuda[f\"{key1}_0\"]\n \n kp2 = dict_kpts_cuda[f\"{key2}_{rot}\"]\n desc2 = dict_descs_cuda[f\"{key2}_{rot}\"]\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 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 n_matches = kp1.shape[0]\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 enough_pairs.append([key1, key2])\n cnt_pairs+=1\n print (f'{model_name}> {key1}-{key2}: {n_matches} matches @ {cnt_pairs}th pair({model_name}+lightglue/rot={rot})') \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 enough_pairs\n\ndef detect_lightglue_common_rot(\n img_fnames, model_name, index_pairs, feature_dir, device, file_keypoints, enough_pairs, rot,\n resize_to, detection_threshold, num_features, min_matches,\n):\n t=time()\n enough_pairs = detect_common_rot(\n img_fnames, model_name, rot, enough_pairs, 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, enough_pairs\n\n \n#####\ndef detect_common_rots(img_fnames,\n model_name,\n rots,\n enough_pairs,\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-n16\" : ALIKED,\n \"aliked-n16rot\" : ALIKED,\n \"aliked-n32\" : ALIKED,\n \"superpoint\" : SuperPoint,\n \"doghardnet\" : DoGHardNet,\n \"disk\" : DISK,\n \"sift\" : SIFT,\n }\n extractor_class = dict_model[model_name]\n dtype = torch.float32\n if not \"aliked\" in model_name:\n extractor = extractor_class(\n max_num_keypoints=num_features, detection_threshold=detection_threshold, resize=resize_to\n ).eval().to(device, dtype)\n else:\n extractor = extractor_class(\n max_num_keypoints=num_features, detection_threshold=detection_threshold, resize=resize_to,\n model_name = model_name,\n ).eval().to(device, dtype)\n \n dict_kpts_cuda = {}\n dict_descs_cuda = {}\n with h5py.File(f'{feature_dir}/keypoints_{model_name}.h5', mode='w') as f_kp, \\\n h5py.File(f'{feature_dir}/descriptors_{model_name}.h5', mode='w') as f_desc:\n for img_path in img_fnames:\n img_fname = img_path.split('/')[-1]\n key = img_fname\n for rot_k in range(4):\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}_{rot_k}\"] = kpts\n dict_descs_cuda[f\"{key}_{rot_k}\"] = 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.split(\"-\")[0], {\"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 is_enough = False\n for ep in enough_pairs:\n if (key1 == ep[0]) and (key2 == ep[1]):\n is_enough = True\n if is_enough:\n continue\n\n kp1 = dict_kpts_cuda[f\"{key1}_0\"]\n desc1 = dict_descs_cuda[f\"{key1}_0\"]\n kp2 = []\n desc2 = []\n kp2_rot_k = []\n for rot_k in range(4):\n kp2.append( dict_kpts_cuda[f\"{key2}_{rot_k}\"] )\n desc2.append( dict_descs_cuda[f\"{key2}_{rot_k}\"] )\n kp2_rot_k.append( torch.zeros(dict_kpts_cuda[f\"{key2}_{rot_k}\"].shape[0]) + rot_k )\n kp2 = torch.concat(kp2, dim=0)\n desc2 = torch.concat(desc2, dim=0)\n kp2_rot_k = torch.concat(kp2_rot_k, dim=0).to(device)\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 kp2_rot_k = kp2_rot_k[idxs[:,1]].cpu().numpy()\n print(Counter(kp2_rot_k.tolist()))\n #med_rot_k = np.median(kp2_rot_k)\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 #kp1 = kp1[kp2_rot_k==med_rot_k, :]\n #kp2 = kp2[kp2_rot_k==med_rot_k, :]\n n_matches = kp1.shape[0]\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 enough_pairs.append([key1, key2])\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 enough_pairs\n\ndef detect_lightglue_common_rots(\n img_fnames, model_name, index_pairs, feature_dir, device, file_keypoints, enough_pairs, rots,\n resize_to, detection_threshold, num_features, min_matches,\n):\n t=time()\n enough_pairs = detect_common_rots(\n img_fnames, model_name, rots, enough_pairs, 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, enough_pairs\n\n\n#################\ndef detect_common(img_fnames,\n model_name,\n rots,\n enough_pairs,\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-n16\" : ALIKED,\n \"aliked-n16rot\" : ALIKED,\n \"aliked-n32\" : ALIKED,\n \"superpoint\" : SuperPoint,\n \"doghardnet\" : DoGHardNet,\n \"disk\" : DISK,\n \"sift\" : SIFT,\n }\n extractor_class = dict_model[model_name]\n dtype = torch.float32\n if not \"aliked\" in model_name:\n extractor = extractor_class(\n max_num_keypoints=num_features, detection_threshold=detection_threshold, resize=resize_to\n ).eval().to(device, dtype)\n else:\n extractor = extractor_class(\n max_num_keypoints=num_features, detection_threshold=detection_threshold, resize=resize_to,\n model_name = model_name,\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.split(\"-\")[0], {\"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 is_enough = False\n for ep in enough_pairs:\n if (key1 == ep[0]) and (key2 == ep[1]):\n is_enough = True\n if is_enough:\n continue\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 enough_pairs.append([key1, key2])\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 enough_pairs\n\ndef detect_lightglue_common(\n img_fnames, model_name, index_pairs, feature_dir, device, file_keypoints, enough_pairs, rots,\n resize_to, detection_threshold, num_features, min_matches,\n):\n t=time()\n enough_pairs = detect_common(\n img_fnames, model_name, rots, enough_pairs, 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, enough_pairs\n","metadata":{"execution":{"iopub.status.busy":"2024-05-14T16:56:50.535683Z","iopub.execute_input":"2024-05-14T16:56:50.535936Z","iopub.status.idle":"2024-05-14T16:56:50.769527Z","shell.execute_reply.started":"2024-05-14T16:56:50.535913Z","shell.execute_reply":"2024-05-14T16:56:50.768597Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def convert_image(image):\n image = image / 255\n return image.transpose((2, 0, 1)) # convert [H, W, C] into [C, H, W]\n \nclass LightGlueDataset_Crop(Dataset):\n def __init__(self, fnames1, fnames2, resize_to, crops, H21s):\n self.fnames1 = fnames1\n self.fnames2 = fnames2\n self.resize_to = resize_to\n self.crops = crops\n self.H21s = H21s\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 crop = self.crops[idx]\n im1_x0, im1_x1, im1_y0, im1_y1, im2_x0, im2_x1, im2_y0, im2_y1 = crop\n\n im1, im2 = cv2.cvtColor(cv2.imread(fname1), cv2.COLOR_BGR2RGB), cv2.cvtColor(cv2.imread(fname2), cv2.COLOR_BGR2RGB)\n\n h1, w1 = im1.shape[0:2]\n h2, w2 = im2.shape[0:2]\n \n # Image1\n if im1_x0 != -1:\n crop1_x0, crop1_x1, crop1_y0, crop1_y1 = max(0, im1_x0), min(w1, im1_x1), max(0, im1_y0), min(h1, im1_y1) \n im1 = im1[crop1_y0:crop1_y1, crop1_x0:crop1_x1, :]\n else:\n crop1_x0 = 0\n crop1_y0 = 0\n image1 = convert_image(im1)\n image1 = torch.from_numpy(image1)\n \n # Image2\n if im1_x0 != -1:\n H21 = self.H21s[idx]\n im2 = cv2.warpPerspective(im2, H21, (im1.shape[1], im1.shape[0]))\n else:\n crop2_x0 = 0\n crop2_y0 = 0\n image2 = convert_image(im2)\n image2 = torch.from_numpy(image2)\n\n crop = [crop1_x0, crop1_y0]\n ori_shape1 = [h1, w1]\n ori_shape2 = [h2, w2]\n return image1, image2, torch.tensor([idx]), torch.tensor(ori_shape1), torch.tensor(ori_shape2), torch.tensor(crop)\n\ndef get_lightglue_dataloader_crop(images1, images2, resize_to, crops, H21s, batch_size=1):\n dataset = LightGlueDataset_Crop(images1, images2, resize_to, crops, H21s)\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_lightglue_common_crop(\n img_fnames, model_name, index_pairs, feature_dir, device, file_keypoints, ref_file_keypoints,\n resize_to, detection_threshold, num_features, min_matches,\n): \n t=time()\n\n dict_model = {\n \"aliked-n16\" : ALIKED,\n \"aliked-n16rot\" : ALIKED,\n \"aliked-n32\" : ALIKED,\n \"superpoint\" : SuperPoint,\n \"doghardnet\" : DoGHardNet,\n \"disk\" : DISK,\n \"sift\" : SIFT,\n }\n extractor_class = dict_model[model_name]\n dtype=torch.float32\n if not \"aliked\" in model_name:\n extractor = extractor_class(\n max_num_keypoints=num_features, detection_threshold=detection_threshold, resize=resize_to\n ).eval().to(device, dtype)\n else:\n extractor = extractor_class(\n max_num_keypoints=num_features, detection_threshold=detection_threshold, resize=resize_to,\n model_name = model_name,\n ).eval().to(device, dtype)\n \n if not os.path.isdir(feature_dir):\n os.makedirs(feature_dir)\n \n if model_name != \"superpoint\":\n lg_matcher = KF.LightGlueMatcher(model_name.split(\"-\")[0], {\"width_confidence\": -1,\n \"depth_confidence\": -1,\n \"mp\": True if 'cuda' in str(device) else False}).eval().to(device)\n else:\n lg_matcher = KF.LightGlueMatcher(model_name, {\"width_confidence\": 0.95,\n \"depth_confidence\": 0.99,\n \"filter_threshold\": 0.1,\n \"mp\": True if 'cuda' in str(device) else False}).eval().to(device)\n\n fnames1, fnames2, idxs1, idxs2, crops, H12s, H21s = [], [], [], [], [], [], []\n with h5py.File(ref_file_keypoints, mode='r') as f_match_base:\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 # calculate crop area\n try:\n mkpts = np.array(f_match_base[key1][key2]).astype(np.int32)\n #print(mkpts.shape)\n pad=50\n im1_x0 = min(mkpts[:, 0]) - pad\n im1_x1 = max(mkpts[:, 0]) + pad\n im1_y0 = min(mkpts[:, 1]) - pad\n im1_y1 = max(mkpts[:, 1]) + pad\n im2_x0 = min(mkpts[:, 2]) - pad\n im2_x1 = max(mkpts[:, 2]) + pad\n im2_y0 = min(mkpts[:, 3]) - pad\n im2_y1 = max(mkpts[:, 3]) + pad\n \n # Homography matrix\n mkpts[:, 0] -= im1_x0\n mkpts[:, 1] -= im1_y0\n H21, inlier = cv2.findHomography(mkpts[:, 2:], mkpts[:, :2], cv2.RANSAC,5.0)\n if H21 is None:\n continue\n H12 = np.linalg.inv(H21)\n if H12 is None:\n continue\n\n crops.append( [im1_x0, im1_x1, im1_y0, im1_y1, im2_x0, im2_x1, im2_y0, im2_y1] )\n H12s.append(H12)\n H21s.append(H21)\n \n except:\n #print(f\"skipped key1={key1}, key2={key2}: not found reference kpts\")\n #crops.append( [-1, -1, -1, -1, -1, -1, -1, -1] )\n continue\n fnames1.append(fname1)\n fnames2.append(fname2)\n idxs1.append(idx1)\n idxs2.append(idx2)\n print(len(index_pairs), \"--->\", len(crops))\n dataloader = get_lightglue_dataloader_crop( fnames1, fnames2, resize_to, crops, H21s)\n\n cnt_pairs = 0\n\n with h5py.File(file_keypoints, mode='w') as f_match:\n for X in dataloader:\n try:\n image1, image2, idx, ori_shape_1, ori_shape_2, crop = X\n\n fname1, fname2 = fnames1[idx], fnames2[idx]\n key1, key2 = fname1.split('/')[-1], fname2.split('/')[-1]\n \n with torch.inference_mode(): \n image1 = image1.to(device).to(dtype)\n feats1 = extractor.extract(image1[0]) # auto-resize the image, disable with resize=None\n kpts1 = feats1['keypoints'].reshape(-1, 2)#.detach().cpu().numpy()\n descs1 = feats1['descriptors'].reshape(kpts1.shape[0], -1)#.detach().cpu().numpy()\n del image1\n\n image2 = image2.to(device).to(dtype)\n feats2 = extractor.extract(image2[0]) # auto-resize the image, disable with resize=None\n kpts2 = feats2['keypoints'].reshape(-1, 2)#.detach().cpu().numpy()\n descs2 = feats2['descriptors'].reshape(kpts2.shape[0], -1)#.detach().cpu().numpy()\n del image2\n\n dists, idxs = lg_matcher(descs1,\n descs2,\n KF.laf_from_center_scale_ori(kpts1[None]),\n KF.laf_from_center_scale_ori(kpts2[None]))\n if len(idxs) == 0:\n continue\n n_matches = len(idxs)\n kpts1 = kpts1[idxs[:,0], :].cpu().numpy().reshape(-1, 2).astype(np.float32)\n kpts2 = kpts2[idxs[:,1], :].cpu().numpy().reshape(-1, 2).astype(np.float32)\n\n crop1_x0, crop1_y0 = crop[0].numpy().copy()\n \n # Scaling coords\n kpts1[:,0] = kpts1[:,0] + crop1_x0 # X\n kpts1[:,1] = kpts1[:,1] + crop1_y0 # Y\n \n H12 = H12s[idx]\n kpts2 = cv2.perspectiveTransform(kpts2[:, np.newaxis, :2], H12)\n kpts2 = kpts2[:, 0, :]\n \n n_matches = kpts1.shape[0]\n \n group = f_match.require_group(key1)\n if n_matches >= min_matches:\n group.create_dataset(key2, data=np.concatenate([kpts1, kpts2], axis=1).astype(np.float32))\n cnt_pairs+=1\n print (f'{model_name}> {key1}-{key2}: {n_matches} matches @ {cnt_pairs}th pair({model_name}+LightGlue/crop)') \n else:\n print (f'{model_name}> {key1}-{key2}: {n_matches} matches --> skipped')\n except:\n print(f\"skipped. key1={key1}, key2={key2}\")\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-14T16:56:50.77088Z","iopub.execute_input":"2024-05-14T16:56:50.771212Z","iopub.status.idle":"2024-05-14T16:56:50.815328Z","shell.execute_reply.started":"2024-05-14T16:56:50.771176Z","shell.execute_reply":"2024-05-14T16:56:50.814599Z"},"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, detection_threshold, num_features, min_matches,\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-14T16:56:50.95853Z","iopub.execute_input":"2024-05-14T16:56:50.958901Z","iopub.status.idle":"2024-05-14T16:56:51.109243Z","shell.execute_reply.started":"2024-05-14T16:56:50.958867Z","shell.execute_reply":"2024-05-14T16:56:51.108102Z"},"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, min_matches):\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-14T16:56:51.110613Z","iopub.execute_input":"2024-05-14T16:56:51.111333Z","iopub.status.idle":"2024-05-14T16:56:51.228125Z","shell.execute_reply.started":"2024-05-14T16:56:51.111295Z","shell.execute_reply":"2024-05-14T16:56:51.227065Z"},"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-14T16:56:51.229635Z","iopub.execute_input":"2024-05-14T16:56:51.229974Z","iopub.status.idle":"2024-05-14T16:56:51.347043Z","shell.execute_reply.started":"2024-05-14T16:56:51.229949Z","shell.execute_reply":"2024-05-14T16:56:51.345971Z"},"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, detection_threshold, num_features, min_matches\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-14T16:56:51.348488Z","iopub.execute_input":"2024-05-14T16:56:51.3489Z","iopub.status.idle":"2024-05-14T16:56:51.483076Z","shell.execute_reply.started":"2024-05-14T16:56:51.348872Z","shell.execute_reply":"2024-05-14T16:56:51.481884Z"},"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-14T16:56:51.485645Z","iopub.execute_input":"2024-05-14T16:56:51.486277Z","iopub.status.idle":"2024-05-14T16:56:51.585622Z","shell.execute_reply.started":"2024-05-14T16:56:51.486236Z","shell.execute_reply":"2024-05-14T16:56:51.584514Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Keypoints wrapper function","metadata":{}},{"cell_type":"code","source":"def switch_gpu(device, device_index):\n if device is None:\n return device0, 0\n elif device == device0:\n return device1, 1\n else:\n return device0, 0\n \ndef wrapper_keypoints(\n img_fnames, index_pairs, feature_dir, timings, rots\n):\n #############################################################\n # get keypoints (1st)\n #############################################################\n files_keypoints = []\n device = None\n device_index = None\n\n with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executors_gpu0, \\\n concurrent.futures.ThreadPoolExecutor(max_workers=1) as executors_gpu1:\n executors_gpus = [executors_gpu0, executors_gpu1]\n l_futures = []\n enough_pairs = []\n\n if CONFIG.use_aliked_lightglue_rots:\n for config in CONFIG.list_params_aliked_lightglue_rots:\n device, device_index = switch_gpu(device, device_index)\n model_name = config[\"model_name\"]\n file_keypoints = f'{feature_dir}/matches_lightglue_{model_name}_rots.h5'\n l_futures.append(\n executors_gpus[device_index].submit(\n detect_lightglue_common_rots,\n img_fnames, model_name, index_pairs, feature_dir, device, file_keypoints, enough_pairs, rots,\n config[\"resize_to\"],\n config[\"detection_threshold\"], \n config[\"num_features\"], \n config[\"min_matches\"], \n )\n )\n files_keypoints.append(file_keypoints)\n\n if CONFIG.use_aliked_lightglue:\n for config in CONFIG.list_params_aliked_lightglue:\n model_name = config[\"model_name\"]\n device, device_index = switch_gpu(device, device_index)\n file_keypoints = f'{feature_dir}/matches_lightglue_{model_name}.h5'\n l_futures.append(\n executors_gpus[device_index].submit(\n detect_lightglue_common,\n img_fnames, model_name, index_pairs, feature_dir, device, file_keypoints, enough_pairs, rots,\n config[\"resize_to\"],\n config[\"detection_threshold\"], \n config[\"num_features\"], \n config[\"min_matches\"], \n )\n )\n files_keypoints.append(file_keypoints)\n\n if CONFIG.use_aliked_lightglue_4rots:\n for config in CONFIG.list_params_aliked_lightglue:\n model_name = config[\"model_name\"]\n for rot in range(4):\n device, device_index = switch_gpu(device, device_index)\n file_keypoints = f'{feature_dir}/matches_lightglue_{model_name}_rot{rot}.h5'\n l_futures.append(\n executors_gpus[device_index].submit(\n detect_lightglue_common_rot,\n img_fnames, model_name, index_pairs, feature_dir, device, file_keypoints, enough_pairs, rot,\n config[\"resize_to\"],\n config[\"detection_threshold\"], \n config[\"num_features\"], \n config[\"min_matches\"], \n )\n )\n files_keypoints.append(file_keypoints)\n\n if CONFIG.use_doghardnet_lightglue:\n device, device_index = switch_gpu(device, device_index)\n model_name = \"doghardnet\"\n file_keypoints = f'{feature_dir}/matches_lightglue_{model_name}.h5'\n l_futures.append(\n executors_gpus[device_index].submit(\n detect_lightglue_common,\n img_fnames, model_name, index_pairs, feature_dir, device, file_keypoints, enough_pairs, rots,\n CONFIG.params_doghardnet_lightglue[\"resize_to\"],\n CONFIG.params_doghardnet_lightglue[\"detection_threshold\"], \n CONFIG.params_doghardnet_lightglue[\"num_features\"], \n CONFIG.params_doghardnet_lightglue[\"min_matches\"], \n )\n )\n files_keypoints.append(file_keypoints)\n\n if CONFIG.use_disk_lightglue:\n device, device_index = switch_gpu(device, device_index)\n model_name = \"disk\"\n file_keypoints = f'{feature_dir}/matches_lightglue_{model_name}.h5'\n l_futures.append(\n executors_gpus[device_index].submit(\n detect_lightglue_common,\n img_fnames, model_name, index_pairs, feature_dir, device, file_keypoints, enough_pairs, rots,\n CONFIG.params_disk_lightglue[\"resize_to\"],\n CONFIG.params_disk_lightglue[\"detection_threshold\"], \n CONFIG.params_disk_lightglue[\"num_features\"], \n CONFIG.params_disk_lightglue[\"min_matches\"], \n )\n )\n files_keypoints.append(file_keypoints)\n\n if CONFIG.use_sift_lightglue:\n device, device_index = switch_gpu(device, device_index)\n model_name = \"sift\"\n enough_pairs = []\n file_keypoints = f'{feature_dir}/matches_lightglue_{model_name}.h5'\n l_futures.append(\n executors_gpus[device_index].submit(\n detect_lightglue_common,\n img_fnames, model_name, index_pairs, feature_dir, device, file_keypoints, enough_pairs, rots,\n CONFIG.params_sift_lightglue[\"resize_to\"],\n CONFIG.params_sift_lightglue[\"detection_threshold\"], \n CONFIG.params_sift_lightglue[\"num_features\"], \n CONFIG.params_sift_lightglue[\"min_matches\"], \n )\n )\n files_keypoints.append(file_keypoints)\n\n if CONFIG.use_loftr:\n device, device_index = switch_gpu(device, device_index)\n file_keypoints = f'{feature_dir}/matches_loftr_{CONFIG.params_loftr[\"resize_small_edge_to\"]}pix.h5'\n l_futures.append(\n executors_gpus[device_index].submit(\n detect_loftr,\n img_fnames, index_pairs, feature_dir, device, file_keypoints,\n CONFIG.params_loftr[\"resize_small_edge_to\"],\n CONFIG.params_loftr[\"min_matches\"],\n )\n )\n files_keypoints.append(file_keypoints)\n\n if CONFIG.use_dkm:\n device, device_index = switch_gpu(device, device_index)\n file_keypoints = f'{feature_dir}/matches_dkm.h5'\n l_futures.append(\n executors_gpus[device_index].submit(\n detect_dkm,\n img_fnames, index_pairs, feature_dir, device, file_keypoints,\n CONFIG.params_dkm[\"resize_to\"],\n CONFIG.params_dkm[\"detection_threshold\"],\n CONFIG.params_dkm[\"num_features\"],\n CONFIG.params_dkm[\"min_matches\"],\n )\n )\n files_keypoints.append(file_keypoints)\n\n if CONFIG.use_matchformer:\n file_keypoints = f'{feature_dir}/matches_matchformer_{CONFIG.params_matchformer[\"resize_to\"]}pix.h5'\n l_futures.append(\n executors_gpus[device_index].submit(\n detect_matchformer,\n img_fnames, index_pairs, feature_dir, device, file_keypoints,\n CONFIG.params_matchformer[\"resize_to\"],\n CONFIG.params_matchformer[\"num_features\"], \n CONFIG.params_matchformer[\"min_matches\"]\n )\n )\n files_keypoints.append(file_keypoints)\n\n # wait to complete\n for l_future in l_futures:\n gc.collect()\n results = l_future.result()\n gc.collect()\n\n save_1st_file = f'{feature_dir}/matches_merge_1st.h5'\n !rm -rf {save_1st_file}\n matches_merger(\n img_fnames,\n index_pairs,\n files_keypoints,\n save_1st_file,\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 )\n files_keypoints = [save_1st_file]\n\n #############################################################\n # get keypoints (2nd)\n #############################################################\n with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executors_gpu0, \\\n concurrent.futures.ThreadPoolExecutor(max_workers=1) as executors_gpu1:\n executors_gpus = [executors_gpu0, executors_gpu1]\n l_futures = []\n\n if CONFIG.use_aliked_lightglue_crop:\n device, device_index = switch_gpu(device, device_index)\n model_name = CONFIG.params_aliked_lightglue_crop[\"model_name\"]\n file_keypoints = f'{feature_dir}/matches_lightglue_{model_name}_crop.h5'\n l_futures.append(\n executors_gpus[device_index].submit(\n detect_lightglue_common_crop,\n img_fnames, model_name, index_pairs, feature_dir, device, file_keypoints, save_1st_file,\n CONFIG.params_aliked_lightglue_crop[\"resize_to\"],\n CONFIG.params_aliked_lightglue_crop[\"detection_threshold\"],\n CONFIG.params_aliked_lightglue_crop[\"num_features\"],\n CONFIG.params_aliked_lightglue_crop[\"min_matches\"], \n )\n )\n gc.collect()\n files_keypoints.append(file_keypoints)\n\n if CONFIG.use_doghardnet_lightglue_crop:\n device, device_index = switch_gpu(device, device_index)\n model_name = \"doghardnet\"\n file_keypoints = f'{feature_dir}/matches_lightglue_{model_name}_crop.h5'\n l_futures.append(\n executors_gpus[device_index].submit(\n detect_lightglue_common_crop,\n img_fnames, model_name, index_pairs, feature_dir, device, file_keypoints, save_1st_file,\n CONFIG.params_doghardnet_lightglue_crop[\"resize_to\"],\n CONFIG.params_doghardnet_lightglue_crop[\"detection_threshold\"],\n CONFIG.params_doghardnet_lightglue_crop[\"num_features\"],\n CONFIG.params_doghardnet_lightglue_crop[\"min_matches\"], \n )\n )\n gc.collect()\n files_keypoints.append(file_keypoints)\n\n if CONFIG.use_disk_lightglue_crop:\n device, device_index = switch_gpu(device, device_index)\n model_name = \"disk\"\n file_keypoints = f'{feature_dir}/matches_lightglue_{model_name}_crop.h5'\n l_futures.append(\n executors_gpus[device_index].submit(\n detect_lightglue_common_crop,\n img_fnames, model_name, index_pairs, feature_dir, device, file_keypoints, save_1st_file,\n CONFIG.params_disk_lightglue_crop[\"resize_to\"],\n CONFIG.params_disk_lightglue_crop[\"detection_threshold\"],\n CONFIG.params_disk_lightglue_crop[\"num_features\"],\n CONFIG.params_disk_lightglue_crop[\"min_matches\"], \n )\n )\n gc.collect()\n files_keypoints.append(file_keypoints)\n \n if CONFIG.use_sift_lightglue_crop:\n device, device_index = switch_gpu(device, device_index)\n model_name = \"sift\"\n file_keypoints = f'{feature_dir}/matches_lightglue_{model_name}_crop.h5'\n l_futures.append(\n executors_gpus[device_index].submit(\n detect_lightglue_common_crop,\n img_fnames, model_name, index_pairs, feature_dir, device, file_keypoints, save_1st_file,\n CONFIG.params_sift_lightglue_crop[\"resize_to\"],\n CONFIG.params_sift_lightglue_crop[\"detection_threshold\"],\n CONFIG.params_sift_lightglue_crop[\"num_features\"],\n CONFIG.params_sift_lightglue_crop[\"min_matches\"], \n )\n )\n gc.collect()\n files_keypoints.append(file_keypoints)\n \n # wait to complete\n for l_future in l_futures:\n gc.collect()\n results = l_future.result()\n gc.collect()\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 = False,#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-14T16:56:51.587525Z","iopub.execute_input":"2024-05-14T16:56:51.587926Z","iopub.status.idle":"2024-05-14T16:56:51.722047Z","shell.execute_reply.started":"2024-05-14T16:56:51.587882Z","shell.execute_reply":"2024-05-14T16:56:51.720881Z"},"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, camera_model, miter):\n scene_result = {}\n #############################################################\n # regist keypoints from h5 into colmap db\n #############################################################\n database_path = f'{feature_dir}/colmap{miter}.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, camera_model=camera_model, database_path=database_path)\n output_path = f'{feature_dir}/colmap{miter}_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 mapper_options.ba_local_num_images = 12\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":{"execution":{"iopub.status.busy":"2024-05-14T16:56:51.723878Z","iopub.execute_input":"2024-05-14T16:56:51.724191Z","iopub.status.idle":"2024-05-14T16:56:51.841223Z","shell.execute_reply.started":"2024-05-14T16:56:51.724165Z","shell.execute_reply":"2024-05-14T16:56:51.840292Z"},"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":{"execution":{"iopub.status.busy":"2024-05-14T16:56:51.84242Z","iopub.execute_input":"2024-05-14T16:56:51.842842Z","iopub.status.idle":"2024-05-14T16:56:51.939499Z","shell.execute_reply.started":"2024-05-14T16:56:51.842808Z","shell.execute_reply":"2024-05-14T16:56:51.938576Z"},"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-14T16:56:51.941051Z","iopub.execute_input":"2024-05-14T16:56:51.941437Z","iopub.status.idle":"2024-05-14T16:56:52.044181Z","shell.execute_reply.started":"2024-05-14T16:56:51.941406Z","shell.execute_reply":"2024-05-14T16:56:52.043312Z"},"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(list)\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, device0)\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=device0)\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, timings, rots\n )\n timings['feature_matching'] = keypoints_timings['feature_matching']\n gc.collect()\n\n #############################################################\n # reconstruction trigger\n #############################################################\n # simple-radial\n futures[dataset].append(\n executors.submit(\n reconstruct_from_db, \n dataset, scene, feature_dir, img_dir, timings, data_dict[dataset][scene], 'simple-radial', 0\n )\n )\n # simple-pinhole\n futures[dataset].append(\n executors.submit(\n reconstruct_from_db, \n dataset, scene, feature_dir, img_dir, timings, data_dict[dataset][scene], 'simple-pinhole', 1\n )\n )\n # simple-radial\n futures[dataset].append(\n executors.submit(\n reconstruct_from_db, \n dataset, scene, feature_dir, img_dir, timings, data_dict[dataset][scene], 'simple-radial', 2\n )\n )\n \n #############################################################\n # reconstruction results\n ############################################################# \n for dataset in datasets:\n scene = dataset\n best_size = -1\n best_result = None\n history = []\n for future in futures[dataset]:\n gc.collect()\n result = future.result()\n gc.collect()\n if result is not None:\n history.append(len(result))\n if len(result) > best_size:\n best_size = len(result)\n best_result = result\n if best_result is not None:\n out_results[dataset][scene] = best_result\n print(f\"dataset={dataset} : history of map size = {history}\")\n \n create_submission(out_results, data_dict)\n gc.collect()","metadata":{"execution":{"iopub.status.busy":"2024-05-14T16:56:52.045949Z","iopub.execute_input":"2024-05-14T16:56:52.046336Z","iopub.status.idle":"2024-05-14T16:57:59.639602Z","shell.execute_reply.started":"2024-05-14T16:56:52.046303Z","shell.execute_reply":"2024-05-14T16:57:59.638683Z"},"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-14T16:57:59.64097Z","iopub.execute_input":"2024-05-14T16:57:59.641688Z","iopub.status.idle":"2024-05-14T16:58:00.614283Z","shell.execute_reply.started":"2024-05-14T16:57:59.641654Z","shell.execute_reply":"2024-05-14T16:58:00.61313Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if CONFIG.DRY_RUN:\n !cp submission.csv submission_bk.csv\n !cp /kaggle/input/image-matching-challenge-2024/sample_submission.csv submission.csv","metadata":{},"execution_count":null,"outputs":[]}]}