{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Baseline submission\n\nA notebook to generate a valid submission. Implements three local feature/matcher methods: LoFTR, DISK, and KeyNetAffNetHardNet.\n\nRemember to enable a GPU accelerator and disable internet access, then press \"submit\" on the right pane.","metadata":{}},{"cell_type":"code","source":"# Get some utility functions from https://github.com/cvg/Hierarchical-Localization/\n\n%cd /kaggle/input/hierarchical-localization-master\n\nfrom hloc import extract_features, match_features, match_dense\n\n%cd /kaggle/working/","metadata":{"execution":{"iopub.status.busy":"2023-06-03T02:12:22.207596Z","iopub.execute_input":"2023-06-03T02:12:22.208212Z","iopub.status.idle":"2023-06-03T02:12:22.220988Z","shell.execute_reply.started":"2023-06-03T02:12:22.208161Z","shell.execute_reply":"2023-06-03T02:12:22.219585Z"},"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 h5py\nfrom IPython.display import clear_output\nfrom collections import defaultdict\nfrom copy import deepcopy\nfrom pathlib import Path\n\n# CV/ML\nimport cv2\nimport torch\nimport torch.nn.functional as F\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\n# 3D reconstruction\nimport pycolmap","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-06-03T02:12:22.223780Z","iopub.execute_input":"2023-06-03T02:12:22.225095Z","iopub.status.idle":"2023-06-03T02:12:22.234106Z","shell.execute_reply.started":"2023-06-03T02:12:22.225053Z","shell.execute_reply":"2023-06-03T02:12:22.232985Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print('Kornia version', K.__version__)\nprint('Pycolmap version', pycolmap.__version__)\n\n# Sparse\n# SP\nfeature_conf = extract_features.confs['superpoint_max']\nfeature_conf['model']['nms_radius'] = 5\nfeature_conf['model']['max_keypoints'] = 8000\nfeature_conf['preprocessing']['resize_max'] = 2200\n\n# R2D2\n\"\"\"feature_conf = extract_features.confs['r2d2']\nfeature_conf['model']['max_keypoints'] = 5000 # 5000\nfeature_conf['preprocessing']['resize_max'] = 2000 # 1600\"\"\"\n\n# SG\nmatcher_conf = match_features.confs['superglue']\nmatcher_conf['model']['sinkhorn_iterations'] = 50\n\n# other matchers\n#matcher_conf = match_features.confs['NN-ratio'] # adalam NN-ratio NN-mutual\n\n# Dense\ndense_conf = match_dense.confs['loftr_aachen']\ndense_conf['preprocessing']['resize_max'] = 1024\ndense_conf['max_error'] = 2\ndense_conf['cell_size'] = 8\nmax_keypoints = 8192\n\n# config\nmatcher = False\npairs = 50\nthr = 0.5\n\ndevice=torch.device('cuda')\nprint(feature_conf)\nprint(matcher_conf)\nprint(dense_conf)","metadata":{"execution":{"iopub.status.busy":"2023-06-03T02:12:22.237171Z","iopub.execute_input":"2023-06-03T02:12:22.237961Z","iopub.status.idle":"2023-06-03T02:12:22.251234Z","shell.execute_reply.started":"2023-06-03T02:12:22.237929Z","shell.execute_reply":"2023-06-03T02:12:22.249922Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def arr_to_str(a):\n    return ';'.join([str(x) for x in a.reshape(-1)])\n\n\ndef load_torch_image(fname, device=torch.device('cpu')):\n    img = K.image_to_tensor(cv2.imread(fname), False).float() / 255.\n    img = K.color.bgr_to_rgb(img.to(device))\n    return img","metadata":{"execution":{"iopub.status.busy":"2023-06-03T02:12:22.254222Z","iopub.execute_input":"2023-06-03T02:12:22.254757Z","iopub.status.idle":"2023-06-03T02:12:22.263967Z","shell.execute_reply.started":"2023-06-03T02:12:22.254717Z","shell.execute_reply":"2023-06-03T02:12:22.262940Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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():\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\n\n\ndef get_img_pairs_exhaustive(img_fnames):\n    index_pairs = []\n    for i in range(len(img_fnames)):\n        for j in range(i+1, len(img_fnames)):\n            index_pairs.append((i,j))\n    matching_list = sorted(list(set(index_pairs)))\n    matching_list_path = os.path.join(feature_dir, 'pair_list.txt')\n    with open(matching_list_path, 'w') as f:\n        for (idx1, idx2) in matching_list:\n            img_name1, img_name2 = img_fnames[idx1].split('/')[-1], img_fnames[idx2].split('/')[-1]\n            f.write(img_name1 + ' ' + img_name2 + '\\n')\n    return Path(matching_list_path)\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    if num_imgs <= exhaustive_if_less:\n        return get_img_pairs_exhaustive(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    print('finish global desc extraction')\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    matching_list_path = os.path.join(feature_dir, 'pair_list.txt')\n    with open(matching_list_path, 'w') as f:\n        for (idx1, idx2) in matching_list:\n            img_name1, img_name2 = fnames[idx1].split('/')[-1], fnames[idx2].split('/')[-1]\n            f.write(img_name1 + ' ' + img_name2 + '\\n')\n    return Path(matching_list_path)","metadata":{"execution":{"iopub.status.busy":"2023-06-03T02:12:22.265906Z","iopub.execute_input":"2023-06-03T02:12:22.266746Z","iopub.status.idle":"2023-06-03T02:12:22.291865Z","shell.execute_reply.started":"2023-06-03T02:12:22.266705Z","shell.execute_reply":"2023-06-03T02:12:22.290760Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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":{"execution":{"iopub.status.busy":"2023-06-03T02:12:22.329788Z","iopub.execute_input":"2023-06-03T02:12:22.330548Z","iopub.status.idle":"2023-06-03T02:12:22.368518Z","shell.execute_reply.started":"2023-06-03T02:12:22.330499Z","shell.execute_reply":"2023-06-03T02:12:22.367284Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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(str(h5_path), 'r')\n\n    camera_id = None\n    fname_to_id = {}\n    for filename in tqdm(list(keypoint_f.keys())):\n        keypoints = keypoint_f[filename]['keypoints'].__array__() # 原来里面是括号\n        keypoints += 0.5\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(str(h5_path), '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]['matches0'].__array__()\n                idx = np.where(matches != -1)[0]\n                matches = np.stack([idx, matches[idx]], -1)\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                       features,\n                       matches,\n                       database_path = 'colmap.db',\n                       img_ext='.jpg',\n                       ):\n    db = COLMAPDatabase.connect(database_path)\n    db.create_tables()\n    single_camera = False\n    fname_to_id = add_keypoints(db, features, img_dir, img_ext, 'simple-radial', single_camera)\n    add_matches(\n        db,\n        matches,\n        fname_to_id,\n    )\n\n    db.commit()\n    return","metadata":{"execution":{"iopub.status.busy":"2023-06-03T02:12:22.371407Z","iopub.execute_input":"2023-06-03T02:12:22.371784Z","iopub.status.idle":"2023-06-03T02:12:22.398543Z","shell.execute_reply.started":"2023-06-03T02:12:22.371744Z","shell.execute_reply":"2023-06-03T02:12:22.397312Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"src = '/kaggle/input/image-matching-challenge-2023'","metadata":{"execution":{"iopub.status.busy":"2023-06-03T02:12:22.400077Z","iopub.execute_input":"2023-06-03T02:12:22.400523Z","iopub.status.idle":"2023-06-03T02:12:22.410455Z","shell.execute_reply.started":"2023-06-03T02:12:22.400471Z","shell.execute_reply":"2023-06-03T02:12:22.409408Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Get data from csv.\n\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            #dataset, scene, image, _, _ = l.strip().split(',')\n            image, dataset, scene, _, _ = l.strip().split(',')\n            if dataset not in data_dict.keys():\n                data_dict[dataset] = {}\n            if scene not in data_dict[dataset].keys():\n                data_dict[dataset][scene] = []\n            data_dict[dataset][scene].append(image)","metadata":{"execution":{"iopub.status.busy":"2023-06-03T02:12:22.412926Z","iopub.execute_input":"2023-06-03T02:12:22.414301Z","iopub.status.idle":"2023-06-03T02:12:22.424802Z","shell.execute_reply.started":"2023-06-03T02:12:22.414261Z","shell.execute_reply":"2023-06-03T02:12:22.423788Z"},"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":"2023-06-03T02:12:22.426224Z","iopub.execute_input":"2023-06-03T02:12:22.427183Z","iopub.status.idle":"2023-06-03T02:12:22.433949Z","shell.execute_reply.started":"2023-06-03T02:12:22.427121Z","shell.execute_reply":"2023-06-03T02:12:22.432862Z"},"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":"2023-06-03T02:12:22.435306Z","iopub.execute_input":"2023-06-03T02:12:22.436154Z","iopub.status.idle":"2023-06-03T02:12:22.443371Z","shell.execute_reply.started":"2023-06-03T02:12:22.436107Z","shell.execute_reply":"2023-06-03T02:12:22.442281Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# 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":"2023-06-03T02:12:22.444672Z","iopub.execute_input":"2023-06-03T02:12:22.445532Z","iopub.status.idle":"2023-06-03T02:12:22.456848Z","shell.execute_reply.started":"2023-06-03T02:12:22.445484Z","shell.execute_reply":"2023-06-03T02:12:22.455746Z"},"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}/{scene}/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}/test/{x}' for x in data_dict[dataset][scene]]\n            print (f\"Got {len(img_fnames)} images\")\n            feature_dir = f'featureout/{dataset}_{scene}'\n            if not os.path.isdir(feature_dir):\n                os.makedirs(feature_dir, exist_ok=True)\n            t=time()\n            index_pairs = get_image_pairs_shortlist(img_fnames,\n                                  sim_th = thr, # should be strict\n                                  min_pairs = pairs, # we select at least min_pairs PER IMAGE with biggest similarity\n                                  exhaustive_if_less = pairs,\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            t=time()\n            if not matcher:\n                features = extract_features.main(feature_conf, Path(img_dir), Path(feature_dir))\n                gc.collect()\n                t=time() -t \n                timings['feature_detection'].append(t)\n                print(f'Features detected in  {t:.4f} sec')\n                t=time()\n                matches = match_features.main(matcher_conf, index_pairs, feature_conf['output'], feature_dir)\n            else:\n                features, matches = match_dense.main(dense_conf, index_pairs, Path(img_dir), Path(feature_dir), \n                                                     max_kps=max_keypoints, overwrite=False)\n            t=time() -t \n            timings['feature_matching'].append(t)\n            print(f'Features matched in  {t:.4f} sec')\n            database_path = f'{feature_dir}/colmap.db'\n            if os.path.isfile(database_path):\n                os.remove(database_path)\n            gc.collect()\n\n            import_into_colmap(img_dir, features, matches, database_path)\n            output_path = f'{feature_dir}/colmap_rec'\n\n            t=time()\n            pycolmap.match_exhaustive(database_path)\n            t=time() - t \n            timings['RANSAC'].append(t)\n            print(f'RANSAC in  {t:.4f} sec')\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 # 3\n            mapper_options.ba_local_max_num_iterations = 25 # 25\n            mapper_options.ba_global_max_num_iterations = 50 # 50\n            mapper_options.ba_local_max_refinements = 2 # 2\n            mapper_options.ba_global_max_refinements = 5 # 5\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            imgs_registered  = 0\n            best_idx = None\n            print (\"Looking for the best reconstruction\")\n            if isinstance(maps, dict):\n                for idx1, rec in maps.items():\n                    print (idx1, rec.summary())\n                    if len(rec.images) > imgs_registered:\n                        imgs_registered = len(rec.images)\n                        best_idx = idx1\n            if best_idx is not None:\n                print (maps[best_idx].summary())\n                for k, im in maps[best_idx].images.items():\n                    key1 = f'{dataset}/{scene}/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            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":{"execution":{"iopub.status.busy":"2023-06-03T02:12:22.458732Z","iopub.execute_input":"2023-06-03T02:12:22.459501Z"},"collapsed":true,"jupyter":{"outputs_hidden":true},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"create_submission(out_results, data_dict)","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}