{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":49349,"databundleVersionId":5447706,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":5373920,"sourceType":"datasetVersion","datasetId":3117886},{"sourceId":5850511,"sourceType":"datasetVersion","datasetId":3364321},{"sourceId":12005133,"sourceType":"datasetVersion","datasetId":7552209},{"sourceId":12010068,"sourceType":"datasetVersion","datasetId":7555674},{"sourceId":12010214,"sourceType":"datasetVersion","datasetId":7555876},{"sourceId":12010218,"sourceType":"datasetVersion","datasetId":7555879},{"sourceId":12025632,"sourceType":"datasetVersion","datasetId":7552874},{"sourceId":12025645,"sourceType":"datasetVersion","datasetId":7566072},{"sourceId":12027050,"sourceType":"datasetVersion","datasetId":7567063},{"sourceId":12027277,"sourceType":"datasetVersion","datasetId":7567218},{"sourceId":176463227,"sourceType":"kernelVersion"},{"sourceId":4534,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":3326,"modelId":986},{"sourceId":418016,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":340991,"modelId":362238}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"https://www.kaggle.com/code/eduardtrulls/imc-2023-submission-example","metadata":{}},{"cell_type":"code","source":"!pip install -q /kaggle/input/einops-whl/einops-0.8.1-py3-none-any.whl --no-index --find-links /kaggle/input/einops-v0-8-0\n!pip install /kaggle/input/imc2023-vggt-whl/* --no-deps --no-index --find-links /kaggle/input/imc2023-vggt-whl\n!pip install /kaggle/input/roma-whl/roma-1.5.2.1-py3-none-any.whl --no-index --find-links /kaggle/input/roma-1.5.2.1\n# !python -m pip install --no-index --find-links=/kaggle/input/pkg-check-orientation/ check_orientation==0.0.5 > /dev/null","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-02T05:23:57.477063Z","iopub.execute_input":"2025-06-02T05:23:57.477944Z","iopub.status.idle":"2025-06-02T05:24:05.428661Z","shell.execute_reply.started":"2025-06-02T05:23:57.477912Z","shell.execute_reply":"2025-06-02T05:24:05.427863Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# General utilities\nimport os\nfrom tqdm import tqdm\nfrom time import time\nfrom fastprogress import progress_bar\nimport gc\nimport math\nimport numpy as np\nfrom IPython.display import clear_output\nfrom collections import defaultdict\nfrom copy import deepcopy\nimport matplotlib.pyplot as plt\nimport concurrent.futures\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\nfrom transformers import AutoImageProcessor, AutoModel\n\n# 3D reconstruction\nimport pycolmap\n\nprint(\"Kornia version\", K.__version__)\nprint(\"Pycolmap version\", pycolmap.__version__)","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","execution":{"iopub.status.busy":"2025-06-02T05:24:05.430612Z","iopub.execute_input":"2025-06-02T05:24:05.430863Z","iopub.status.idle":"2025-06-02T05:24:05.437979Z","shell.execute_reply.started":"2025-06-02T05:24:05.430841Z","shell.execute_reply":"2025-06-02T05:24:05.437331Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import sys\nfrom pathlib import Path\nimport torch\n\n# 這一層目錄裡面應該直接包含  mast3r/  資料夾\nMAST3R_ROOT = Path(\"/kaggle/input/mast3r-code/mast3r\")\n\n# 如果還沒加過，就加進搜尋路徑\nif str(MAST3R_ROOT) not in sys.path:\n    sys.path.append(str(MAST3R_ROOT))\n\n# 現在就可以正常 import 了\n\nfrom mast3r.cloud_opt.sparse_ga import sparse_global_alignment\nfrom mast3r.model import AsymmetricMASt3R\nfrom mast3r.image_pairs import make_pairs\nfrom dust3r.utils.image import load_images\nfrom dust3r.utils.device import to_numpy\nfrom mast3r.model import load_model\nfrom mast3r.fast_nn import extract_correspondences_nonsym, bruteforce_reciprocal_nns\nfrom dust3r.inference import inference  # Import here to avoid circular imports\nfrom dust3r.utils.geometry import find_reciprocal_matches, xy_grid, geotrf  # noqa","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-02T05:24:05.438845Z","iopub.execute_input":"2025-06-02T05:24:05.439089Z","iopub.status.idle":"2025-06-02T05:24:05.459611Z","shell.execute_reply.started":"2025-06-02T05:24:05.439067Z","shell.execute_reply":"2025-06-02T05:24:05.458879Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Global Configs","metadata":{}},{"cell_type":"code","source":"# Mode can only be train or test. This will be used to find the image directory.\n# Use \"test\" for submission \nMODE = \"train\"\nMODE = \"test\"\n\n# Option to change path for local testing\nis_local = True\nis_local = False\n\nif is_local:\n    NUM_CORES = 2\n    SRC = \"./kaggle/input/image-matching-challenge-2023\"\n    MODEL_DIR = \"./kaggle/input/kornia-local-feature-weights/\"\n    DISK_PATH = \"./loftr_disk.ckpt\"\n    HARDNET_PT = \"./kaggle/input/kornia-local-feature-weights/hardnet8v2.pt\"\nelse:\n    NUM_CORES = 2\n    SRC = \"/kaggle/input/image-matching-challenge-2023\"\n    MODEL_DIR = \"/kaggle/input/kornia-local-feature-weights/\"\n    DISK_PATH = \"/kaggle/input/disk/pytorch/depth-supervision/1/loftr_outdoor.ckpt\"\n    HARDNET_PT = \"/kaggle/input/hardnet8v2/hardnet8v2.pt\"\n    MAST3R_PATH = \"/kaggle/input/mast3r/pytorch/default/1\"\n\nLOG_MESSAGE = \"Final submission\"\nMATCHES_CAP = None\n\nDEBUG = True\nDEBUG = False\n\n# DEBUG_SCENE = [\"cyprus\", \"kyiv-puppet-theater\"]\n# DEBUG_SCENE = [\"cyprus\"]\nDEBUG_SCENE = [\"kyiv-puppet-theater\"]\n# DEBUG_SCENE = [\"kyiv-puppet-theater\", \"cyprus\", \"wall\", \"chairs\"]\n# DEBUG_SCENE = [\"bike\"]\n# DEBUG_SCENE = [\"wall\"]\n\n# Longer edge limit of the input image\nhardnet_res = 1600\n\nMODEL_DICT = {\n    \"Keynet\": {\"enable\": True, \"resize_long_edge_to\": hardnet_res, \"pair_only\": False},\n    \"GFTT\": {\"enable\": True, \"resize_long_edge_to\": hardnet_res},\n    \"DoG\": {\"enable\": True, \"resize_long_edge_to\": hardnet_res},\n    \"Harris\": {\"enable\": True, \"resize_long_edge_to\": hardnet_res},\n    \"MASt3R\": {\"enable\": True, \"resize_long_edge_to\": 512}  # MASt3R uses 512x512 input\n}\n\n# Find fundamental matrix parameters\nFM_PARAMS = {\"ransacReprojThreshold\": 5, \"confidence\": 0.9999, \"maxIters\": 50000, \"removeOutliers\": True}\n\n# Remove a \"match\" if the number of matches is lower than MATCH_FILTER_RATIO*max_num_matches\n# e.g. img1 and img2 have max 10000 matches with some other images, img2 and img1 only have 99 matches. The matches btw img1 and img2 won't be selected.\nMATCH_FILTER_RATIO = 0.01\n\n# for logging\nLOG_DICT = dict()\nLOG_DICT[\"mode\"] = MODE\nLOG_DICT[\"log_message\"] = LOG_MESSAGE\nLOG_DICT[\"matches_cap\"] = MATCHES_CAP\nLOG_DICT[\"debug\"] = DEBUG\nLOG_DICT[\"debug_scene\"] = DEBUG_SCENE\n\nif MODE == \"test\":\n    DEBUG = False\ndevice = torch.device(\"cuda\")\nprint(torch.cuda.is_available())\n","metadata":{"execution":{"iopub.status.busy":"2025-06-02T05:24:05.461385Z","iopub.execute_input":"2025-06-02T05:24:05.461640Z","iopub.status.idle":"2025-06-02T05:24:05.477276Z","shell.execute_reply.started":"2025-06-02T05:24:05.461624Z","shell.execute_reply":"2025-06-02T05:24:05.476708Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Get datadict from submission file","metadata":{}},{"cell_type":"code","source":"# Get datadict from csv.\nif MODE == \"train\":\n    sample_path = f\"{SRC}/train/train_labels.csv\"\nelse:\n    sample_path = f\"{SRC}/sample_submission.csv\"\n\ndata_dict = {}\nwith open(sample_path, \"r\") as f:\n    for i, l in enumerate(f):\n        # Skip header.\n        if l and i > 0:\n            if MODE == \"train\":\n                dataset, scene, image, _, _ = l.strip().split(\",\")\n            else:\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            \nall_scenes = []\nscene_len = []\nfor dataset in data_dict:\n    for scene in data_dict[dataset]:\n        print(f\"{dataset} / {scene} -> {len(data_dict[dataset][scene])} images\")\n        if DEBUG and (scene not in DEBUG_SCENE):\n            continue\n        all_scenes.append((dataset, scene))\n        scene_len.append(len(data_dict[dataset][scene]))\n\n# sort all scenes by length, lowest first\nall_scenes = [x for _, x in sorted(zip(scene_len, all_scenes), reverse=True)]\n\n# Print reconst order\nprint(\"Reconstruction order: \")\nfor scene in all_scenes:\n    print(f\" --{scene[0]} / {scene[1]}\")","metadata":{"execution":{"iopub.status.busy":"2025-06-02T05:24:05.477989Z","iopub.execute_input":"2025-06-02T05:24:05.478194Z","iopub.status.idle":"2025-06-02T05:24:05.500359Z","shell.execute_reply.started":"2025-06-02T05:24:05.478178Z","shell.execute_reply":"2025-06-02T05:24:05.499683Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Submission Utils","metadata":{}},{"cell_type":"code","source":"def arr_to_str(a):\n    return \";\".join([str(x) for x in a.reshape(-1)])\n\n\n# Function to create a submission file.\ndef create_submission(out_results, data_dict, mode=\"test\"):\n    if mode == \"train\":\n        file_name = \"submission_train.csv\"\n    else:\n        file_name = \"submission.csv\"\n\n    with open(file_name, \"w\") as f:\n        f.write(\"image_path,dataset,scene,rotation_matrix,translation_vector\\n\")\n        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(\n                        f\"{image},{dataset},{scene},{arr_to_str(R)},{arr_to_str(T)}\\n\"\n                    )","metadata":{"execution":{"iopub.status.busy":"2025-06-02T05:24:05.501358Z","iopub.execute_input":"2025-06-02T05:24:05.501646Z","iopub.status.idle":"2025-06-02T05:24:05.515951Z","shell.execute_reply.started":"2025-06-02T05:24:05.501624Z","shell.execute_reply":"2025-06-02T05:24:05.515256Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Image Loading and Resize","metadata":{}},{"cell_type":"code","source":"def load_torch_image(fname, device=torch.device(\"cpu\")):\n    img = K.image_to_tensor(cv2.imread(fname), False).float() / 255.0\n    img = K.color.bgr_to_rgb(img.to(device))\n    return img\n\n\ndef resize_torch_image(\n    timg, resize_long_edge_to=None, align=None, disable_enlarge=True\n):\n    h, w = timg.shape[2:]\n    raw_size = torch.tensor(timg.shape[2:])\n    if resize_long_edge_to is None:\n        scale = 1\n    else:\n        scale = float(resize_long_edge_to) / float(max(raw_size[0], raw_size[1]))\n\n    if disable_enlarge:\n        scale = min(scale, 1)\n\n    h_resized = int(h * scale)\n    w_resized = int(w * scale)\n\n    if align is not None:\n        assert align > 0\n        h_resized = h_resized - h_resized % align\n        w_resized = w_resized - w_resized % align\n    scale_h = h_resized / h\n    scale_w = w_resized / w\n\n    timg_resized = K.geometry.resize(timg, (h_resized, w_resized), antialias = True)\n    return timg_resized, scale_h, scale_w\n\n\ndef get_roi_image(timg, roi):\n    min_h = int(roi[\"roi_min_h\"])\n    min_w = int(roi[\"roi_min_w\"])\n    max_h = int(roi[\"roi_max_h\"])\n    max_w = int(roi[\"roi_max_w\"])\n    roi_img = timg[:, :, min_h:max_h, min_w:max_w]\n    roi_w_scale = (max_w - min_w) / timg.shape[3]\n    roi_h_scale = (max_h - min_h) / timg.shape[2]\n    return roi_img, min_h, min_w","metadata":{"execution":{"iopub.status.busy":"2025-06-02T05:24:05.516732Z","iopub.execute_input":"2025-06-02T05:24:05.516981Z","iopub.status.idle":"2025-06-02T05:24:05.535479Z","shell.execute_reply.started":"2025-06-02T05:24:05.516958Z","shell.execute_reply":"2025-06-02T05:24:05.534754Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Rotation detection","metadata":{}},{"cell_type":"code","source":"# from torchvision.io import read_image as T_read_image\n# from torchvision.io import ImageReadMode\n# from torchvision import transforms as T\n# from check_orientation.pre_trained_models import create_model\n\n# def 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\n# class 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\n# def 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\n# def 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":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-02T05:24:05.536353Z","iopub.execute_input":"2025-06-02T05:24:05.536571Z","iopub.status.idle":"2025-06-02T05:24:05.552565Z","shell.execute_reply.started":"2025-06-02T05:24:05.536556Z","shell.execute_reply":"2025-06-02T05:24:05.551877Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Visualization Utils","metadata":{}},{"cell_type":"code","source":"# Visualzation block\ndef draw_keypoints(img, keypoints, color=(0, 255, 0)):\n    max_edge = max(img.shape[0], img.shape[1])\n    good_radius = 4\n    for kp in keypoints:\n        x, y = kp\n        cv2.circle(img, (int(x), int(y)), color=color, radius=good_radius, thickness=-1)\n\n\ndef draw_roi(img, roi, color=(0, 255, 255)):\n    x1, y1, x2, y2 = (\n        roi[\"roi_min_w\"],\n        roi[\"roi_min_h\"],\n        roi[\"roi_max_w\"],\n        roi[\"roi_max_h\"],\n    )\n    cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), color=color, thickness=2)\n\n\ndef plot_images_with_keypoints(fname1, fname2, kpts1, kpts2, matches, rois=None):\n    print(fname1, fname2)\n    # Draw keypoints on the images\n    image1 = cv2.imread(fname1)\n    image2 = cv2.imread(fname2)\n    print(image1.shape, image2.shape)\n\n    # draw_keypoints(image1, kpts1)\n    # draw_keypoints(image2, kpts2)\n    if rois is not None:\n        draw_roi(image1, rois[0])\n        draw_roi(image2, rois[1])\n    print(image1.shape, image2.shape)\n    print(\"Number of matches:\", len(matches))\n    print(\"Number of keypoints:\", len(kpts1), len(kpts2))\n    #print the first match\n    print(matches[0])\n    # Resize image1 and image2 to have the same smaller height\n    display_h = 840\n    h1, w1 = image1.shape[:2]\n    h2, w2 = image2.shape[:2]\n    # new_h = min(h1, h2)\n    scale1 = display_h / h1\n    scale2 = display_h / h2\n    new_w1 = int(w1 * scale1)\n    new_w2 = int(w2 * scale2)\n\n    image1 = cv2.resize(image1, (new_w1, display_h))\n    image2 = cv2.resize(image2, (new_w2, display_h))\n\n    # Create a new image by horizontally concatenating the two images\n    concatenated_img = cv2.hconcat([image1, image2])\n\n    # Draw lines between the matching keypoints\n    for match in matches:\n        img1_idx = match[0]\n        img2_idx = match[1]\n        (x1, y1) = kpts1[img1_idx] * scale1\n        (x2, y2) = kpts2[img2_idx] * scale2\n        pt1 = (int(x1), int(y1))\n        pt2 = (int(x2) + image1.shape[1], int(y2))\n        cv2.line(concatenated_img, pt1, pt2, (0, 0, 255), 2)\n\n    # Plot the concatenated image\n    plt.figure(figsize=(20, 12))\n    plt.imshow(cv2.cvtColor(concatenated_img, cv2.COLOR_BGR2RGB))\n    plt.axis(\"off\")\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2025-06-02T05:24:05.553319Z","iopub.execute_input":"2025-06-02T05:24:05.553594Z","iopub.status.idle":"2025-06-02T05:24:05.572998Z","shell.execute_reply.started":"2025-06-02T05:24:05.553568Z","shell.execute_reply":"2025-06-02T05:24:05.572307Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Colmap database","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(\n    MAX_IMAGE_ID\n)\n\nCREATE_TWO_VIEW_GEOMETRIES_TABLE = \"\"\"\nCREATE TABLE IF NOT EXISTS two_view_geometries (\n    pair_id INTEGER PRIMARY KEY NOT NULL,\n    rows INTEGER NOT NULL,\n    cols INTEGER NOT NULL,\n    data BLOB,\n    config INTEGER NOT NULL,\n    F BLOB,\n    E BLOB,\n    H BLOB)\n\"\"\"\n\nCREATE_KEYPOINTS_TABLE = \"\"\"CREATE TABLE IF NOT EXISTS keypoints (\n    image_id INTEGER PRIMARY KEY NOT NULL,\n    rows INTEGER NOT NULL,\n    cols INTEGER NOT NULL,\n    data BLOB,\n    FOREIGN KEY(image_id) REFERENCES images(image_id) ON DELETE CASCADE)\n\"\"\"\n\nCREATE_MATCHES_TABLE = \"\"\"CREATE TABLE IF NOT EXISTS matches (\n    pair_id INTEGER PRIMARY KEY NOT NULL,\n    rows INTEGER NOT NULL,\n    cols INTEGER NOT NULL,\n    data BLOB)\"\"\"\n\nCREATE_NAME_INDEX = \"CREATE UNIQUE INDEX IF NOT EXISTS index_name ON images(name)\"\n\nCREATE_ALL = \"; \".join(\n    [\n        CREATE_CAMERAS_TABLE,\n        CREATE_IMAGES_TABLE,\n        CREATE_KEYPOINTS_TABLE,\n        CREATE_DESCRIPTORS_TABLE,\n        CREATE_MATCHES_TABLE,\n        CREATE_TWO_VIEW_GEOMETRIES_TABLE,\n        CREATE_NAME_INDEX,\n    ]\n)\n\n\ndef image_ids_to_pair_id(image_id1, image_id2):\n    if image_id1 > image_id2:\n        image_id1, image_id2 = image_id2, image_id1\n    return image_id1 * MAX_IMAGE_ID + image_id2\n\n\ndef pair_id_to_image_ids(pair_id):\n    image_id2 = pair_id % MAX_IMAGE_ID\n    image_id1 = (pair_id - image_id2) / MAX_IMAGE_ID\n    return image_id1, image_id2\n\n\ndef array_to_blob(array):\n    if IS_PYTHON3:\n        return array.tostring()\n    else:\n        return np.getbuffer(array)\n\n\ndef blob_to_array(blob, dtype, shape=(-1,)):\n    if IS_PYTHON3:\n        return np.fromstring(blob, dtype=dtype).reshape(*shape)\n    else:\n        return np.frombuffer(blob, dtype=dtype).reshape(*shape)\n\n\nclass COLMAPDatabase(sqlite3.Connection):\n    @staticmethod\n    def connect(database_path):\n        return sqlite3.connect(database_path, factory=COLMAPDatabase)\n\n    def __init__(self, *args, **kwargs):\n        super(COLMAPDatabase, self).__init__(*args, **kwargs)\n\n        self.create_tables = lambda: self.executescript(CREATE_ALL)\n        self.create_cameras_table = lambda: self.executescript(CREATE_CAMERAS_TABLE)\n        self.create_descriptors_table = lambda: self.executescript(\n            CREATE_DESCRIPTORS_TABLE\n        )\n        self.create_images_table = lambda: self.executescript(CREATE_IMAGES_TABLE)\n        self.create_two_view_geometries_table = lambda: self.executescript(\n            CREATE_TWO_VIEW_GEOMETRIES_TABLE\n        )\n        self.create_keypoints_table = lambda: self.executescript(CREATE_KEYPOINTS_TABLE)\n        self.create_matches_table = lambda: self.executescript(CREATE_MATCHES_TABLE)\n        self.create_name_index = lambda: self.executescript(CREATE_NAME_INDEX)\n\n    def add_camera(\n        self, model, width, height, params, prior_focal_length=False, camera_id=None\n    ):\n        params = np.asarray(params, np.float64)\n        cursor = self.execute(\n            \"INSERT INTO cameras VALUES (?, ?, ?, ?, ?, ?)\",\n            (\n                camera_id,\n                model,\n                width,\n                height,\n                array_to_blob(params),\n                prior_focal_length,\n            ),\n        )\n        return cursor.lastrowid\n\n    def add_image(\n        self, name, camera_id, prior_q=np.zeros(4), prior_t=np.zeros(3), image_id=None\n    ):\n        cursor = self.execute(\n            \"INSERT INTO images VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)\",\n            (\n                image_id,\n                name,\n                camera_id,\n                prior_q[0],\n                prior_q[1],\n                prior_q[2],\n                prior_q[3],\n                prior_t[0],\n                prior_t[1],\n                prior_t[2],\n            ),\n        )\n        return cursor.lastrowid\n\n    def add_keypoints(self, image_id, keypoints):\n        assert len(keypoints.shape) == 2\n        assert keypoints.shape[1] in [2, 4, 6]\n\n        keypoints = np.asarray(keypoints, np.float32)\n        self.execute(\n            \"INSERT INTO keypoints VALUES (?, ?, ?, ?)\",\n            (image_id,) + keypoints.shape + (array_to_blob(keypoints),),\n        )\n\n    def add_descriptors(self, image_id, descriptors):\n        descriptors = np.ascontiguousarray(descriptors, np.uint8)\n        self.execute(\n            \"INSERT INTO descriptors VALUES (?, ?, ?, ?)\",\n            (image_id,) + descriptors.shape + (array_to_blob(descriptors),),\n        )\n\n    def add_matches(self, image_id1, image_id2, matches):\n        assert len(matches.shape) == 2\n        assert matches.shape[1] == 2\n\n        if image_id1 > image_id2:\n            matches = matches[:, ::-1]\n\n        pair_id = image_ids_to_pair_id(image_id1, image_id2)\n        matches = np.asarray(matches, np.uint32)\n        self.execute(\n            \"INSERT INTO matches VALUES (?, ?, ?, ?)\",\n            (pair_id,) + matches.shape + (array_to_blob(matches),),\n        )\n\n    def add_two_view_geometry(\n        self,\n        image_id1,\n        image_id2,\n        matches,\n        F=np.eye(3),\n        E=np.eye(3),\n        H=np.eye(3),\n        config=2,\n    ):\n        assert len(matches.shape) == 2\n        assert matches.shape[1] == 2\n\n        if image_id1 > image_id2:\n            matches = matches[:, ::-1]\n\n        pair_id = image_ids_to_pair_id(image_id1, image_id2)\n        matches = np.asarray(matches, np.uint32)\n        F = np.asarray(F, dtype=np.float64)\n        E = np.asarray(E, dtype=np.float64)\n        H = np.asarray(H, dtype=np.float64)\n        self.execute(\n            \"INSERT INTO two_view_geometries VALUES (?, ?, ?, ?, ?, ?, ?, ?)\",\n            (pair_id,)\n            + matches.shape\n            + (\n                array_to_blob(matches),\n                config,\n                array_to_blob(F),\n                array_to_blob(E),\n                array_to_blob(H),\n            ),\n        )","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2025-06-02T05:24:05.575509Z","iopub.execute_input":"2025-06-02T05:24:05.575726Z","iopub.status.idle":"2025-06-02T05:24:05.595708Z","shell.execute_reply.started":"2025-06-02T05:24:05.575710Z","shell.execute_reply":"2025-06-02T05:24:05.595066Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# DB operation","metadata":{}},{"cell_type":"code","source":"# Modified from https://github.com/cvlab-epfl/disk/blob/37f1f7e971cea3055bb5ccfc4cf28bfd643fa339/colmap/h5_to_db.py\n\n#  Copyright [2020] [Michał Tyszkiewicz, Pascal Fua, Eduard Trulls]\n#\n#   Licensed under the Apache License, Version 2.0 (the \"License\");\n#   you may not use this file except in compliance with the License.\n#   You may obtain a copy of the License at\n#\n#       http://www.apache.org/licenses/LICENSE-2.0\n#\n#   Unless required by applicable law or agreed to in writing, software\n#   distributed under the License is distributed on an \"AS IS\" BASIS,\n#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n#   See the License for the specific language governing permissions and\n#   limitations under the License.\n\nimport os, argparse, h5py, warnings\nimport numpy as np\nfrom tqdm import tqdm\nfrom PIL import Image, ExifTags\n\n\ndef get_focal(image_path, err_on_default=False):\n    image = Image.open(image_path)\n    max_size = max(image.size)\n\n    exif = image.getexif()\n    \n    #\n    # Modified to add exif_ifd to exif dict\n    #\n    exif_ifd = exif.get_ifd(0x8769)\n    exif.update(exif_ifd)\n\n    focal = None\n    is_from_exif = False\n    if exif is not None:\n        focal_35mm = None\n        # https://github.com/colmap/colmap/blob/d3a29e203ab69e91eda938d6e56e1c7339d62a99/src/util/bitmap.cc#L299\n        for tag, value in exif.items():\n            focal_35mm = None\n            if ExifTags.TAGS.get(tag, None) == \"FocalLengthIn35mmFilm\":\n                focal_35mm = float(value)\n                is_from_exif = True\n                break\n\n        if focal_35mm is not None:\n            focal = focal_35mm / 35.0 * max_size\n\n    if focal is None:\n        if err_on_default:\n            raise RuntimeError(\"Failed to find focal length\")\n\n        # failed to find it in exif, use prior\n        FOCAL_PRIOR = 1.2\n        focal = FOCAL_PRIOR * max_size\n    \n    #\n    # Modified to return a bool indicating if the focal length is from exif \n    #\n    return focal, is_from_exif\n\n\ndef create_camera(db, image_path, camera_model):\n    image = Image.open(image_path)\n    width, height = image.size\n\n    focal, is_from_exif = get_focal(image_path)\n\n    if camera_model == \"simple-pinhole\":\n        model = 0  # simple pinhole\n        param_arr = np.array([focal, width / 2, height / 2])\n    if camera_model == \"pinhole\":\n        model = 1  # pinhole\n        param_arr = np.array([focal, focal, width / 2, height / 2])\n    elif camera_model == \"simple-radial\":\n        model = 2  # simple radial\n        param_arr = np.array([focal, width / 2, height / 2, 0.1])\n    elif camera_model == \"opencv\":\n        model = 4  # opencv\n        param_arr = np.array([focal, focal, width / 2, height / 2, 0.0, 0.0, 0.0, 0.0])\n\n    #\n    # Modified to set prior_focal_length if the focal length is from exif\n    #\n    return db.add_camera(\n        model, width, height, param_arr, prior_focal_length=is_from_exif\n    )\n\n\ndef add_kpts_matches(db, img_dir, kpts, matches, fms = None):\n    fname_to_id = {}\n\n    # Add keypoints\n    for filename in tqdm(kpts):\n        path = os.path.join(img_dir, filename)\n        camera_model = \"simple-radial\"\n        camera_id = create_camera(db, path, camera_model)\n        image_id = db.add_image(filename, camera_id)\n        fname_to_id[filename] = image_id\n        db.add_keypoints(image_id, kpts[filename])\n\n    n_keys = len(matches)\n    n_total = (n_keys * (n_keys - 1)) // 2\n    # Add matches\n    added = set()\n    with tqdm(total=n_total) as pbar:\n        for key1 in matches:\n            for key2 in matches[key1]:\n                id_1 = fname_to_id[key1]\n                id_2 = fname_to_id[key2]\n                pair_id = image_ids_to_pair_id(id_1, id_2)\n                if pair_id in added:\n                    warnings.warn(f\"Pair {pair_id} ({id_1}, {id_2}) already added!\")\n                    continue\n                \n                # Remove duplicate matches\n                matches_array = matches[key1][key2]\n                matches_tuples = set(map(tuple, matches_array))\n                unique_matches = np.array(list(matches_tuples))\n                \n                # Sort matches by first column to ensure consistent ordering\n                unique_matches = unique_matches[np.argsort(unique_matches[:, 0])]\n                \n                # Remove matches where the same point in image1 matches to multiple points in image2\n                _, unique_indices = np.unique(unique_matches[:, 0], return_index=True)\n                unique_matches = unique_matches[unique_indices]\n                \n                # Remove matches where the same point in image2 matches to multiple points in image1\n                _, unique_indices = np.unique(unique_matches[:, 1], return_index=True)\n                unique_matches = unique_matches[unique_indices]\n                \n                if len(unique_matches) >= 15:  # Keep only if we still have enough matches\n                    db.add_matches(id_1, id_2, unique_matches)\n                    added.add(pair_id)\n                    if fms is not None:\n                        db.add_two_view_geometry(id_1, id_2, unique_matches, fms[key1][key2], np.eye(3), np.eye(3))\n                pbar.update(1)\n\n    db.commit()\n","metadata":{"execution":{"iopub.status.busy":"2025-06-02T05:24:05.596379Z","iopub.execute_input":"2025-06-02T05:24:05.596588Z","iopub.status.idle":"2025-06-02T05:24:05.612309Z","shell.execute_reply.started":"2025-06-02T05:24:05.596574Z","shell.execute_reply":"2025-06-02T05:24:05.611687Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_unique_idxs(A, dim=0):\n    # https://stackoverflow.com/questions/72001505/how-to-get-unique-elements-and-their-firstly-appeared-indices-of-a-pytorch-tenso\n    unique, idx, counts = torch.unique(\n        A, dim=dim, sorted=True, return_inverse=True, return_counts=True\n    )\n    _, ind_sorted = torch.sort(idx, stable=True)\n    cum_sum = counts.cumsum(0)\n    cum_sum = torch.cat((torch.tensor([0], device=cum_sum.device), cum_sum[:-1]))\n    first_indices = ind_sorted[cum_sum]\n    return first_indices","metadata":{"execution":{"iopub.status.busy":"2025-06-02T05:24:05.613273Z","iopub.execute_input":"2025-06-02T05:24:05.613562Z","iopub.status.idle":"2025-06-02T05:24:05.630796Z","shell.execute_reply.started":"2025-06-02T05:24:05.613544Z","shell.execute_reply":"2025-06-02T05:24:05.630190Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# AffNetHardNet Models","metadata":{}},{"cell_type":"code","source":"# Making kornia local features loading w/o internet\nclass AffNetHardNet(KF.LocalFeature):\n    \"\"\"Convenience module, which implements KeyNet detector + AffNet + HardNet descriptor.\n\n    .. image:: _static/img/keynet_affnet.jpg\n    \"\"\"\n\n    def __init__(\n        self,\n        num_features: int = 5000,\n        upright: bool = False,\n        device=torch.device(\"cpu\"),\n        scale_laf: float = 1.0,\n        detector = \"keynet\"\n    ):\n        detector_options = [\"keynet\", \"GFTT\", \"Hessian\", \"Harris\", \"DoG\"]\n        if detector not in detector_options:\n            raise ValueError(\"Detector must be one of {}\".format(detector_options))\n        \n        ori_module = (\n            KF.PassLAF()\n            if upright\n            else KF.LAFOrienter(angle_detector=KF.OriNet(False)).eval()\n        )\n        if not upright:\n            weights = torch.load(os.path.join(MODEL_DIR, \"OriNet.pth\"))[\"state_dict\"]\n            ori_module.angle_detector.load_state_dict(weights)\n\n        config = {\n            # Extraction Parameters\n            \"nms_size\": 15,\n            \"pyramid_levels\": 4,\n            \"up_levels\": 1,\n            \"scale_factor_levels\": math.sqrt(2),\n            \"s_mult\": 22.0,\n        }\n\n        if detector == \"keynet\":\n            detector = KF.KeyNetDetector(\n            False,\n            num_features=num_features,\n            ori_module=ori_module,\n            aff_module=KF.LAFAffNetShapeEstimator(False).eval(),\n            ).to(device)\n            kn_weights = torch.load(os.path.join(MODEL_DIR, \"keynet_pytorch.pth\"))[\n            \"state_dict\"\n            ]\n            detector.model.load_state_dict(kn_weights)\n        elif detector == \"GFTT\":\n            detector = KF.MultiResolutionDetector(\n                KF.CornerGFTT(),\n                num_features=num_features,\n                config=config,\n                ori_module=ori_module,\n                aff_module=KF.LAFAffNetShapeEstimator(False).eval(),\n            ).to(device)\n        elif detector == \"Harris\":\n            detector = KF.MultiResolutionDetector(\n                KF.CornerHarris(0.04),\n                num_features=num_features,\n                config=config,\n                ori_module=ori_module,\n                aff_module=KF.LAFAffNetShapeEstimator(False).eval(),\n            ).to(device)\n        elif detector == \"DoG\":\n            detector = KF.MultiResolutionDetector(\n                KF.BlobDoGSingle(),\n                num_features=num_features,\n                config=config,\n                ori_module=ori_module,\n                aff_module=KF.LAFAffNetShapeEstimator(False).eval(),\n            ).to(device)\n        affnet_weights = torch.load(os.path.join(MODEL_DIR, \"AffNet.pth\"))[\"state_dict\"]\n        detector.aff.load_state_dict(affnet_weights)\n\n        # hardnet = KF.HardNet(False).eval()\n        # hn_weights = torch.load(os.path.join(MODEL_DIR, \"HardNetLib.pth\"))[\"state_dict\"]\n        # hardnet.load_state_dict(hn_weights)\n        # descriptor = KF.LAFDescriptor(\n        #     hardnet, patch_size=32, grayscale_descriptor=True\n        # ).to(device)\n        hardnet8 = KF.HardNet8(False).eval()\n        hn8_weights = torch.load(HARDNET_PT)\n        hardnet8.load_state_dict(hn8_weights)\n        descriptor = KF.LAFDescriptor(\n            hardnet8, patch_size=32, grayscale_descriptor=True\n        ).to(device)\n        super().__init__(detector, descriptor, scale_laf)\n","metadata":{"execution":{"iopub.status.busy":"2025-06-02T05:24:05.631584Z","iopub.execute_input":"2025-06-02T05:24:05.631811Z","iopub.status.idle":"2025-06-02T05:24:05.646685Z","shell.execute_reply.started":"2025-06-02T05:24:05.631795Z","shell.execute_reply":"2025-06-02T05:24:05.645959Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_unique_matches(f_match_kpts):\n    kpts = defaultdict(list)\n    match_indexes = defaultdict(dict)\n    total_kpts = defaultdict(int)\n    for key1 in f_match_kpts:\n        for key2 in f_match_kpts[key1]:\n            matches = f_match_kpts[key1][key2]\n            kpts[key1].append(matches[:, :2])\n            kpts[key2].append(matches[:, 2:])\n            current_match = torch.arange(len(matches)).reshape(-1, 1).repeat(1, 2)\n            current_match[:, 0] += total_kpts[key1]\n            current_match[:, 1] += total_kpts[key2]\n            total_kpts[key1] += len(matches)\n            total_kpts[key2] += len(matches)\n            match_indexes[key1][key2] = current_match\n\n    for key in kpts:\n        kpts[key] = np.round(np.concatenate(kpts[key], axis=0))\n\n    unique_kpts = {}\n    unique_match_idxs = {}\n    out_match = defaultdict(dict)\n\n    for key in kpts.keys():\n        uniq_kps, uniq_reverse_idxs = torch.unique(\n            torch.from_numpy(kpts[key]), dim=0, return_inverse=True\n        )\n        unique_match_idxs[key] = uniq_reverse_idxs\n        unique_kpts[key] = uniq_kps.numpy()\n\n    for key1 in match_indexes:\n        for key2 in match_indexes[key1]:\n            m2 = deepcopy(match_indexes[key1][key2])\n            m2[:, 0] = unique_match_idxs[key1][m2[:, 0]]\n            m2[:, 1] = unique_match_idxs[key2][m2[:, 1]]\n            mkpts = np.concatenate(\n                [\n                    unique_kpts[key1][m2[:, 0]],\n                    unique_kpts[key2][m2[:, 1]],\n                ],\n                axis=1,\n            )\n            unique_idxs_current = get_unique_idxs(torch.from_numpy(mkpts), dim=0)\n            m2_semiclean = m2[unique_idxs_current]\n            unique_idxs_current1 = get_unique_idxs(m2_semiclean[:, 0], dim=0)\n            m2_semiclean = m2_semiclean[unique_idxs_current1]\n            unique_idxs_current2 = get_unique_idxs(m2_semiclean[:, 1], dim=0)\n            m2_semiclean2 = m2_semiclean[unique_idxs_current2]\n            out_match[key1][key2] = m2_semiclean2.numpy()\n    return unique_kpts, out_match","metadata":{"execution":{"iopub.status.busy":"2025-06-02T05:24:05.647622Z","iopub.execute_input":"2025-06-02T05:24:05.647884Z","iopub.status.idle":"2025-06-02T05:24:05.667649Z","shell.execute_reply.started":"2025-06-02T05:24:05.647861Z","shell.execute_reply":"2025-06-02T05:24:05.666827Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Scene feature detector","metadata":{}},{"cell_type":"code","source":"class AffNetHardNetDetector:\n    def __init__(\n        self,\n        model,\n        device=torch.device(\"cuda\"),\n        resize_long_edge_to=600,\n        matcher=\"adalam\",\n        min_matches=15,\n        rgb_input = False\n    ):\n        self.rgb_input = rgb_input\n        print(\"Init AffNetHardNetDetector\")\n        self.model = model\n        self.device = device\n        self.resize_long_edge_to = resize_long_edge_to\n        print(\"Longer edge will be resized to\", self.resize_long_edge_to)\n\n    def detect_features(self, img_fnames):\n        f_lafs = dict()\n        f_descs = dict()\n        f_kpts = dict()\n        f_raw_size = dict()\n        f_matches = dict()\n        # Get features\n        print(\"Detecting AffNetHardNet features\")\n        for img_path in tqdm(img_fnames):\n            img_fname = img_path.split(\"/\")[-1]\n            key = img_fname\n            f_matches[key] = dict()\n            with torch.inference_mode():\n                timg = load_torch_image(img_path, device=device)\n                raw_size = torch.tensor(timg.shape[2:])\n                timg_resized, h_scale, w_scale = resize_torch_image(\n                    timg, self.resize_long_edge_to, disable_enlarge=True\n                )\n                if self.rgb_input:\n                    lafs, resps, descs = self.model(timg_resized)\n                else:\n                    lafs, resps, descs = self.model(K.color.rgb_to_grayscale(timg_resized))\n                \n                # Recover scale?\n                lafs[:, :, 0, :] *= 1 / w_scale\n                lafs[:, :, 1, :] *= 1 / h_scale\n                desc_dim = descs.shape[-1]\n                # Move keypoints to cpu for later colmap operations\n                kpts = KF.get_laf_center(lafs).reshape(-1, 2).detach().cpu().numpy()\n                descs = descs.reshape(-1, desc_dim).detach()\n                f_lafs[key] = lafs.detach()\n                f_kpts[key] = kpts\n                f_descs[key] = descs\n                f_raw_size[key] = raw_size\n        gc.collect()\n        torch.cuda.empty_cache()\n        return f_lafs, f_kpts, f_descs, f_raw_size\n\n","metadata":{"execution":{"iopub.status.busy":"2025-06-02T05:24:05.668462Z","iopub.execute_input":"2025-06-02T05:24:05.668724Z","iopub.status.idle":"2025-06-02T05:24:05.688741Z","shell.execute_reply.started":"2025-06-02T05:24:05.668701Z","shell.execute_reply":"2025-06-02T05:24:05.687933Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Scene LAF matcher","metadata":{}},{"cell_type":"code","source":"class LafMatcher:\n    def __init__(self, min_matches=15, device=\"cuda\", matcher=\"adalam\", min_pairs=50, distances_threshold=0.3):\n        self.adalam_config = KF.adalam.get_adalam_default_config()\n        self.adalam_config[\"force_seed_mnn\"] = True\n        self.adalam_config[\"search_expansion\"] = 16\n        self.adalam_config[\"ransac_iters\"] = 256\n        self.adalam_config[\"device\"] = device\n        self.min_matches = min_matches\n        self.matcher = matcher\n        self.min_pairs = min_pairs\n        self.distances_threshold = distances_threshold\n        self.tolerance = 500\n        \n    def get_pairs(self, img_fnames):\n        \"\"\"Get image pairs using DINOv2 embeddings for similarity\"\"\"\n        print(\"Getting image pairs using DINOv2...\")\n        # Load DINOv2 model\n        processor = AutoImageProcessor.from_pretrained('/kaggle/input/dinov2/pytorch/base/1/', use_fast=True)\n        model = AutoModel.from_pretrained('/kaggle/input/dinov2/pytorch/base/1/').eval().to(self.adalam_config[\"device\"])\n        embeddings = []\n        \n        # Get embeddings for all images\n        for img_path in tqdm(img_fnames):\n            image = K.io.load_image(img_path, K.io.ImageLoadType.RGB32, device=self.adalam_config[\"device\"])[None, ...]\n            with torch.inference_mode():\n                inputs = processor(images=image, return_tensors=\"pt\", do_rescale=False, do_resize=True, \n                                do_center_crop=False, size=224).to(self.adalam_config[\"device\"])\n                outputs = model(**inputs)\n                embedding = F.normalize(outputs.last_hidden_state.max(dim=1)[0])\n            embeddings.append(embedding)\n            \n        embeddings = torch.cat(embeddings, dim=0)\n        distances = torch.cdist(embeddings, embeddings).cpu()\n        distances_ = (distances <= self.distances_threshold).numpy()\n        np.fill_diagonal(distances_, False)\n        \n        # Ensure minimum number of pairs per image\n        z = distances_.sum(axis=1)\n        idxs0 = np.where(z == 0)[0]\n        for idx0 in idxs0:\n            t = np.argsort(distances[idx0])[1:self.min_pairs]\n            distances_[idx0, t] = True\n\n        s = np.where(distances >= self.tolerance)\n        distances_[s] = False\n            \n        # Convert to pairs format\n        pairs = []\n        for i in range(len(img_fnames)):\n            for j in range(i + 1, len(img_fnames)):\n                if distances_[i][j]:\n                    pairs.append([i, j])\n                    \n        print(f\"Found {len(pairs)} pairs\")\n        return pairs\n\n    def match(self, img_fnames, f_lafs, f_kpts, f_descs, f_raw_size, get_roi=False):\n        \"\"\"Match features between image pairs\"\"\"\n        index_pairs = self.get_pairs(img_fnames)\n        \n        f_matches = defaultdict(dict)\n        f_rois = defaultdict(dict)\n        print(\"Matching features for selected pairs\")\n        \n        for idx1, idx2 in tqdm(index_pairs):\n            fname1, fname2 = img_fnames[idx1], img_fnames[idx2]\n            key1, key2 = fname1.split(\"/\")[-1], fname2.split(\"/\")[-1]\n            lafs1 = f_lafs[key1]\n            lafs2 = f_lafs[key2]\n            desc1 = f_descs[key1]\n            desc2 = f_descs[key2]\n            \n            if self.matcher == \"adalam\":\n                hw1, hw2 = f_raw_size[key1], f_raw_size[key2]\n                dists, idxs = KF.match_adalam(\n                    desc1,\n                    desc2,\n                    lafs1,\n                    lafs2,\n                    hw1=hw1,\n                    hw2=hw2,\n                    config=self.adalam_config,\n                )\n            else:\n                dists, idxs = KF.match_smnn(desc1, desc2, 0.98)\n\n            if dists.mean().cpu().numpy() < 0.5:\n                first_indices = get_unique_idxs(idxs[:, 1])\n                idxs = idxs[first_indices]\n                dists = dists[first_indices]\n                n_matches = len(idxs)\n                if n_matches >= self.min_matches:\n                    f_matches[key1][key2] = idxs.detach().cpu().numpy().reshape(-1, 2)\n\n                    # Compute ROI if requested\n                    if get_roi:\n                        mkpts1 = f_kpts[key1][idxs.cpu().numpy()[:, 0]]\n                        mkpts2 = f_kpts[key2][idxs.cpu().numpy()[:, 1]]\n                        roi_min_w_1, roi_max_w_1 = np.percentile(mkpts1[:, 0], [5, 95])\n                        roi_min_h_1, roi_max_h_1 = np.percentile(mkpts1[:, 1], [5, 95])\n                        roi_area_1 = (roi_max_w_1 - roi_min_w_1) * (roi_max_h_1 - roi_min_h_1)\n                        roi1 = {\n                            \"roi_min_w\": roi_min_w_1,\n                            \"roi_min_h\": roi_min_h_1,\n                            \"roi_max_w\": roi_max_w_1,\n                            \"roi_max_h\": roi_max_h_1,\n                            \"area\": roi_area_1,\n                        }\n                        roi_min_w_2, roi_max_w_2 = np.percentile(mkpts2[:, 0], [5, 95])\n                        roi_min_h_2, roi_max_h_2 = np.percentile(mkpts2[:, 1], [5, 95])\n                        roi_area_2 = (roi_max_w_2 - roi_min_w_2) * (roi_max_h_2 - roi_min_h_2)\n                        roi2 = {\n                            \"roi_min_w\": roi_min_w_2,\n                            \"roi_min_h\": roi_min_h_2,\n                            \"roi_max_w\": roi_max_w_2,\n                            \"roi_max_h\": roi_max_h_2,\n                            \"area\": roi_area_2,\n                        }\n                        f_rois[key1][key2] = [roi1, roi2]\n\n        print(f\"Successfully matched {len(f_matches)} pairs\")\n        torch.cuda.empty_cache()\n        gc.collect()\n        return index_pairs, f_kpts, f_matches, f_rois\n    ","metadata":{"execution":{"iopub.status.busy":"2025-06-02T05:24:05.689695Z","iopub.execute_input":"2025-06-02T05:24:05.689960Z","iopub.status.idle":"2025-06-02T05:24:05.711373Z","shell.execute_reply.started":"2025-06-02T05:24:05.689940Z","shell.execute_reply":"2025-06-02T05:24:05.710674Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# MASt3RDetector","metadata":{}},{"cell_type":"code","source":"class MASt3RDetector:\n    def __init__(self, model, device=torch.device(\"cuda\"), resize_long_edge_to=512, min_pairs=50, distances_threshold=0.3, min_matches=15):\n        from dust3r.inference import inference\n        from mast3r.cloud_opt.sparse_ga import extract_correspondences, symmetric_inference\n        from mast3r.utils.misc import hash_md5\n        from dust3r.utils.device import to_cpu\n        import os\n        \n        self.inference = inference\n        self.symmetric_inference = symmetric_inference\n        self.extract_correspondences = extract_correspondences\n        self.hash_md5 = hash_md5\n        self.to_cpu = to_cpu\n        self.model = model\n        self.device = device\n        self.resize_long_edge_to = resize_long_edge_to\n        self.min_pairs = min_pairs\n        self.distances_threshold = distances_threshold\n        self.min_matches = min_matches\n        self.tolerance = 500\n        print(\"Init MASt3RDetector\")\n        print(\"Longer edge will be resized to\", self.resize_long_edge_to)\n\n    def remove_duplicate_matches(self, matches):\n        \"\"\"Remove duplicate matches by keeping only unique pairs.\"\"\"\n        unique_matches = {}\n        for key1 in matches:\n            unique_matches[key1] = {}\n            for key2 in matches[key1]:\n                # Convert matches to a set of tuples for uniqueness\n                match_set = set(map(tuple, matches[key1][key2]))\n                # Convert back to numpy array\n                unique_matches[key1][key2] = np.array(list(match_set))\n        return unique_matches\n\n    def get_pairs(self, img_fnames):\n        \"\"\"Get image pairs using DINOv2 embeddings for similarity\"\"\"\n        print(\"Getting image pairs using DINOv2...\")\n        # Load DINOv2 model\n        processor = AutoImageProcessor.from_pretrained('/kaggle/input/dinov2/pytorch/base/1/')\n        model = AutoModel.from_pretrained('/kaggle/input/dinov2/pytorch/base/1/').eval().to(self.device)\n        embeddings = []\n        \n        # Get embeddings for all images\n        for img_path in tqdm(img_fnames):\n            image = K.io.load_image(img_path, K.io.ImageLoadType.RGB32, device=self.device)[None, ...]\n            with torch.inference_mode():\n                inputs = processor(images=image, return_tensors=\"pt\", do_rescale=False, do_resize=True, \n                                do_center_crop=False, size=224).to(self.device)\n                outputs = model(**inputs)\n                embedding = F.normalize(outputs.last_hidden_state.max(dim=1)[0])\n            embeddings.append(embedding)\n            \n        embeddings = torch.cat(embeddings, dim=0)\n        distances = torch.cdist(embeddings, embeddings).cpu()\n        distances_ = (distances <= self.distances_threshold).numpy()\n        np.fill_diagonal(distances_, False)\n        \n        # Ensure minimum number of pairs per image\n        z = distances_.sum(axis=1)\n        idxs0 = np.where(z == 0)[0]\n        for idx0 in idxs0:\n            t = np.argsort(distances[idx0])[1:self.min_pairs]\n            distances_[idx0, t] = True\n        \n        s = np.where(distances >= self.tolerance)\n        distances_[s] = False\n            \n        # Convert to pairs format\n        pairs = []\n        for i in range(len(img_fnames)):\n            for j in range(i + 1, len(img_fnames)):\n                if distances_[i][j]:\n                    pairs.append([i, j])\n                    \n        print(f\"Found {len(pairs)} pairs using DINOv2\")\n        return pairs\n\n    def convert_matches_to_colmap_format(self, img0, img1, matches_im0, matches_im1):\n        \"\"\"Convert 2D matches to COLMAP format using ravel indices\"\"\"\n        matches = [matches_im0.astype(np.float64), matches_im1.astype(np.float64)]\n        ravel_matches = []\n        \n        for j in range(2):\n            H, W = [img0, img1][j]['true_shape'][0]\n            qx, qy = matches[j].round().astype(np.int32).T\n            # Convert 2D coordinates to ravel indices\n            ravel_matches_j = qx.clip(min=0, max=W - 1, out=qx) + W * qy.clip(min=0, max=H - 1, out=qy)\n            ravel_matches.append(ravel_matches_j)\n            \n        # Stack matches and ensure proper ordering\n        colmap_matches = np.stack([ravel_matches[0], ravel_matches[1]], axis=-1)\n        # Remove duplicates\n        colmap_matches = np.unique(colmap_matches, axis=0)\n        return colmap_matches\n\n    @torch.no_grad()\n    def detect_features(self, img_fnames):\n        f_kpts = dict()\n        f_descs = dict()\n        f_matches = defaultdict(dict)\n        \n        # Convert image paths to MASt3R format\n        imgs = []\n        for i, fname in enumerate(img_fnames):\n            # Load image to get shape\n            img = load_torch_image(fname, device=self.device)\n            # Get original size for scaling\n            orig_h, orig_w = img.shape[2:]\n            # Resize for MASt3R\n            img_resized, h_scale, w_scale = resize_torch_image(img, self.resize_long_edge_to, align=16)\n            \n            # Create dense grid of keypoints\n            h, w = img_resized.shape[2:]\n            y, x = np.mgrid[0:h, 0:w]\n            pts = np.stack([x, y], axis=-1).reshape(-1, 2)\n            \n            # Scale back to original image size\n            pts = pts * np.array([1.0/w_scale, 1.0/h_scale])\n            \n            # Store keypoints\n            key = fname.split(\"/\")[-1]\n            f_kpts[key] = pts\n            \n            imgs.append({\n                'instance': fname, \n                'img': None,  # Don't store image yet\n                'idx': i,\n                'true_shape': np.int32([img_resized.shape[2:]]),\n                'h_scale': float(h_scale),  # Store scale factors\n                'w_scale': float(w_scale),\n                'orig_h': orig_h,  # Store original dimensions\n                'orig_w': orig_w\n            })\n            del img, img_resized\n            \n        # Get pairs using DINOv2\n        pairs = self.get_pairs(img_fnames)\n        print(f\"Processing {len(pairs)} pairs...\")\n\n        # Process pairs\n        for idx1, idx2 in tqdm(pairs):\n            # Get file names\n            key1 = img_fnames[idx1].split(\"/\")[-1]\n            key2 = img_fnames[idx2].split(\"/\")[-1]\n            \n            # Get corresponding img objects\n            img1 = next(img for img in imgs if img['idx'] == idx1)\n            img2 = next(img for img in imgs if img['idx'] == idx2)\n\n            # Load images for this pair\n            img = load_torch_image(img1['instance'], device=self.device)\n            img_resized, _, _ = resize_torch_image(img, self.resize_long_edge_to, align=16)\n            img1['img'] = img_resized\n            \n            img = load_torch_image(img2['instance'], device=self.device)\n            img_resized, _, _ = resize_torch_image(img, self.resize_long_edge_to, align=16)\n            img2['img'] = img_resized\n\n            # Run symmetric inference\n            res = self.symmetric_inference(self.model, img1, img2, device=self.device)\n            \n            # Extract points and confidences\n            descs = [r['desc'][0] for r in res]\n            qonfs = [r['desc_conf'][0] for r in res]\n            (xy1, xy2, confs) = self.extract_correspondences(descs, qonfs, device=self.device, subsample=8)\n            \n            # Filter by confidence\n            mask = confs >= 2.5\n            xy1 = xy1[mask].cpu().numpy()\n            xy2 = xy2[mask].cpu().numpy()\n            \n            if len(xy1) >= self.min_matches:\n                # Convert to indices in the dense grid\n                h1, w1 = img1['true_shape'][0]\n                h2, w2 = img2['true_shape'][0]\n                idx1 = np.round(xy1[:, 0]).astype(int) + w1 * np.round(xy1[:, 1]).astype(int)\n                idx2 = np.round(xy2[:, 0]).astype(int) + w2 * np.round(xy2[:, 1]).astype(int)\n                \n                # Store matches\n                f_matches[key1][key2] = np.column_stack([idx1, idx2])\n\n            # Clear GPU memory\n            del img1['img']\n            del img2['img']\n            torch.cuda.empty_cache()\n\n        # Remove duplicate matches before returning\n        f_matches = self.remove_duplicate_matches(f_matches)\n\n        # Final cleanup\n        del self.model\n        torch.cuda.empty_cache()\n        \n        return f_kpts, f_descs, f_matches\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-02T05:24:05.712306Z","iopub.execute_input":"2025-06-02T05:24:05.712828Z","iopub.status.idle":"2025-06-02T05:24:05.735329Z","shell.execute_reply.started":"2025-06-02T05:24:05.712800Z","shell.execute_reply":"2025-06-02T05:24:05.734603Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Matches and pairs operation","metadata":{}},{"cell_type":"code","source":"def merge_kpts_matches(kpts, matches, new_kpts, new_matches, cap = None):\n    # merge kpts\n    prev_len = dict()\n    for new_key in new_kpts:\n        if new_key in kpts:\n            old_len = len(kpts[new_key])\n            kpts[new_key] = np.concatenate([kpts[new_key], new_kpts[new_key]], axis=0)\n        else:\n            old_len = 0\n            kpts[new_key] = new_kpts[new_key]\n        prev_len[new_key] = old_len\n\n    for new_key1 in new_matches:\n        for new_key2 in new_matches[new_key1]:\n            old_len1 = prev_len[new_key1]\n            old_len2 = prev_len[new_key2]\n            new_match = new_matches[new_key1][new_key2] + [old_len1, old_len2]\n            if cap is not None and len(new_match) > cap:\n                keep = np.random.choice(len(new_match), cap, replace=False)\n                new_match = new_match[keep, :]\n            if new_key1 in matches and new_key2 in matches[new_key1]:\n\n                matches[new_key1][new_key2] = np.concatenate(\n                    [\n                        matches[new_key1][new_key2],\n                        new_match,\n                    ],\n                    axis=0,\n                )\n            else:\n                if new_key1 not in matches:\n                    matches[new_key1] = dict()\n                matches[new_key1][new_key2] = new_match\n    return kpts, matches\n\n\ndef keep_matches(matches, max_num=None):\n    if max_num is None:\n        return matches\n    if len(matches) > max_num:\n        # radnomly select max_num matches\n        matches = np.random.choice(matches, max_num, replace=False)\n    return matches\n\n\ndef keep_pairs(index_pairs, max_num_pairs=20):\n    new_count = 0\n    old_count = 0\n    new_idx_count = defaultdict(int)\n    new_pairs = defaultdict(list)\n    for key1 in index_pairs:\n        # sort pairs by number of pairs\n        index_pairs[key1] = sorted(index_pairs[key1], key=lambda x: x[2], reverse=True)\n        for pair in index_pairs[key1]:\n            old_count += 1\n            idx1 = key1\n            idx2 = pair[0]\n\n            if new_idx_count[key1] < max_num_pairs:\n                new_pairs[idx1].append(pair)\n                new_count += 1\n                new_idx_count[idx1] += 1\n                new_idx_count[idx2] += 1\n            else:\n                continue\n\n    if DEBUG:\n        print(f\"origin pairs: {old_count}, kept pairs: {new_count}\")\n    return index_pairs\n\ndef select_matches(matches, keep_ratio = 0.01):\n    max_matches = defaultdict(int)\n    old_matches_count = 0\n    for key1 in matches:\n        for key2 in matches[key1]:\n            max_matches[key1] = max(max_matches[key1], len(matches[key1][key2]))\n            max_matches[key2] = max(max_matches[key2], len(matches[key1][key2]))\n            old_matches_count +=1\n\n    new_matches_count = 0\n    new_matches = defaultdict(dict)\n    for key1 in matches:\n        for key2 in matches[key1]:\n            n_matches = len(matches[key1][key2])\n            if n_matches > max_matches[key1] * keep_ratio or n_matches > max_matches[key2] * keep_ratio:\n                new_matches[key1][key2] = matches[key1][key2]\n                new_matches_count+=1\n    if DEBUG:\n        print(f\"origin matches: {old_matches_count}, kept matches: {new_matches_count}\")\n    return new_matches","metadata":{"execution":{"iopub.status.busy":"2025-06-02T05:24:05.736205Z","iopub.execute_input":"2025-06-02T05:24:05.736698Z","iopub.status.idle":"2025-06-02T05:24:05.760092Z","shell.execute_reply.started":"2025-06-02T05:24:05.736673Z","shell.execute_reply":"2025-06-02T05:24:05.759433Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Get Fundamental matrices from matches and keypoints","metadata":{}},{"cell_type":"code","source":"def get_fms(kpts, matches):\n    prev_len = dict()\n    fms = defaultdict(dict)\n    print(\"Get Fundamental Matrix\")\n    for key1 in tqdm(matches):\n        for key2 in matches[key1]:\n            match = matches[key1][key2]\n            mkpts1 = kpts[key1][match[:, 0]]\n            mkpts2 = kpts[key2][match[:, 1]]\n            Fm, inliers = cv2.findFundamentalMat(mkpts1, mkpts2, cv2.USAC_MAGSAC, FM_PARAMS[\"ransacReprojThreshold\"], FM_PARAMS[\"confidence\"], FM_PARAMS[\"maxIters\"])\n            #keep inliers matches\n            #print how many matches are inliers\n            # print(f\"key1: {key1}, key2: {key2}, inliers: {len(new_match)}/{len(match)}\")\n            if FM_PARAMS[\"removeOutliers\"] == True:\n                new_match = match[inliers.ravel() == 1]\n                matches[key1][key2] = new_match\n            fms[key1][key2] = Fm\n    # print(Fm.shape)\n    return kpts, matches, fms","metadata":{"execution":{"iopub.status.busy":"2025-06-02T05:24:05.761000Z","iopub.execute_input":"2025-06-02T05:24:05.761260Z","iopub.status.idle":"2025-06-02T05:24:05.781511Z","shell.execute_reply.started":"2025-06-02T05:24:05.761238Z","shell.execute_reply":"2025-06-02T05:24:05.780890Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Model Setup","metadata":{}},{"cell_type":"code","source":"if MODEL_DICT[\"Keynet\"][\"enable\"]:\n    keynet_model = (\n        AffNetHardNet(num_features=8000, upright=False, device=device, detector=\"keynet\")\n        .to(device)\n        .eval()\n    )\n    keynet_detector = AffNetHardNetDetector(keynet_model, resize_long_edge_to=MODEL_DICT[\"Keynet\"][\"resize_long_edge_to\"])\n    laf_matcher = LafMatcher(device=device)\n    \nif MODEL_DICT[\"GFTT\"][\"enable\"]:\n    gftt_model = (\n        AffNetHardNet(num_features=8000, upright=False, device=device, detector=\"GFTT\")\n        .to(device)\n        .eval()\n    )\n    gftt_detector = AffNetHardNetDetector(gftt_model, resize_long_edge_to=MODEL_DICT[\"GFTT\"][\"resize_long_edge_to\"])\n    laf_matcher = LafMatcher(device=device)\n\nif MODEL_DICT[\"DoG\"][\"enable\"]:\n    DoG_model = (\n        AffNetHardNet(num_features=8000, upright=False, device=device, detector=\"DoG\")\n        .to(device)\n        .eval()\n    )\n    DoG_detector = AffNetHardNetDetector(DoG_model, resize_long_edge_to=MODEL_DICT[\"DoG\"][\"resize_long_edge_to\"])\n    laf_matcher = LafMatcher(device=device)\n\nif MODEL_DICT[\"Harris\"][\"enable\"]:\n    harris_model = (\n        AffNetHardNet(num_features=8000, upright=False, device=device, detector=\"Harris\")\n        .to(device)\n        .eval()\n    )\n    harris_detector = AffNetHardNetDetector(harris_model, resize_long_edge_to=MODEL_DICT[\"Harris\"][\"resize_long_edge_to\"])\n    laf_matcher = LafMatcher(device=device)\n\n\nif MODEL_DICT[\"MASt3R\"][\"enable\"]:\n    def load_model(model_path, device, verbose=True):\n        if verbose:\n            print('... loading model from', model_path)\n        ckpt = torch.load(model_path, map_location='cpu', weights_only=False )\n        args = ckpt['args'].model.replace(\"ManyAR_PatchEmbed\", \"PatchEmbedDust3R\")\n        if 'landscape_only' not in args:\n            args = args[:-1] + ', landscape_only=False)'\n        else:\n            args = args.replace(\" \", \"\").replace('landscape_only=True', 'landscape_only=False')\n        assert \"landscape_only=False\" in args\n        if verbose:\n            print(f\"instantiating : {args}\")\n        inf = float('inf')\n        net = eval(args)\n        s = net.load_state_dict(ckpt['model'], strict=False)\n        if verbose:\n            print(s)\n        return net.to(device)\n    model_path = \"/kaggle/input/mast3r-vitlarge-basedecoder-512-catmlpdpt-metric/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric.pth\"\n    mast3r_model = load_model(model_path, device='cpu').to(device)\n    mast3r_detector = MASt3RDetector(\n        mast3r_model,\n        device=device,\n        resize_long_edge_to=MODEL_DICT[\"MASt3R\"][\"resize_long_edge_to\"]\n    )","metadata":{"execution":{"iopub.status.busy":"2025-06-02T05:24:05.782352Z","iopub.execute_input":"2025-06-02T05:24:05.782606Z","iopub.status.idle":"2025-06-02T05:24:08.097043Z","shell.execute_reply.started":"2025-06-02T05:24:05.782590Z","shell.execute_reply":"2025-06-02T05:24:08.096222Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Util to check if the pairs are identital\ndef compare_pairs(pairs1, pairs2):\n    pair1_dict = dict()\n    pair2_dict = dict()\n    for idx1 in tqdm(range(len(pairs1) - 1)):\n        for pair in pairs1[idx1]:\n            pair1_dict[(idx1, pair[0])] = 0\n    for idx2 in tqdm(range(len(pairs2) - 1)):\n        for pair in pairs2[idx2]:\n            pair2_dict[(idx2, pair[0])] = 0\n\n    for key in tqdm(pair1_dict):\n        if key not in pair2_dict:\n            print(f\"Key{key} not in pair2_dict\")\n    for key in tqdm(pair2_dict):\n        if key not in pair1_dict:\n            print(f\"Key{key} not in pair1_dict\")","metadata":{"execution":{"iopub.status.busy":"2025-06-02T05:24:08.097919Z","iopub.execute_input":"2025-06-02T05:24:08.098148Z","iopub.status.idle":"2025-06-02T05:24:08.103295Z","shell.execute_reply.started":"2025-06-02T05:24:08.098130Z","shell.execute_reply":"2025-06-02T05:24:08.102592Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Function to generate scene db for reconstruction","metadata":{}},{"cell_type":"code","source":"def generate_scene_db(dataset, scene):\n    feature_det_start = time()\n    # Process a scene and write matches and keypoints to the database\n    img_dir = f\"{SRC}/{MODE}/{dataset}/{scene}/images\"\n    if not os.path.exists(img_dir):\n        print(\"Image dir does not exist:\", img_dir)\n        return\n\n    img_fnames = [f\"{SRC}/{MODE}/{x}\" for x in data_dict[dataset][scene]]\n    print(f\"Got {len(img_fnames)} images\")\n\n    matches = dict()\n    kpts = dict()\n\n    if MODEL_DICT[\"Keynet\"][\"enable\"]:\n        f_lafs, f_kpts, f_descs, f_raw_size = keynet_detector.detect_features(\n            img_fnames\n        )\n        keynet_pairs, keynet_kpts, keynet_matches, keynet_rois = laf_matcher.match(\n            img_fnames, f_lafs, f_kpts, f_descs, f_raw_size\n        )\n        if not MODEL_DICT[\"Keynet\"][\"pair_only\"]:\n            kpts, matches = merge_kpts_matches(kpts, matches, keynet_kpts, keynet_matches, MATCHES_CAP)\n\n    if MODEL_DICT[\"GFTT\"][\"enable\"]:\n        gftt_lafs, gftt_kpts, gftt_descs, gftt_raw_size = gftt_detector.detect_features(\n            img_fnames\n        )\n        index_pairs, gftt_kpts, gftt_matches, gftt_rois = laf_matcher.match(\n            img_fnames, gftt_lafs, gftt_kpts, gftt_descs, gftt_raw_size\n        )\n        kpts, matches = merge_kpts_matches(kpts, matches, gftt_kpts, gftt_matches, MATCHES_CAP)\n\n    if MODEL_DICT[\"DoG\"][\"enable\"]:\n        DoG_lafs, DoG_kpts, DoG_descs, DoG_raw_size = DoG_detector.detect_features(\n            img_fnames\n        )\n        index_pairs, DoG_kpts, DoG_matches, DoG_rois = laf_matcher.match(\n            img_fnames, DoG_lafs, DoG_kpts, DoG_descs, DoG_raw_size\n        )\n        kpts, matches = merge_kpts_matches(kpts, matches, DoG_kpts, DoG_matches, MATCHES_CAP)\n        \n    if MODEL_DICT[\"Harris\"][\"enable\"]:\n        harris_lafs, harris_kpts, harris_descs, harris_raw_size = harris_detector.detect_features(\n            img_fnames\n        )\n        harris_pairs, harris_kpts, harris_matches, harris_rois = laf_matcher.match(\n            img_fnames, harris_lafs, harris_kpts, harris_descs, harris_raw_size\n        )\n        kpts, matches = merge_kpts_matches(kpts, matches, harris_kpts, harris_matches, MATCHES_CAP)\n        # compare_pairs(keynet_pairs, harris_pairs)\n\n    if MODEL_DICT[\"MASt3R\"][\"enable\"]:\n        mast3r_kpts, mast3r_descs, mast3r_matches = mast3r_detector.detect_features(img_fnames)\n        kpts, matches = merge_kpts_matches(kpts, matches, mast3r_kpts, mast3r_matches, MATCHES_CAP)\n    \n    # Get fundamental matrices\n    kpts, matches, fms = get_fms(kpts, matches)\n    \n    matches = select_matches(matches, MATCH_FILTER_RATIO)\n\n    if DEBUG:\n        import random\n        random.seed(0)\n        for i in range(5):\n            print(matches.keys())\n           \n            key1 = random.choice(list(matches.keys()))\n            key2 = random.choice(list(matches[key1].keys()))\n            print(key1, key2)\n            fname1, fname2 = os.path.join(img_dir, key1), os.path.join(img_dir, key2)\n\n            print(\"Plot Combined matches\")\n            plot_images_with_keypoints(\n                fname1, fname2, kpts[key1], kpts[key2], matches[key1][key2]\n            )\n    # Write to database\n    feature_dir = f\"featureout/{dataset}_{scene}\"\n    if not os.path.isdir(feature_dir):\n        os.makedirs(feature_dir, exist_ok=True)\n    database_path = f\"{feature_dir}/colmap.db\"\n    if os.path.isfile(database_path):\n        os.remove(database_path)\n\n    db = COLMAPDatabase.connect(database_path)\n    db.create_tables()\n    single_camera = False\n    print(\"Add kpts and matches to database\")\n    add_kpts_matches(db, img_dir, kpts, matches, fms)\n    feature_det_end = time()\n    matching_time = feature_det_end - feature_det_start\n    torch.cuda.empty_cache()\n    gc.collect()\n\n    return matching_time","metadata":{"execution":{"iopub.status.busy":"2025-06-02T05:24:08.104355Z","iopub.execute_input":"2025-06-02T05:24:08.104801Z","iopub.status.idle":"2025-06-02T05:24:08.125491Z","shell.execute_reply.started":"2025-06-02T05:24:08.104784Z","shell.execute_reply":"2025-06-02T05:24:08.124752Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Function of reconstruction","metadata":{}},{"cell_type":"code","source":"def reconstruct_from_db(dataset, scene):\n    scene_result = {}\n    reconst_start = time()\n\n    img_dir = f\"{SRC}/{MODE}/{dataset}/{scene}/images\"\n    if not os.path.exists(img_dir):\n        print(\"Image dir does not exist:\", img_dir)\n        return\n\n    feature_dir = f\"featureout/{dataset}_{scene}\"\n    database_path = f\"{feature_dir}/colmap.db\"\n    db = COLMAPDatabase.connect(database_path)\n    output_path = f\"{feature_dir}/colmap_rec\"\n    t = time()\n    gc.collect()\n\n    t = time() - t\n    print(f\"RANSAC in  {t:.4f} sec\")\n    t = time()\n    \n    # Create pipeline options with min_num_reg_images = 3\n    pipeline_options = pycolmap.IncrementalPipelineOptions()\n    # Basic parameters\n    pipeline_options.min_model_size = 3\n    pipeline_options.max_model_overlap = 20\n    pipeline_options.min_focal_length_ratio = 0.1\n    pipeline_options.max_focal_length_ratio = 10\n    pipeline_options.max_extra_param = 1.0\n    \n    # Triangulation parameters\n    pipeline_options.min_tri_angle = 1.5\n    pipeline_options.max_reproj_error = 4.0\n    pipeline_options.min_track_length = 3\n    pipeline_options.max_track_length = 100\n    \n    # Bundle adjustment parameters\n    pipeline_options.abs_pose_max_error = 12.0\n    pipeline_options.abs_pose_min_num_inliers = 30\n    pipeline_options.abs_pose_min_inlier_ratio = 0.25\n    pipeline_options.filter_max_reproj_error = 4.0\n    pipeline_options.filter_min_tri_angle = 1.5\n    \n    # Local bundle adjustment parameters\n    pipeline_options.local_ba_min_tri_angle = 6.0\n    \n    pipeline_options.min_tri_angle = 2.0\n    pipeline_options.filter_min_tri_angle = 2.0\n    pipeline_options.local_ba_min_tri_angle = 8.0\n    pipeline_options.max_reproj_error = 3.0\n    pipeline_options.filter_max_reproj_error = 3.0\n\n    os.makedirs(output_path, exist_ok=True)\n    maps = pycolmap.incremental_mapping(\n        database_path=database_path,\n        image_path=img_dir,\n        output_path=output_path,\n        options=pipeline_options\n    )\n    print(maps)\n    t = time() - t\n    print(f\"Reconstruction done in  {t:.4f} sec\")\n    imgs_registered = 0\n    best_idx = None\n    print(\"Looking for the best reconstruction\")\n    if isinstance(maps, dict):\n        for idx1, rec in maps.items():\n            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            scene_result[key1] = {}\n            scene_result[key1][\"R\"] = deepcopy(im.cam_from_world.rotation.matrix())\n            scene_result[key1][\"t\"] = deepcopy(np.array(im.cam_from_world.translation))\n\n    gc.collect()\n    reconst_end = time()\n    reconst_time = reconst_end - reconst_start\n    return scene_result, reconst_time","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2025-06-02T05:24:08.126511Z","iopub.execute_input":"2025-06-02T05:24:08.127210Z","iopub.status.idle":"2025-06-02T05:24:08.148194Z","shell.execute_reply.started":"2025-06-02T05:24:08.127189Z","shell.execute_reply":"2025-06-02T05:24:08.147520Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Main Loop","metadata":{}},{"cell_type":"code","source":"# Main loop to add kpts and matches\ndatasets = []\ntime_dict = dict()\n\nfor dataset in data_dict:\n    datasets.append(dataset)\n\nif DEBUG:\n    matching_start = time()\n    for dataset, scene in all_scenes:\n        print(dataset, scene)\n        time_dict[\"matching-\" + scene] = generate_scene_db(dataset, scene)\n\n    matching_end = time()\n    time_dict[\"matching-TOTAL\"] = matching_end - matching_start\nelse:\n    # Run db generation and reconstuction with multiprocessing if not DEBUG\n    out_results = defaultdict(dict)\n    total_start = time()\n    with concurrent.futures.ProcessPoolExecutor(max_workers=NUM_CORES) as executors:\n        futures = defaultdict(dict)\n\n        for dataset, scene in all_scenes:\n            print(dataset, scene)\n            time_dict[\"matching-\" + scene] = generate_scene_db(dataset, scene)\n            futures[dataset][scene] = executors.submit(reconstruct_from_db, dataset, scene)\n\n        for dataset, scene in all_scenes:\n            result = futures[dataset][scene].result()\n            if result is not None:\n                out_results[dataset][scene], time_dict[\"reconst-\" + scene] = result\n    total_end = time()\n    time_dict[\"TOTAL\"] = total_end - total_start","metadata":{"_kg_hide-output":true,"scrolled":true,"execution":{"iopub.status.busy":"2025-06-02T05:24:08.149002Z","iopub.execute_input":"2025-06-02T05:24:08.149237Z","iopub.status.idle":"2025-06-02T05:26:22.493827Z","shell.execute_reply.started":"2025-06-02T05:24:08.149214Z","shell.execute_reply":"2025-06-02T05:26:22.493140Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Reconstruction Loop","metadata":{}},{"cell_type":"code","source":"# Main loop for reconstruction\nif DEBUG:\n    reconst_start = time()\n    out_results = defaultdict(dict)\n    with concurrent.futures.ProcessPoolExecutor(max_workers=NUM_CORES) as executors:\n        futures = defaultdict(dict)\n        for dataset, scene in all_scenes:\n            futures[dataset][scene] = executors.submit(reconstruct_from_db, dataset, scene)\n            print(\"Submitted\", dataset, scene)\n        for dataset, scene in all_scenes:\n            result = futures[dataset][scene].result()\n            if result is not None:\n                out_results[dataset][scene], time_dict[\"reconst-\" + scene] = result\n\n    reconst_end = time()\n    time_dict[\"reconst-TOTAL\"] = reconst_end - reconst_start","metadata":{"scrolled":true,"trusted":true,"execution":{"iopub.status.busy":"2025-06-02T05:26:22.495277Z","iopub.execute_input":"2025-06-02T05:26:22.495549Z","iopub.status.idle":"2025-06-02T05:28:26.713165Z","shell.execute_reply.started":"2025-06-02T05:26:22.495528Z","shell.execute_reply":"2025-06-02T05:28:26.712434Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"create_submission(out_results, data_dict, MODE)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-02T05:28:26.714502Z","iopub.execute_input":"2025-06-02T05:28:26.715232Z","iopub.status.idle":"2025-06-02T05:28:26.723891Z","shell.execute_reply.started":"2025-06-02T05:28:26.715206Z","shell.execute_reply":"2025-06-02T05:28:26.723123Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Evaluation Code\nfrom https://www.kaggle.com/code/eduardtrulls/imc2023-evaluation","metadata":{}},{"cell_type":"code","source":"import numpy as np\nfrom dataclasses import dataclass\nfrom time import time\n\n\ndef arr_to_str(a):\n    return \";\".join([str(x) for x in a.reshape(-1)])\n\n\n# Evaluation metric.\n\n\n@dataclass\nclass Camera:\n    rotmat: np.array\n    tvec: np.array\n\n\ndef quaternion_from_matrix(matrix):\n    M = np.array(matrix, dtype=np.float64, copy=False)[:4, :4]\n    m00 = M[0, 0]\n    m01 = M[0, 1]\n    m02 = M[0, 2]\n    m10 = M[1, 0]\n    m11 = M[1, 1]\n    m12 = M[1, 2]\n    m20 = M[2, 0]\n    m21 = M[2, 1]\n    m22 = M[2, 2]\n\n    # Symmetric matrix K.\n    K = np.array(\n        [\n            [m00 - m11 - m22, 0.0, 0.0, 0.0],\n            [m01 + m10, m11 - m00 - m22, 0.0, 0.0],\n            [m02 + m20, m12 + m21, m22 - m00 - m11, 0.0],\n            [m21 - m12, m02 - m20, m10 - m01, m00 + m11 + m22],\n        ]\n    )\n    K /= 3.0\n\n    # Quaternion is eigenvector of K that corresponds to largest eigenvalue.\n    w, V = np.linalg.eigh(K)\n    q = V[[3, 0, 1, 2], np.argmax(w)]\n\n    if q[0] < 0.0:\n        np.negative(q, q)\n    return q\n\n\ndef evaluate_R_t(R_gt, t_gt, R, t, eps=1e-15):\n    t = t.flatten()\n    t_gt = t_gt.flatten()\n\n    q_gt = quaternion_from_matrix(R_gt)\n    q = quaternion_from_matrix(R)\n    q = q / (np.linalg.norm(q) + eps)\n    q_gt = q_gt / (np.linalg.norm(q_gt) + eps)\n    loss_q = np.maximum(eps, (1.0 - np.sum(q * q_gt) ** 2))\n    err_q = np.arccos(1 - 2 * loss_q)\n\n    GT_SCALE = np.linalg.norm(t_gt)\n    t = GT_SCALE * (t / (np.linalg.norm(t) + eps))\n    err_t = min(np.linalg.norm(t_gt - t), np.linalg.norm(t_gt + t))\n\n    return np.degrees(err_q), err_t\n\n\ndef compute_dR_dT(R1, T1, R2, T2):\n    \"\"\"Given absolute (R, T) pairs for two cameras, compute the relative pose difference, from the first.\"\"\"\n\n    dR = np.dot(R2, R1.T)\n    dT = T2 - np.dot(dR, T1)\n    return dR, dT\n\n\ndef compute_mAA(err_q, err_t, ths_q, ths_t):\n    \"\"\"Compute the mean average accuracy over a set of thresholds. Additionally returns the metric only over rotation and translation.\"\"\"\n\n    acc, acc_q, acc_t = [], [], []\n    for th_q, th_t in zip(ths_q, ths_t):\n        cur_acc_q = err_q <= th_q\n        cur_acc_t = err_t <= th_t\n        cur_acc = cur_acc_q & cur_acc_t\n\n        acc.append(cur_acc.astype(np.float32).mean())\n        acc_q.append(cur_acc_q.astype(np.float32).mean())\n        acc_t.append(cur_acc_t.astype(np.float32).mean())\n    return np.array(acc), np.array(acc_q), np.array(acc_t)\n\n\ndef dict_from_csv(csv_path, has_header):\n    csv_dict = {}\n    with open(csv_path, \"r\") as f:\n        for i, l in enumerate(f):\n            if has_header and i == 0:\n                continue\n            if l:\n                image, dataset, scene, R_str, T_str = l.strip().split(\",\")\n                R = np.fromstring(R_str.strip(), sep=\";\").reshape(3, 3)\n                T = np.fromstring(T_str.strip(), sep=\";\")\n                if dataset not in csv_dict:\n                    csv_dict[dataset] = {}\n                if scene not in csv_dict[dataset]:\n                    csv_dict[dataset][scene] = {}\n                csv_dict[dataset][scene][image] = Camera(rotmat=R, tvec=T)\n    return csv_dict\n\n\ndef eval_submission(\n    submission_csv_path,\n    ground_truth_csv_path,\n    rotation_thresholds_degrees_dict,\n    translation_thresholds_meters_dict,\n    verbose=False,\n):\n    \"\"\"Compute final metric given submission and ground truth files. Thresholds are specified per dataset.\"\"\"\n\n    submission_dict = dict_from_csv(submission_csv_path, has_header=True)\n    gt_dict = dict_from_csv(ground_truth_csv_path, has_header=True)\n\n    # Check that all necessary keys exist in the submission file\n    for dataset in gt_dict:\n        assert dataset in submission_dict, f\"Unknown dataset: {dataset}\"\n        for scene in gt_dict[dataset]:\n            assert (\n                scene in submission_dict[dataset]\n            ), f\"Unknown scene: {dataset}->{scene}\"\n            for image in gt_dict[dataset][scene]:\n                assert (\n                    image in submission_dict[dataset][scene]\n                ), f\"Unknown image: {dataset}->{scene}->{image}\"\n\n    # Iterate over all the scenes\n    if verbose:\n        t = time()\n        print(\"*** METRICS ***\")\n\n    metrics_per_dataset = []\n    for dataset in gt_dict:\n        metrics_per_scene = []\n        for scene in gt_dict[dataset]:\n            err_q_all = []\n            err_t_all = []\n            images = [camera for camera in gt_dict[dataset][scene]]\n            # Process all pairs in a scene\n            for i in range(len(images)):\n                for j in range(i + 1, len(images)):\n                    gt_i = gt_dict[dataset][scene][images[i]]\n                    gt_j = gt_dict[dataset][scene][images[j]]\n                    dR_gt, dT_gt = compute_dR_dT(\n                        gt_i.rotmat, gt_i.tvec, gt_j.rotmat, gt_j.tvec\n                    )\n\n                    pred_i = submission_dict[dataset][scene][images[i]]\n                    pred_j = submission_dict[dataset][scene][images[j]]\n                    dR_pred, dT_pred = compute_dR_dT(\n                        pred_i.rotmat, pred_i.tvec, pred_j.rotmat, pred_j.tvec\n                    )\n\n                    err_q, err_t = evaluate_R_t(dR_gt, dT_gt, dR_pred, dT_pred)\n                    err_q_all.append(err_q)\n                    err_t_all.append(err_t)\n\n            mAA, mAA_q, mAA_t = compute_mAA(\n                err_q=err_q_all,\n                err_t=err_t_all,\n                ths_q=rotation_thresholds_degrees_dict[(dataset, scene)],\n                ths_t=translation_thresholds_meters_dict[(dataset, scene)],\n            )\n            if verbose:\n                print(\n                    f\"{dataset} / {scene} ({len(images)} images, {len(err_q_all)} pairs) -> mAA={np.mean(mAA):.06f}, mAA_q={np.mean(mAA_q):.06f}, mAA_t={np.mean(mAA_t):.06f}\"\n                )\n            metrics_per_scene.append(np.mean(mAA))\n\n        metrics_per_dataset.append(np.mean(metrics_per_scene))\n        if verbose:\n            print(f\"{dataset} -> mAA={np.mean(metrics_per_scene):.06f}\")\n            print()\n\n    if verbose:\n        print(\n            f\"Final metric -> mAA={np.mean(metrics_per_dataset):.06f} (t: {time() - t} sec.)\"\n        )\n        print()\n\n    return np.mean(metrics_per_dataset)\n\n\n# Set rotation thresholds per scene.\n\nrotation_thresholds_degrees_dict = {\n    **{\n        (\"haiper\", scene): np.linspace(1, 10, 10)\n        for scene in [\"bike\", \"chairs\", \"fountain\"]\n    },\n    **{(\"heritage\", scene): np.linspace(1, 10, 10) for scene in [\"cyprus\", \"dioscuri\"]},\n    **{(\"heritage\", \"wall\"): np.linspace(0.2, 10, 10)},\n    **{(\"urban\", \"kyiv-puppet-theater\"): np.linspace(1, 10, 10)},\n}\n\ntranslation_thresholds_meters_dict = {\n    **{\n        (\"haiper\", scene): np.geomspace(0.05, 0.5, 10)\n        for scene in [\"bike\", \"chairs\", \"fountain\"]\n    },\n    **{\n        (\"heritage\", scene): np.geomspace(0.1, 2, 10)\n        for scene in [\"cyprus\", \"dioscuri\"]\n    },\n    **{(\"heritage\", \"wall\"): np.geomspace(0.05, 1, 10)},\n    **{(\"urban\", \"kyiv-puppet-theater\"): np.geomspace(0.5, 5, 10)},\n}","metadata":{"_kg_hide-input":true,"trusted":true,"execution":{"iopub.status.busy":"2025-06-02T05:28:26.724958Z","iopub.execute_input":"2025-06-02T05:28:26.725293Z","iopub.status.idle":"2025-06-02T05:28:26.754505Z","shell.execute_reply.started":"2025-06-02T05:28:26.725268Z","shell.execute_reply":"2025-06-02T05:28:26.753691Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Run evaluation and log outputs ","metadata":{}},{"cell_type":"code","source":"%%capture cap --no-stderr\n\nimport datetime\ncurrent_time = datetime.datetime.now().strftime(\"%Y-%m-%d_%H-%M-%S\")\nprint(\"=========================================\")\nprint(current_time)\nprint(FM_PARAMS)\nprint(LOG_DICT)\nprint(MODEL_DICT)\nprint(\"Match filter ratio = \", MATCH_FILTER_RATIO)\nfor dataset in out_results:\n    for scene in out_results[dataset]:\n        print(\n            f\"Registered: {dataset} / {scene} -> {len(out_results[dataset][scene])}/{len(data_dict[dataset][scene])} images\"\n        )\n\nprint(time_dict)\n\nif MODE == \"train\":\n    with open(f\"{SRC}/train/train_labels.csv\", \"r\") as fr, open(\n        \"ground_truth.csv\", \"w\"\n    ) as fw:\n        for i, l in enumerate(fr):\n            if i == 0:\n                fw.write(\n                    \"image_path,dataset,scene,rotation_matrix,translation_vector\\n\"\n                )\n            else:\n                dataset, scene, image, R, T = l.strip().split(\",\")\n                fw.write(f\"{image},{dataset},{scene},{R},{T}\\n\")\n\n    eval_submission(\n        submission_csv_path=\"submission_train.csv\",\n        ground_truth_csv_path=\"ground_truth.csv\",\n        rotation_thresholds_degrees_dict=rotation_thresholds_degrees_dict,\n        translation_thresholds_meters_dict=translation_thresholds_meters_dict,\n        verbose=True,\n    )\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-02T05:28:26.758563Z","iopub.execute_input":"2025-06-02T05:28:26.759260Z","iopub.status.idle":"2025-06-02T05:28:28.679382Z","shell.execute_reply.started":"2025-06-02T05:28:26.759232Z","shell.execute_reply":"2025-06-02T05:28:28.678757Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Write log to a text file","metadata":{}},{"cell_type":"code","source":"print(cap)\nif is_local:\n    with open(\"log.txt\", \"a\") as f:\n        f.write(str(cap))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-02T05:28:28.680535Z","iopub.execute_input":"2025-06-02T05:28:28.680768Z","iopub.status.idle":"2025-06-02T05:28:28.685825Z","shell.execute_reply.started":"2025-06-02T05:28:28.680751Z","shell.execute_reply":"2025-06-02T05:28:28.685070Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# mast3r","metadata":{}},{"cell_type":"code","source":"# def run_mast3r_inference(image_paths, save_dir, model, device, \n#                         scenegraph_type=\"swin\", winsize=1, win_cyclic=False,\n#                         optim_level=\"refine+depth\", lr1=0.07, niter1=1000, lr2=0.014, niter2=400,\n#                         min_conf_thr=1.5, matching_conf_thr=5, shared_intrinsics=True):\n#     \"\"\"Run MASt3R inference and save results in COLMAP format\n    \n#     Args:\n#         image_paths: List of paths to input images\n#         save_dir: Directory to save results\n#         model: MASt3R model\n#         device: Device to run inference on\n#         scenegraph_type: Type of scene graph to use (\"swin\", \"logwin\", \"oneref\", \"retrieval\")\n#         winsize: Window size for scene graph\n#         win_cyclic: Whether to use cyclic window\n#         optim_level: Optimization level (\"coarse\", \"fine\", \"coarse+fine\")\n#         lr1: Learning rate for first optimization stage\n#         niter1: Number of iterations for first optimization stage\n#         lr2: Learning rate for second optimization stage\n#         niter2: Number of iterations for second optimization stage\n#         min_conf_thr: Minimum confidence threshold for points\n#         matching_conf_thr: Matching confidence threshold\n#         shared_intrinsics: Whether to use shared intrinsics\n#     \"\"\"\n#     # Load images\n#     imgs = load_images(image_paths, size=512, verbose=True)\n    \n#     # Create scene graph\n#     scene_graph_params = [scenegraph_type]\n#     if scenegraph_type in [\"swin\", \"logwin\"]:\n#         scene_graph_params.append(str(winsize))\n#         if not win_cyclic:\n#             scene_graph_params.append('noncyclic')\n#     elif scenegraph_type == \"oneref\":\n#         scene_graph_params.append(\"0\")  # Use first image as reference\n    \n#     scene_graph = '-'.join(scene_graph_params)\n    \n#     # Create image pairs\n#     pairs = make_pairs(imgs, scene_graph=scene_graph, prefilter=None, symmetrize=True)\n    \n#     # Run sparse global alignment\n#     if optim_level == 'coarse':\n#         niter2 = 0\n    \n#     scene = sparse_global_alignment(\n#         image_paths, pairs, save_dir,\n#         model, lr1=lr1, niter1=niter1, lr2=lr2, niter2=niter2, device=device,\n#         opt_depth='depth' in optim_level, shared_intrinsics=shared_intrinsics,\n#         matching_conf_thr=matching_conf_thr\n#     )\n    \n#     return scene\n\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-02T05:28:28.686791Z","iopub.execute_input":"2025-06-02T05:28:28.687302Z","iopub.status.idle":"2025-06-02T05:28:28.704022Z","shell.execute_reply.started":"2025-06-02T05:28:28.687275Z","shell.execute_reply":"2025-06-02T05:28:28.703447Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# def run_mast3r_reconstruction():\n#     \"\"\"Run MASt3R reconstruction and create submission.csv\"\"\"\n#     print(\"Running MASt3R reconstruction...\")\n    \n#     # Initialize output results dictionary\n#     out_results = defaultdict(dict)\n    \n#     # Process each scene\n#     for dataset, scene in all_scenes:\n#         print(f\"\\nProcessing {dataset}/{scene}\")\n        \n#         # Get image paths\n#         img_dir = f\"{SRC}/{MODE}/{dataset}/{scene}/images\"\n#         if not os.path.exists(img_dir):\n#             print(\"Image dir does not exist:\", img_dir)\n#             continue\n            \n#         img_fnames = [f\"{SRC}/{MODE}/{x}\" for x in data_dict[dataset][scene]]\n#         print(f\"Got {len(img_fnames)} images\")\n        \n#         # Create save directory\n#         save_dir = f\"mast3r_out/{dataset}_{scene}\"\n#         os.makedirs(save_dir, exist_ok=True)\n        \n#         try:\n#             # Run MASt3R inference\n#             scene_obj = run_mast3r_inference(\n#                 img_fnames,\n#                 save_dir,\n#                 model,\n#                 device='cuda',\n#                 scenegraph_type='swin',\n#                 winsize=2,  # 增加窗口大小\n#                 optim_level='refine+depth',\n#                 lr1=0.07,\n#                 niter1=1000,  # 增加迭代次數\n#                 lr2=0.014,   # 調整學習率\n#                 niter2=500,  # 增加迭代次數\n#                 shared_intrinsics=False,\n#                 matching_conf_thr=5.0  # 增加匹配置信度閾值\n#             )\n            \n#             c2w_poses = scene_obj.get_im_poses()\n            \n#             for i, img_path in enumerate(img_fnames):\n#                 img_name = img_path.split(\"/\")[-1]\n                \n#                 w2c = np.linalg.inv(c2w_poses[i])\n\n#                 R = w2c[:3, :3]\n#                 t = w2c[:3, 3]\n                \n#                 out_results[dataset][scene][img_name] = {\n#                     \"R\": R,  # world2cam rotation\n#                     \"t\": t   # world2cam translation\n#                 }\n                \n#             print(f\"Registered: {dataset}/{scene} -> {len(out_results[dataset][scene])}/{len(img_fnames)} images\")\n            \n#         except Exception as e:\n#             print(f\"Error processing {dataset}/{scene}:\")\n#             print(e)\n#             continue\n            \n#         gc.collect()\n#         torch.cuda.empty_cache()\n    \n#     # Create submission file\n#     create_submission(out_results, data_dict, MODE)\n    \n#     # Evaluate if in training mode\n#     if MODE == \"train\":\n#         with open(f\"{SRC}/train/train_labels.csv\", \"r\") as fr, open(\"ground_truth.csv\", \"w\") as fw:\n#             for i, l in enumerate(fr):\n#                 if i == 0:\n#                     fw.write(\n#                         \"image_path,dataset,scene,rotation_matrix,translation_vector\\n\"\n#                     )\n#                 else:\n#                     dataset, scene, image, R, T = l.strip().split(\",\")\n#                     fw.write(f\"{image},{dataset},{scene},{R},{T}\\n\")\n\n#         eval_submission(\n#             submission_csv_path=\"submission_train.csv\",\n#             ground_truth_csv_path=\"ground_truth.csv\",\n#             rotation_thresholds_degrees_dict=rotation_thresholds_degrees_dict,\n#             translation_thresholds_meters_dict=translation_thresholds_meters_dict,\n#             verbose=True\n#         )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-02T05:28:28.704779Z","iopub.execute_input":"2025-06-02T05:28:28.705001Z","iopub.status.idle":"2025-06-02T05:28:28.723648Z","shell.execute_reply.started":"2025-06-02T05:28:28.704975Z","shell.execute_reply":"2025-06-02T05:28:28.722838Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# run_mast3r_reconstruction()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-02T05:28:28.724665Z","iopub.execute_input":"2025-06-02T05:28:28.725133Z","iopub.status.idle":"2025-06-02T05:28:28.740933Z","shell.execute_reply.started":"2025-06-02T05:28:28.725105Z","shell.execute_reply":"2025-06-02T05:28:28.740365Z"}},"outputs":[],"execution_count":null}]}