{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":71885,"databundleVersionId":8143495,"sourceType":"competition"},{"sourceId":3414836,"sourceType":"datasetVersion","datasetId":2058261},{"sourceId":5373920,"sourceType":"datasetVersion","datasetId":3117886},{"sourceId":7884485,"sourceType":"datasetVersion","datasetId":4628051},{"sourceId":170475544,"sourceType":"kernelVersion"},{"sourceId":170565695,"sourceType":"kernelVersion"},{"sourceId":174129945,"sourceType":"kernelVersion"},{"sourceId":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-index --find-links=/kaggle/input/pkg-check-orientation/ check_orientation==0.0.5 > /dev/null","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#DISK\n!cp /kaggle/input/disk/pytorch/depth-supervision/1/* /root/.cache/torch/hub/checkpoints/\n!cp /kaggle/input/lightglue/pytorch/disk/1/* /root/.cache/torch/hub/checkpoints/\n!cp /kaggle/input/disk-lightglue-weight/disk_lightglue.pth /root/.cache/torch/hub/checkpoints/disk_lightglue_v0-1_arxiv.pth\n!cp /kaggle/input/disk-lightglue-weight/depth-save.pth /root/.cache/torch/hub/checkpoints/depth-save.pth\n\n!cp /kaggle/input/disk-lightglue-weight/disk_lightglue.pth /root/.cache/torch/hub/checkpoints/disk_lightglue_v0-1_arxiv-pth","metadata":{"execution":{"iopub.status.busy":"2024-05-19T00:21:24.142087Z","iopub.execute_input":"2024-05-19T00:21:24.142784Z","iopub.status.idle":"2024-05-19T00:21:29.417717Z","shell.execute_reply.started":"2024-05-19T00:21:24.142741Z","shell.execute_reply":"2024-05-19T00:21:29.416383Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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","metadata":{"execution":{"iopub.status.busy":"2024-05-19T00:21:29.420467Z","iopub.execute_input":"2024-05-19T00:21:29.420911Z","iopub.status.idle":"2024-05-19T00:23:04.639748Z","shell.execute_reply.started":"2024-05-19T00:21:29.420866Z","shell.execute_reply":"2024-05-19T00:23:04.638351Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# 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/","metadata":{"execution":{"iopub.status.busy":"2024-05-19T00:23:04.641146Z","iopub.execute_input":"2024-05-19T00:23:04.641502Z","iopub.status.idle":"2024-05-19T00:23:18.042608Z","shell.execute_reply.started":"2024-05-19T00:23:04.641473Z","shell.execute_reply":"2024-05-19T00:23:18.041362Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%matplotlib inline","metadata":{"execution":{"iopub.status.busy":"2024-05-19T00:23:18.045232Z","iopub.execute_input":"2024-05-19T00:23:18.045568Z","iopub.status.idle":"2024-05-19T00:23:18.052432Z","shell.execute_reply.started":"2024-05-19T00:23:18.045539Z","shell.execute_reply":"2024-05-19T00:23:18.051344Z"},"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\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-19T00:23:18.053738Z","iopub.execute_input":"2024-05-19T00:23:18.054148Z","iopub.status.idle":"2024-05-19T00:23:27.06059Z","shell.execute_reply.started":"2024-05-19T00:23:18.054114Z","shell.execute_reply":"2024-05-19T00:23:27.059639Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from lightglue import match_pair\nfrom lightglue import ALIKED, SuperPoint, DoGHardNet, LightGlue,DISK\nfrom lightglue.utils import load_image, rbd\nfrom kornia.feature import LoFTR","metadata":{"execution":{"iopub.status.busy":"2024-05-19T00:23:27.062039Z","iopub.execute_input":"2024-05-19T00:23:27.062484Z","iopub.status.idle":"2024-05-19T00:23:27.073476Z","shell.execute_reply.started":"2024-05-19T00:23:27.062443Z","shell.execute_reply":"2024-05-19T00:23:27.072541Z"},"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-19T00:23:27.074481Z","iopub.execute_input":"2024-05-19T00:23:27.074871Z","iopub.status.idle":"2024-05-19T00:23:27.086616Z","shell.execute_reply.started":"2024-05-19T00:23:27.074844Z","shell.execute_reply":"2024-05-19T00:23:27.085636Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torchvision.io import read_image as T_read_image\nfrom torchvision.io import ImageReadMode\nfrom torchvision import transforms as T\nfrom check_orientation.pre_trained_models import create_model\n\ndef convert_rot_k(index):\n    if index == 0:\n        return 0\n    elif index == 1:\n        return 3\n    elif index == 2:\n        return 2\n    else:\n        return 1\n\nclass CheckRotationDataset(Dataset):\n    def __init__(self, files, transform=None):\n        self.transform = transform\n        self.files = files\n\n    def __len__(self):\n        return len(self.files)\n\n    def __getitem__(self, idx):\n        imgPath = self.files[idx]\n        image = T_read_image(imgPath, mode=ImageReadMode.RGB)\n        if self.transform:\n            image = self.transform(image)\n        return image\n\ndef get_CheckRotation_dataloader(images, batch_size=1):\n    transform = T.Compose([\n        T.Resize((224, 224)),\n        T.ConvertImageDtype(torch.float),\n        T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))\n    ])\n\n    dataset = CheckRotationDataset(images, transform=transform)\n    dataloader = DataLoader(\n        dataset=dataset,\n        shuffle=False,\n        batch_size=batch_size,\n        pin_memory=True,\n        num_workers=2,\n        drop_last=False\n    )\n    return dataloader\n\ndef exec_rotation_detection(img_files, device):\n    model = create_model(\"swsl_resnext50_32x4d\")\n    model.eval().to(device);\n    \n    dataloader = get_CheckRotation_dataloader(img_files)\n    \n    rots = []\n    for idx, image in enumerate(dataloader):\n        image = image.to(torch.float32).to(device)\n        with torch.no_grad():\n            prediction = model(image).detach().cpu().numpy()\n            detected_rot = prediction[0].argmax()\n            rot_k = convert_rot_k(detected_rot)\n            rots.append(rot_k)\n            print(f\"{os.path.basename(img_files[idx])} > rot_k={rot_k}\")\n    return rots","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Configurations","metadata":{}},{"cell_type":"code","source":"class CONFIG:\n    use_aliked_lightglue = True\n    use_doghardnet_lightglue =False\n    use_superpoint_lightglue = False\n    use_loftr =False\n    use_dkm = False\n    use_superglue = False\n    ROTATION_CORRECTION = True\n    \n    params_aliked_lightglue = {\n        \"num_features\" : 8192,\n        \"detection_threshold\" : 0.0095,\n        \"min_matches\" : 66,\n        \"resize_to\" : 1360,\n    }\n    params_doghardnet_lightglue = {\n        \"num_features\" : 8192,\n        \"detection_threshold\" : 0.001,\n        \"min_matches\" : 15,\n        \"resize_to\" : 1024,\n    }\n    \n    params_superpoint_lightglue = {\n        \"num_features\" : 4096,\n        \"detection_threshold\" : 0.005,\n        \"min_matches\" : 15,\n        \"resize_to\" : 1024,\n    }\n    \n    params_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_sg = {\n        \"sg_config\" : \n        {\n            \"superpoint\": {\n                \"nms_radius\": 4, \n                \"keypoint_threshold\": 0.0001,\n                \"max_keypoints\": 4096\n            },\n            \"superglue\": {\n                \"weights\": \"outdoor\",\n                \"sinkhorn_iterations\": 10,\n                \"match_threshold\": 0.2,\n            },\n        },\n        \"resize_to\": 1240,\n        \"min_matches\": 15,\n    }","metadata":{"execution":{"iopub.status.busy":"2024-05-19T00:23:27.087726Z","iopub.execute_input":"2024-05-19T00:23:27.088082Z","iopub.status.idle":"2024-05-19T00:23:27.098624Z","shell.execute_reply.started":"2024-05-19T00:23:27.088056Z","shell.execute_reply":"2024-05-19T00:23:27.097563Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device=torch.device('cuda')","metadata":{"execution":{"iopub.status.busy":"2024-05-19T00:23:27.099844Z","iopub.execute_input":"2024-05-19T00:23:27.100187Z","iopub.status.idle":"2024-05-19T00:23:27.112212Z","shell.execute_reply.started":"2024-05-19T00:23:27.100162Z","shell.execute_reply":"2024-05-19T00:23:27.111217Z"},"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-19T00:23:27.11394Z","iopub.execute_input":"2024-05-19T00:23:27.114259Z","iopub.status.idle":"2024-05-19T00:23:27.146827Z","shell.execute_reply.started":"2024-05-19T00:23:27.114234Z","shell.execute_reply":"2024-05-19T00:23:27.145859Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# h5 to colmap db","metadata":{}},{"cell_type":"code","source":"# Code to interface DISK with Colmap.\n# Forked from https://github.com/cvlab-epfl/disk/blob/37f1f7e971cea3055bb5ccfc4cf28bfd643fa339/colmap/h5_to_db.py\n\n#  Copyright [2020] [Michał Tyszkiewicz, Pascal Fua, Eduard Trulls]\n#\n#   Licensed under the Apache License, Version 2.0 (the \"License\");\n#   you may not use this file except in compliance with the License.\n#   You may obtain a copy of the License at\n#\n#       http://www.apache.org/licenses/LICENSE-2.0\n#\n#   Unless required by applicable law or agreed to in writing, software\n#   distributed under the License is distributed on an \"AS IS\" BASIS,\n#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n#   See the License for the specific language governing permissions and\n#   limitations under the License.\n\nimport os, argparse, h5py, warnings\nimport numpy as np\nfrom tqdm import tqdm\nfrom PIL import Image, ExifTags\n\n\ndef get_focal(image_path, err_on_default=False):\n    image         = Image.open(image_path)\n    max_size      = max(image.size)\n\n    exif = image.getexif()\n    focal = None\n    if exif is not None:\n        focal_35mm = None\n        # https://github.com/colmap/colmap/blob/d3a29e203ab69e91eda938d6e56e1c7339d62a99/src/util/bitmap.cc#L299\n        for tag, value in exif.items():\n            focal_35mm = None\n            if ExifTags.TAGS.get(tag, None) == 'FocalLengthIn35mmFilm':\n                focal_35mm = float(value)\n                break\n\n        if focal_35mm is not None:\n            focal = focal_35mm / 35. * max_size\n    \n    if focal is None:\n        if err_on_default:\n            raise RuntimeError(\"Failed to find focal length\")\n\n        # failed to find it in exif, use prior\n        FOCAL_PRIOR = 1.2\n        focal = FOCAL_PRIOR * max_size\n\n    return focal\n\ndef create_camera(db, image_path, camera_model):\n    image         = Image.open(image_path)\n    width, height = image.size\n\n    focal = get_focal(image_path)\n\n    if camera_model == 'simple-pinhole':\n        model = 0 # simple pinhole\n        param_arr = np.array([focal, width / 2, height / 2])\n    if camera_model == 'pinhole':\n        model = 1 # pinhole\n        param_arr = np.array([focal, focal, width / 2, height / 2])\n    elif camera_model == 'simple-radial':\n        model = 2 # simple radial\n        param_arr = np.array([focal, width / 2, height / 2, 0.1])\n    elif camera_model == 'opencv':\n        model = 4 # opencv\n        param_arr = np.array([focal, focal, width / 2, height / 2, 0., 0., 0., 0.])\n         \n    return db.add_camera(model, width, height, param_arr)\n\n\ndef add_keypoints(db, h5_path, image_path, img_ext, camera_model, single_camera = True):\n    keypoint_f = h5py.File(os.path.join(h5_path, 'keypoints.h5'), 'r')\n\n    camera_id = None\n    fname_to_id = {}\n    for filename in tqdm(list(keypoint_f.keys())):\n        keypoints = keypoint_f[filename][()]\n\n        fname_with_ext = filename# + img_ext\n        path = os.path.join(image_path, fname_with_ext)\n        if not os.path.isfile(path):\n            raise IOError(f'Invalid image path {path}')\n\n        if camera_id is None or not single_camera:\n            camera_id = create_camera(db, path, camera_model)\n        image_id = db.add_image(fname_with_ext, camera_id)\n        fname_to_id[filename] = image_id\n\n        db.add_keypoints(image_id, keypoints)\n\n    return fname_to_id\n\ndef add_matches(db, h5_path, fname_to_id):\n    match_file = h5py.File(os.path.join(h5_path, 'matches.h5'), 'r')\n    \n    added = set()\n    n_keys = len(match_file.keys())\n    n_total = (n_keys * (n_keys - 1)) // 2\n\n    with tqdm(total=n_total) as pbar:\n        for key_1 in match_file.keys():\n            group = match_file[key_1]\n            for key_2 in group.keys():\n                id_1 = fname_to_id[key_1]\n                id_2 = fname_to_id[key_2]\n\n                pair_id = image_ids_to_pair_id(id_1, id_2)\n                if pair_id in added:\n                    warnings.warn(f'Pair {pair_id} ({id_1}, {id_2}) already added!')\n                    continue\n            \n                matches = group[key_2][()]\n                db.add_matches(id_1, id_2, matches)\n\n                added.add(pair_id)\n\n                pbar.update(1)\n                \ndef import_into_colmap(img_dir,\n                       feature_dir ='.featureout',\n                       database_path = 'colmap.db',\n                       img_ext='.jpg'):\n    db = COLMAPDatabase.connect(database_path)\n    db.create_tables()\n    single_camera = False\n    fname_to_id = add_keypoints(db, feature_dir, img_dir, img_ext, 'simple-radial', 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-19T00:23:27.150972Z","iopub.execute_input":"2024-05-19T00:23:27.151464Z","iopub.status.idle":"2024-05-19T00:23:27.175099Z","shell.execute_reply.started":"2024-05-19T00:23:27.151436Z","shell.execute_reply":"2024-05-19T00:23:27.173926Z"},"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.55, # 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":{"execution":{"iopub.status.busy":"2024-05-19T00:23:27.176829Z","iopub.execute_input":"2024-05-19T00:23:27.17772Z","iopub.status.idle":"2024-05-19T00:23:27.196887Z","shell.execute_reply.started":"2024-05-19T00:23:27.177686Z","shell.execute_reply":"2024-05-19T00:23:27.195959Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Keypoints: LightGlue series","metadata":{}},{"cell_type":"code","source":"def CLAHE_Convert(origin_input):\n    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))\n    t = np.asarray(origin_input)\n    t = cv2.cvtColor(t, cv2.COLOR_BGR2HSV)\n    t[:,:,-1] = clahe.apply(t[:,:,-1])\n    t = cv2.cvtColor(t, cv2.COLOR_HSV2BGR)\n    # t = Img.fromarray(t)\n    return t","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def load_torch_image(fname, device=torch.device('cpu')):\n    img = CLAHE_Convert(cv2.imread(fname))\n    img = K.image_to_tensor(img, False).float() / 255.\n    img = K.color.bgr_to_rgb(img.to(device))\n    return img\n\ndef detect_common(img_fnames,\n                  model_name,\n                  feature_dir = '.featureout',\n                  num_features = 4096,\n                  resize_to = 1380,\n                  detection_threshold = 0.008,\n                  device=torch.device('cpu')):\n    \n    dict_model = {\n        'disk':DISK,\n        \"aliked\" : ALIKED,\n        'superpoint':SuperPoint\n    }\n    extractor_class = dict_model[model_name]\n    \n    dtype = torch.float32 # ALIKED has issues with float16\n    extractor = extractor_class(\n        max_num_keypoints=num_features, detection_threshold=detection_threshold, resize=resize_to\n    ).eval().to(device, dtype)\n    if not os.path.isdir(feature_dir):\n        os.makedirs(feature_dir)\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 tqdm(img_fnames):\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                feats0 = extractor.extract(image0)  # auto-resize the image, disable with resize=None\n                kpts = feats0['keypoints'].reshape(-1, 2).detach().cpu().numpy()\n                descs = feats0['descriptors'].reshape(len(kpts), -1).detach().cpu().numpy()\n                f_kp[key] = kpts\n                f_desc[key] = descs\n                print(f\"{model_name} > kpts.shape={kpts.shape}, descs.shape={descs.shape}\")\n    return\n\ndef match_with_lightglue_common(img_fnames, model_name,\n                   index_pairs, file_keypoints,\n                   feature_dir = '.featureout',\n                   device=torch.device('cpu'),\n                   min_matches=15,verbose=True):\n    lg_matcher = KF.LightGlueMatcher(model_name, {\"width_confidence\": -1,\n                                                \"depth_confidence\": -1,\n                                                 \"mp\": True if 'cuda' in str(device) else False}).eval().to(device)\n    \n    cnt_pairs = 0\n    with h5py.File(f'{feature_dir}/keypoints_{model_name}.h5', mode='r') as f_kp, \\\n        h5py.File(f'{feature_dir}/descriptors_{model_name}.h5', mode='r') as f_desc, \\\n        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            kp1 = torch.from_numpy(f_kp[key1][...]).to(device)\n            kp2 = torch.from_numpy(f_kp[key2][...]).to(device)\n            desc1 = torch.from_numpy(f_desc[key1][...]).to(device)\n            desc2 = torch.from_numpy(f_desc[key2][...]).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            n_matches = len(idxs)\n            kp1 = kp1[idxs[:,0], :].cpu().numpy().reshape(-1, 2).astype(np.float32)\n            kp2 = kp2[idxs[:,1], :].cpu().numpy().reshape(-1, 2).astype(np.float32)\n            group  = f_match.require_group(key1)\n            if n_matches >= min_matches:\n                group.create_dataset(key2, data=np.concatenate([kp1, kp2], axis=1))\n                cnt_pairs+=1\n                print (f'{key1}-{key2}: {n_matches} matches @ {cnt_pairs}th pair({model_name}+lightglue)')            \n            else:\n                print (f'{key1}-{key2}: {n_matches} matches --> skipped')            \n    return\n\ndef detect_lightglue_common(\n    img_fnames, model_name, index_pairs, feature_dir, device, file_keypoints,\n    resize_to=1024,\n    detection_threshold=0.01, \n    num_features=4096, \n    min_matches=15,\n):\n    t=time()\n    detect_common(\n        img_fnames, model_name, feature_dir, \n        resize_to=resize_to,\n        num_features=num_features, \n        detection_threshold=detection_threshold, \n        device=device\n    )\n    gc.collect()\n    match_with_lightglue_common(\n        img_fnames, model_name, index_pairs, file_keypoints,\n        feature_dir=feature_dir,\n        min_matches=min_matches,\n        device=device\n    )\n    t=time() -t \n    print(f'Features matched in  {t:.4f} sec ({model_name}+LightGlue)')\n    return t","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-05-19T00:23:27.198219Z","iopub.execute_input":"2024-05-19T00:23:27.198575Z","iopub.status.idle":"2024-05-19T00:23:27.222787Z","shell.execute_reply.started":"2024-05-19T00:23:27.19855Z","shell.execute_reply":"2024-05-19T00:23:27.221986Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Keypoints: SuperGlue","metadata":{}},{"cell_type":"code","source":"import sys\nsys.path.append(\"../input/super-glue-pretrained-network\")\nfrom models.matching import Matching\nfrom models.utils import (compute_pose_error, compute_epipolar_error,\n                          estimate_pose, make_matching_plot,\n                          error_colormap, AverageTimer, pose_auc, read_image,\n                          process_resize, frame2tensor,\n                          rotate_intrinsics, rotate_pose_inplane,\n                          scale_intrinsics)\n\nfrom torch.nn import functional as torchF  # For resizing tensor\n\ndef sg_imread(path):\n    image = cv2.imread(str(path), cv2.IMREAD_GRAYSCALE)\n    return image\n\n# Preprocess\ndef sg_read_image(image, device, resize):\n    w, h = image.shape[1], image.shape[0]\n    w_new, h_new = process_resize(w, h, [resize,])\n    \n    unit_shape = 8\n    w_new = w_new // unit_shape * unit_shape\n    h_new = h_new // unit_shape * unit_shape\n    \n    scales = (float(w) / float(w_new), float(h) / float(h_new))\n    image = cv2.resize(image.astype('float32'), (w_new, h_new))\n\n    inp = frame2tensor(image, \"cpu\")\n    return image, inp, scales, (h, w)\n\nclass SGDataset(Dataset):\n    def __init__(self, fnames1, fnames2, resize_to, device):\n        self.fnames1 = fnames1\n        self.fnames2 = fnames2\n        self.resize_to = resize_to\n        self.device = device\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 = cv2.imread(fname1, cv2.IMREAD_GRAYSCALE), cv2.imread(fname2, cv2.IMREAD_GRAYSCALE)\n        _, image1, scale1, ori_shape1 = sg_read_image(im1, self.device, self.resize_to)\n        _, image2, scale2, ori_shape2 = sg_read_image(im2, self.device, self.resize_to)\n        return image1, image2, torch.tensor([idx]), torch.tensor(ori_shape1), torch.tensor(ori_shape2)\n\ndef get_superglue_dataloader(images1, images2, resize_to, device, batch_size=1):\n    dataset = SGDataset(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 detect_superglue(\n    img_fnames, index_pairs, feature_dir, device, sg_config, file_keypoints, \n    resize_to=750, min_matches=15\n):    \n    t=time()\n\n    matcher_superglue = Matching(sg_config).eval().to(device)\n\n    fnames1, fnames2, idxs1, idxs2 = [], [], [], []\n    for pair_idx in progress_bar(index_pairs):\n        idx1, idx2 = pair_idx\n        fname1, fname2 = img_fnames[idx1], img_fnames[idx2]\n        fnames1.append(fname1)\n        fnames2.append(fname2)\n        idxs1.append(idx1)\n        idxs2.append(idx2)\n        \n    dataloader = get_superglue_dataloader( fnames1, fnames2, resize_to, device)\n\n    cnt_pairs = 0\n\n    with h5py.File(file_keypoints, mode='w') as f_match:\n        for X in dataloader:\n            image1, image2, idx, ori_shape_1, ori_shape_2 = X\n\n            fname1, fname2 = fnames1[idx], fnames2[idx]\n            key1, key2 = fname1.split('/')[-1], fname2.split('/')[-1]\n            \n            pred = matcher_superglue({\"image0\": image1[0].to(device), \"image1\": image2[0].to(device)})\n            pred = {k: v[0].detach().cpu().numpy().copy() for k, v in pred.items()}\n            mkpts1, mkpts2 = pred[\"keypoints0\"], pred[\"keypoints1\"]\n            matches, conf = pred[\"matches0\"], pred[\"matching_scores0\"]\n\n            valid = matches > -1\n            mkpts1 = mkpts1[valid]\n            mkpts2 = mkpts2[matches[valid]]\n            mconf = conf[valid]\n\n            ori_shape_1 = ori_shape_1[0].numpy()\n            ori_shape_2 = ori_shape_2[0].numpy()\n            \n            # Scaling coords\n            mkpts1[:,0] = mkpts1[:,0] * ori_shape_1[1] / image1[0].shape[3]   # X\n            mkpts1[:,1] = mkpts1[:,1] * ori_shape_1[0] / image1[0].shape[2]   # Y\n            mkpts2[:,0] = mkpts2[:,0] * ori_shape_2[1] / image2[0].shape[3]   # X\n            mkpts2[:,1] = mkpts2[:,1] * ori_shape_2[0] / image2[0].shape[2]   # Y  \n            \n            n_matches = mconf.shape[0]\n            \n            group  = f_match.require_group(key1)\n            if n_matches >= min_matches:\n                group.create_dataset(key2, data=np.concatenate([mkpts1, mkpts2], axis=1).astype(np.float32))\n                cnt_pairs+=1\n                print (f'{key1}-{key2}: {n_matches} matches @ {cnt_pairs}th pair(superglue)')            \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-19T00:23:27.224197Z","iopub.execute_input":"2024-05-19T00:23:27.224634Z","iopub.status.idle":"2024-05-19T00:23:27.291624Z","shell.execute_reply.started":"2024-05-19T00:23:27.224607Z","shell.execute_reply":"2024-05-19T00:23:27.290651Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Keypoints: DKM","metadata":{}},{"cell_type":"code","source":"class DKMDataset(Dataset):\n    def __init__(self, fnames1, fnames2, resize_to, device):\n        self.fnames1 = fnames1\n        self.fnames2 = fnames2\n        self.resize_to = resize_to\n        self.device = device\n        self.test_transform = get_tuple_transform_ops(\n            resize=self.resize_to, normalize=True\n        )\n\n        \n    def __len__(self):\n        return len(self.fnames1)\n    \n    def __getitem__(self, idx):\n        fname1 = self.fnames1[idx]\n        fname2 = self.fnames2[idx]\n                \n        im1, im2 = Image.open(fname1), Image.open(fname2)\n        ori_shape_1 = im1.size\n        ori_shape_2 = im2.size\n        image1, image2 = self.test_transform((im1, im2))\n        return image1, image2, torch.tensor([idx]), torch.tensor(ori_shape_1), torch.tensor(ori_shape_2)\n\ndef get_dkm_dataloader(images1, images2, resize_to, device, batch_size=4):\n    dataset = DKMDataset(images1, images2, resize_to, device)\n    dataloader = DataLoader(\n        dataset=dataset,\n        shuffle=False,\n        batch_size=batch_size,\n        pin_memory=True,\n        num_workers=2,\n        drop_last=False\n    )\n    return dataloader\n\ndef get_dkm_mkpts(dkm_model, bimgs1, bimgs2, shapes1, shapes2, detection_threshold=0.5, num_features = 2000, min_matches=15):\n    dense_matches, dense_certainty = dkm_model.match(bimgs1, bimgs2, batched=True)\n    print(\"***\", dense_matches.shape, dense_certainty.shape)\n\n    store_mkpts1, store_mkpts2, store_mconf = [], [], []\n    # drop low confidence pairs\n    for b in range(dense_matches.shape[0]):\n        u_dense_matches = dense_matches[b, dense_certainty[b,...].sqrt() >= detection_threshold, :]\n        u_dense_certainty = dense_certainty[b, dense_certainty[b,...].sqrt() >= detection_threshold]\n    \n        if u_dense_matches.shape[0] > num_features:\n            u_dense_matches, u_dense_certainty = dkm_model.sample( u_dense_matches, u_dense_certainty, num=num_features)\n        \n        u_dense_matches = u_dense_matches.reshape((-1, 4))\n        u_dense_certainty = u_dense_certainty.reshape((-1,))\n    \n        mkpts1 = u_dense_matches[:, :2]\n        mkpts2 = u_dense_matches[:, 2:]\n        \n        w1, h1 = shapes1[b, :]\n        w2, h2 = shapes2[b, :]\n\n        mkpts1[:, 0] = ((mkpts1[:, 0] + 1)/2) * w1\n        mkpts1[:, 1] = ((mkpts1[:, 1] + 1)/2) * h1\n\n        mkpts2[:, 0] = ((mkpts2[:, 0] + 1)/2) * w2\n        mkpts2[:, 1] = ((mkpts2[:, 1] + 1)/2) * h2\n\n        mkpts1 = mkpts1.cpu().detach().numpy()\n        mkpts2 = mkpts2.cpu().detach().numpy()\n        mconf  = u_dense_certainty.sqrt().cpu().detach().numpy()\n\n        \n        if mconf.shape[0] > min_matches:\n            try:\n                # calc Fundamental matrix from keypoints\n                F, inliers = cv2.findFundamentalMat(mkpts1, mkpts2, cv2.USAC_MAGSAC, 0.200, 0.999, 2000)\n                inliers = inliers > 0\n                mkpts1 = mkpts1[inliers[:,0]]\n                mkpts2 = mkpts2[inliers[:,0]]\n                mconf  = mconf[inliers[:,0]]\n                #print(\"---\", mconf.shape)\n                if mconf.shape[0] > 3000:\n                    rand_idx = np.random.choice(range(mconf.shape[0]), 3000, replace=False)\n                    mkpts1 = mkpts1[rand_idx, :]\n                    mkpts2 = mkpts2[rand_idx, :]\n                    mconf  = mconf[rand_idx]\n            except:\n                mkpts1 = np.empty((0,2))\n                mkpts2 = np.empty((0,2))\n                mconf = np.empty((0,))\n        \n        store_mkpts1.append(mkpts1)\n        store_mkpts2.append(mkpts2)\n        store_mconf.append(mconf)\n    return store_mkpts1, store_mkpts2, store_mconf\n\ndef detect_dkm(\n    img_fnames, index_pairs, feature_dir, device, \n    resize_to=(540, 720), \n    detection_threshold=0.4, \n    num_features=2000, \n    min_matches=15,\n):\n    t=time()\n    dkm_model = DKMv3_outdoor(device=device)\n    dkm_model.upsample_preds=False\n\n    fnames1, fnames2 = [], []\n    for pair_idx in progress_bar(index_pairs):\n        idx1, idx2 = pair_idx\n        fname1, fname2 = img_fnames[idx1], img_fnames[idx2]\n        fnames1.append(fname1)\n        fnames2.append(fname2)\n        \n    cnt_pairs = 0\n    with h5py.File(f'{feature_dir}/matches_dkm.h5', mode='w') as f_match:    \n        dataloader = get_dkm_dataloader(fnames1, fnames2, resize_to, device, batch_size=4)\n        for X in tqdm(dataloader):\n            images1, images2, idxs, shapes1, shapes2 = X\n            store_mkpts1, store_mkpts2, store_mconf = get_dkm_mkpts(\n                dkm_model, images1.to(device), images2.to(device), shapes1, shapes2, \n                detection_threshold=detection_threshold, num_features = num_features, min_matches=min_matches,\n            )\n            \n            for b in range(images1.shape[0]):\n                mkpts1 = store_mkpts1[b]\n                mkpts2 = store_mkpts2[b]\n                mconf = store_mconf[b]\n                file1 = fnames1[idxs[b]]\n                file2 = fnames2[idxs[b]]\n                key1, key2 = file1.split('/')[-1], file2.split('/')[-1]\n            \n                n_matches = mconf.shape[0]\n                print (f'{key1}-{key2}: {n_matches} matches @ {cnt_pairs}th pair(dkm)')            \n\n                group  = f_match.require_group(key1)\n                if n_matches >= min_matches:\n                    group.create_dataset(key2, data=np.concatenate([mkpts1, mkpts2], axis=1).astype(np.float32))\n                    cnt_pairs+=1\n    gc.collect()\n    t=time() -t \n    print(f'Features matched in  {t:.4f} sec')\n    return t","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-05-19T00:23:27.293061Z","iopub.execute_input":"2024-05-19T00:23:27.293359Z","iopub.status.idle":"2024-05-19T00:23:27.32359Z","shell.execute_reply.started":"2024-05-19T00:23:27.293335Z","shell.execute_reply":"2024-05-19T00:23:27.322517Z"},"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 = CLAHE_Convert(cv2.imread(fname))\n        original_shape = img.shape\n        ratio = self.resize_small_edge_to / min([img.shape[0], img.shape[1]])\n        w = int(img.shape[1] * ratio) # int( (img.shape[1] * ratio) // self.round_unit * self.round_unit )\n        h = int(img.shape[0] * ratio) # int( (img.shape[0] * ratio) // self.round_unit * self.round_unit )\n        img_resized = cv2.resize(img, (w, h))\n        img_resized = K.image_to_tensor(img_resized, False).float() /255.\n        img_resized = K.color.bgr_to_rgb(img_resized)\n        img_resized = K.color.rgb_to_grayscale(img_resized)\n        return img_resized.to(device), original_shape\n    \n    def __getitem__(self, idx):\n        fname1 = self.fnames1[idx]\n        fname2 = self.fnames2[idx]\n        image1, ori_shape_1 = self.load_torch_image(fname1, device)\n        image2, ori_shape_2 = self.load_torch_image(fname2, device)\n\n        return image1, image2, self.keys1[idx], self.keys2[idx], self.idxs1[idx], self.idxs2[idx], ori_shape_1, ori_shape_2\n\ndef get_loftr_dataloader(images1, images2, idxs1, idxs2, resize_small_edge_to, device, batch_size=1):\n    dataset = LoFTRDataset(images1, images2, idxs1, idxs2, resize_small_edge_to, device)\n    dataloader = DataLoader(\n        dataset=dataset,\n        shuffle=False,\n        batch_size=batch_size,\n        pin_memory=True,\n        num_workers=2,\n        drop_last=False\n    )\n    return dataset\n    \ndef detect_loftr(img_fnames, index_pairs, feature_dir, device, file_keypoints, resize_small_edge_to=750, min_matches=15):\n    t=time()\n\n    matcher = LoFTR(pretrained=None)\n    matcher.load_state_dict(torch.load(\"../input/loftr/pytorch/outdoor/1/loftr_outdoor.ckpt\")['state_dict'])\n    matcher = matcher.to(device).eval()\n\n    fnames1, fnames2, idxs1, idxs2 = [], [], [], []\n    for pair_idx in progress_bar(index_pairs):\n        idx1, idx2 = pair_idx\n        fname1, fname2 = img_fnames[idx1], img_fnames[idx2]\n        fnames1.append(fname1)\n        fnames2.append(fname2)\n        idxs1.append(idx1)\n        idxs2.append(idx2)\n        \n        \n    dataloader = get_loftr_dataloader( fnames1, fnames2, idxs1, idxs2, resize_small_edge_to, device)\n\n    cnt_pairs = 0\n\n    with h5py.File(file_keypoints, mode='w') as f_match:    \n        store_mkpts = {}\n        for X in tqdm(dataloader):\n            image1, image2, key1, key2, idx1, idx2, ori_shape_1, ori_shape_2 = X\n            fname1, fname2 = img_fnames[idx1], img_fnames[idx2]\n\n            with torch.no_grad():\n                correspondences = matcher( {\"image0\": image1.to(device),\"image1\": image2.to(device)} )\n                mkpts1 = correspondences['keypoints0'].cpu().numpy()\n                mkpts2 = correspondences['keypoints1'].cpu().numpy()\n                mconf  = correspondences['confidence'].cpu().numpy()\n\n            mkpts1[:,0] *= (float(ori_shape_1[1]) / float(image1.shape[3]))\n            mkpts1[:,1] *= (float(ori_shape_1[0]) / float(image1.shape[2]))\n\n            mkpts2[:,0] *= (float(ori_shape_2[1]) / float(image2.shape[3]))\n            mkpts2[:,1] *= (float(ori_shape_2[0]) / float(image2.shape[2]))\n            \n            n_matches = mconf.shape[0]\n            \n            group  = f_match.require_group(key1)\n            if n_matches >= min_matches:\n                group.create_dataset(key2, data=np.concatenate([mkpts1, mkpts2], axis=1).astype(np.float32))\n                cnt_pairs+=1\n                print (f'{key1}-{key2}: {n_matches} matches @ {cnt_pairs}th pair(loftr)')\n            else:\n                print (f'{key1}-{key2}: {n_matches} matches --> skipped')\n    gc.collect()\n    t=time() -t \n    print(f'Features matched in  {t:.4f} sec')\n    return t","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-05-19T00:23:27.324876Z","iopub.execute_input":"2024-05-19T00:23:27.325214Z","iopub.status.idle":"2024-05-19T00:23:27.350724Z","shell.execute_reply.started":"2024-05-19T00:23:27.325189Z","shell.execute_reply":"2024-05-19T00:23:27.349701Z"},"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-19T00:23:27.352137Z","iopub.execute_input":"2024-05-19T00:23:27.352522Z","iopub.status.idle":"2024-05-19T00:23:27.382819Z","shell.execute_reply.started":"2024-05-19T00:23:27.352488Z","shell.execute_reply":"2024-05-19T00:23:27.381714Z"},"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 keypoints_merger(\n    img_fnames,\n    index_pairs,\n    files_keypoints,\n    feature_dir = 'featureout',\n):\n    # open h5 files\n    fps = [ h5py.File(file, mode=\"r\") for file in files_keypoints ]\n\n    # temprary file\n    with h5py.File(f'{feature_dir}/merge_tmp.h5', 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            print (f'{key1}-{key2}: {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    \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(f'{feature_dir}/merge_tmp.h5', 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-19T00:23:27.384459Z","iopub.execute_input":"2024-05-19T00:23:27.384889Z","iopub.status.idle":"2024-05-19T00:23:27.413504Z","shell.execute_reply.started":"2024-05-19T00:23:27.384854Z","shell.execute_reply":"2024-05-19T00:23:27.412447Z"},"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-19T00:23:27.41471Z","iopub.execute_input":"2024-05-19T00:23:27.415021Z","iopub.status.idle":"2024-05-19T00:23:27.426733Z","shell.execute_reply.started":"2024-05-19T00:23:27.414987Z","shell.execute_reply":"2024-05-19T00:23:27.425785Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Main","metadata":{}},{"cell_type":"code","source":"src = '/kaggle/input/image-matching-challenge-2024'\nDEBUG = True\nDUMP = True\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 len(data_dict[dataset][scene]) == 21:\n            #    break","metadata":{"execution":{"iopub.status.busy":"2024-05-19T00:23:27.427771Z","iopub.execute_input":"2024-05-19T00:23:27.428048Z","iopub.status.idle":"2024-05-19T00:23:27.442862Z","shell.execute_reply.started":"2024-05-19T00:23:27.428025Z","shell.execute_reply":"2024-05-19T00:23:27.442037Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for 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-19T00:23:27.44399Z","iopub.execute_input":"2024-05-19T00:23:27.44441Z","iopub.status.idle":"2024-05-19T00:23:27.450231Z","shell.execute_reply.started":"2024-05-19T00:23:27.444355Z","shell.execute_reply":"2024-05-19T00:23:27.449138Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"out_results = {}\ntimings = {\"shortlisting\":[],\n           \"feature_detection\": [],\n           \"feature_matching\":[],\n           \"RANSAC\": [],\n           \"Reconstruction\": []}","metadata":{"execution":{"iopub.status.busy":"2024-05-19T00:23:27.451367Z","iopub.execute_input":"2024-05-19T00:23:27.451831Z","iopub.status.idle":"2024-05-19T00:23:27.457861Z","shell.execute_reply.started":"2024-05-19T00:23:27.451799Z","shell.execute_reply":"2024-05-19T00:23:27.456853Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"gc.collect()\ndatasets = []\nfor dataset in data_dict:\n    datasets.append(dataset)\n\nfor 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        # Wrap the meaty part in a try-except block.\n        try:\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            # get image rotations\n            #############################################################\n            t = time()\n            if CONFIG.ROTATION_CORRECTION:\n                rots = exec_rotation_detection(img_fnames, device)\n            else:\n                rots = [ 0 for fname in img_fnames ]\n            t = time()-t\n            timings['rotation_detection'].append(t)\n            print (f'rotation_detection for {len(img_fnames)} images : {t:.4f} sec')\n            gc.collect()\n            #############################################################\n            # get image pairs\n            #############################################################\n            t=time()\n            index_pairs = get_image_pairs_shortlist(img_fnames,\n                                  sim_th = 1.0, # should be strict\n                                  min_pairs = 50, # we select at least min_pairs PER IMAGE with biggest similarity\n                                  exhaustive_if_less = 50,\n                                  device=device)\n            t=time() -t \n            timings['shortlisting'].append(t)\n            print (f'{len(index_pairs)}, pairs to match, {t:.4f} sec')\n            gc.collect()\n            \n            #############################################################\n            # get keypoints\n            #############################################################\n            files_keypoints = []\n            if CONFIG.use_superglue:\n                resize_to = CONFIG.params_sg[\"resize_to\"]\n                file_keypoints = f\"{feature_dir}/matches_superglue_{resize_to}pix.h5\"\n                !rm -rf {file_keypoints}\n                t = detect_superglue(\n                    img_fnames, index_pairs, feature_dir, device, \n                    CONFIG.params_sg[\"sg_config\"], file_keypoints,\n                    resize_to=CONFIG.params_sg[\"resize_to\"], \n                    min_matches=CONFIG.params_sg[\"min_matches\"],\n                )\n                gc.collect()\n                files_keypoints.append( file_keypoints )\n                timings['feature_matching'].append(t)\n\n            if CONFIG.use_aliked_lightglue:\n                model_name = \"aliked\"\n                file_keypoints = f'{feature_dir}/matches_lightglue_{model_name}.h5'\n                t = detect_lightglue_common(\n                    img_fnames, model_name, index_pairs, feature_dir, device, file_keypoints,\n                    resize_to=CONFIG.params_aliked_lightglue[\"resize_to\"],\n                    detection_threshold=CONFIG.params_aliked_lightglue[\"detection_threshold\"],\n                    num_features=CONFIG.params_aliked_lightglue[\"num_features\"],\n                    min_matches=CONFIG.params_aliked_lightglue[\"min_matches\"],\n                )\n                gc.collect()\n                files_keypoints.append(file_keypoints)\n                timings['feature_matching'].append(t)\n                \n            if CONFIG.use_doghardnet_lightglue:\n                model_name = \"doghardnet\"\n                file_keypoints = f'{feature_dir}/matches_lightglue_{model_name}.h5'\n                t = detect_lightglue_common(\n                    img_fnames, model_name, index_pairs, feature_dir, device, file_keypoints,\n                    resize_to=CONFIG.params_doghardnet_lightglue[\"resize_to\"],\n                    detection_threshold=CONFIG.params_doghardnet_lightglue[\"detection_threshold\"],\n                    num_features=CONFIG.params_doghardnet_lightglue[\"num_features\"],\n                    min_matches=CONFIG.params_doghardnet_lightglue[\"min_matches\"],\n                )\n                gc.collect()\n                files_keypoints.append(file_keypoints)\n                timings['feature_matching'].append(t)\n                \n            if CONFIG.use_superpoint_lightglue:\n                model_name = \"superpoint\"\n                file_keypoints = f'{feature_dir}/matches_lightglue_{model_name}.h5'\n                t = detect_lightglue_common(\n                    img_fnames, model_name, index_pairs, feature_dir, device, file_keypoints,\n                    resize_to=CONFIG.params_superpoint_lightglue[\"resize_to\"],\n                    detection_threshold=CONFIG.params_superpoint_lightglue[\"detection_threshold\"],\n                    num_features=CONFIG.params_superpoint_lightglue[\"num_features\"],\n                    min_matches=CONFIG.params_superpoint_lightglue[\"min_matches\"],\n                )\n                gc.collect()\n                files_keypoints.append(file_keypoints)\n                timings['feature_matching'].append(t)\n                \n            if CONFIG.use_loftr:\n                file_keypoints = f'{feature_dir}/matches_loftr_{CONFIG.params_loftr[\"resize_small_edge_to\"]}pix.h5'\n                t = detect_loftr(\n                    img_fnames, index_pairs, feature_dir, device, file_keypoints,\n                    resize_small_edge_to=CONFIG.params_loftr[\"resize_small_edge_to\"],\n                    min_matches=CONFIG.params_loftr[\"min_matches\"],\n                )\n                gc.collect()\n                files_keypoints.append( file_keypoints )\n                timings['feature_matching'].append(t)\n                \n            if CONFIG.use_dkm:\n                file_keypoints = f'{feature_dir}/matches_dkm.h5'\n                t = detect_dkm(\n                    img_fnames, index_pairs, feature_dir, device, file_keypoints,\n                    resize_to=CONFIG.params_dkm[\"resize_to\"], \n                    detection_threshold=CONFIG.params_dkm[\"detection_threshold\"], \n                    num_features=CONFIG.params_dkm[\"num_features\"], \n                    min_matches=CONFIG.params_dkm[\"min_matches\"]\n                )\n                gc.collect()\n                files_keypoints.append(file_keypoints)\n                timings['feature_matching'].append(t)\n                \n            #############################################################\n            # merge keypoints\n            #############################################################\n            keypoints_merger(\n                img_fnames,\n                index_pairs,\n                files_keypoints,\n                feature_dir = feature_dir,\n            )\n            \n            #############################################################\n            # regist keypoints from h5 into colmap db\n            #############################################################\n            database_path = f'{feature_dir}/colmap.db'\n            if os.path.isfile(database_path):\n                os.remove(database_path)\n            gc.collect()\n            import_into_colmap(img_dir, feature_dir=feature_dir,database_path=database_path)\n            output_path = f'{feature_dir}/colmap_rec'\n\n            #############################################################\n            # Calculate fundamental matrix with colmap api\n            #############################################################\n            t=time()\n            options = pycolmap.SiftMatchingOptions()\n            options.confidence = 0.9995\n            options.max_num_trials = 5000\n            pycolmap.match_exhaustive(database_path, sift_options=options)\n            t=time() - t \n            timings['RANSAC'].append(t)\n            print(f'RANSAC in  {t:.4f} sec')\n\n            #############################################################\n            # Execute bundle adjustmnet with colmap api\n            # --> Bundle adjustment Calcs Camera matrix, R and t\n            #############################################################\n            t=time()\n            # By default colmap does not generate a reconstruction if less than 10 images are registered. Lower it to 3.\n            mapper_options = pycolmap.IncrementalMapperOptions()\n            mapper_options.min_model_size = 3\n            os.makedirs(output_path, exist_ok=True)\n            maps = pycolmap.incremental_mapping(database_path=database_path, image_path=img_dir, output_path=output_path, options=mapper_options)\n            print(maps)\n            clear_output(wait=False)\n            t=time() - t\n            timings['Reconstruction'].append(t)\n            print(f'Reconstruction done in  {t:.4f} sec')\n\n            #############################################################\n            # Extract R,t from maps \n            #############################################################            \n            imgs_registered  = 0\n            best_idx = None\n            list_num_images = []            \n            print (\"Looking for the best reconstruction\")\n            if isinstance(maps, dict):\n                for idx1, rec in maps.items():\n                    print (idx1, rec.summary())\n                    list_num_images.append( len(rec.images) )\n                    if len(rec.images) > imgs_registered:\n                        imgs_registered = len(rec.images)\n                        best_idx = idx1\n            list_num_images = np.array(list_num_images)\n            print(f\"list_num_images = {list_num_images}\")\n            if best_idx is not None:\n                print (maps[best_idx].summary())\n                for k, im in maps[best_idx].images.items():\n                    key1 = f'test/{dataset}/images/{im.name}'\n                    out_results[dataset][scene][key1] = {}\n                    out_results[dataset][scene][key1][\"R\"] = deepcopy(im.rotmat())\n                    out_results[dataset][scene][key1][\"t\"] = deepcopy(np.array(im.tvec))\n                \n            print(f'Registered: {dataset} / {scene} -> {len(out_results[dataset][scene])} images')\n            print(f'Total: {dataset} / {scene} -> {len(data_dict[dataset][scene])} images')\n            create_submission(out_results, data_dict)\n            gc.collect()\n        except:\n            pass","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2024-05-19T00:23:27.459096Z","iopub.execute_input":"2024-05-19T00:23:27.45949Z","iopub.status.idle":"2024-05-19T00:33:26.703546Z","shell.execute_reply.started":"2024-05-19T00:23:27.459465Z","shell.execute_reply":"2024-05-19T00:33:26.702321Z"},"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-19T00:33:26.705451Z","iopub.execute_input":"2024-05-19T00:33:26.705784Z","iopub.status.idle":"2024-05-19T00:33:27.756541Z","shell.execute_reply.started":"2024-05-19T00:33:26.70575Z","shell.execute_reply":"2024-05-19T00:33:27.755412Z"},"trusted":true},"execution_count":null,"outputs":[]}]}