{"metadata":{"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":71885,"databundleVersionId":8069805,"sourceType":"competition"},{"sourceId":7884485,"sourceType":"datasetVersion","datasetId":4628051},{"sourceId":7884725,"sourceType":"datasetVersion","datasetId":4628331},{"sourceId":4534,"sourceType":"modelInstanceVersion","modelInstanceId":3326},{"sourceId":17191,"sourceType":"modelInstanceVersion","modelInstanceId":14317},{"sourceId":17555,"sourceType":"modelInstanceVersion","modelInstanceId":14611}],"dockerImageVersionId":30674,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true},"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.10.13"},"papermill":{"default_parameters":{},"duration":1306.94114,"end_time":"2024-03-28T20:52:19.339200","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2024-03-28T20:30:32.398060","version":"2.5.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install --no-index /kaggle/input/imc2024-packages-lightglue-rerun-kornia/* --no-deps\n!mkdir -p /root/.cache/torch/hub/checkpoints\n!cp /kaggle/input/aliked/pytorch/aliked-n16/1/* /root/.cache/torch/hub/checkpoints/\n!cp /kaggle/input/lightglue/pytorch/aliked/1/* /root/.cache/torch/hub/checkpoints/\n!cp /kaggle/input/lightglue/pytorch/aliked/1/aliked_lightglue.pth /root/.cache/torch/hub/checkpoints/aliked_lightglue_v0-1_arxiv-pth","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_kg_hide-input":true,"_kg_hide-output":true,"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","papermill":{"duration":11.643329,"end_time":"2024-03-28T20:30:47.044473","exception":false,"start_time":"2024-03-28T20:30:35.401144","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-04-01T10:41:31.493654Z","iopub.execute_input":"2024-04-01T10:41:31.493986Z","iopub.status.idle":"2024-04-01T10:41:39.640650Z","shell.execute_reply.started":"2024-04-01T10:41:31.493958Z","shell.execute_reply":"2024-04-01T10:41:39.639380Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# General utilities\nimport matplotlib.pyplot as plt\n\nimport os\nfrom tqdm import tqdm\nfrom pathlib import Path\nfrom time import time, sleep\nfrom fastprogress import progress_bar\nimport gc\nimport numpy as np\nimport h5py\nfrom IPython.display import clear_output\nfrom collections import defaultdict\nfrom copy import deepcopy\nfrom typing import Any\nimport itertools\nimport pandas as pd\n\n# CV/MLe\nimport cv2\nimport torch\nfrom torch import Tensor as T\nimport torch.nn.functional as F\nimport kornia as K\nimport kornia.feature as KF\nfrom PIL import Image\nfrom transformers import AutoImageProcessor, AutoModel\n\nimport torch\nfrom lightglue import match_pair\nfrom lightglue import LightGlue, ALIKED\nfrom lightglue.utils import load_image, rbd\n\n# 3D reconstruction\nimport pycolmap\n\n# Data importing into colmap\nimport sys\nsys.path.append(\"/kaggle/input/colmap-db-import\")\n\n# Provided by organizers\nfrom database import *\nfrom h5_to_db import *\n\ndef arr_to_str(a):\n    \"\"\"Returns ;-separated string representing the input\"\"\"\n    return \";\".join([str(x) for x in a.reshape(-1)])\n\ndef load_torch_image(file_name: Path | str, device=torch.device(\"cpu\")):\n    \"\"\"Loads an image and adds batch dimension\"\"\"\n    img = K.io.load_image(file_name, K.io.ImageLoadType.RGB32, device=device)[None, ...]\n    return img\n\ndevice = K.utils.get_cuda_device_if_available(0)\nprint(device)\n\nDEBUG = len([p for p in Path(\"/kaggle/input/image-matching-challenge-2024/test/\").iterdir() if p.is_dir()]) == 2\nprint(\"DEBUG:\", DEBUG)","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"papermill":{"duration":19.706664,"end_time":"2024-03-28T20:31:06.761690","exception":false,"start_time":"2024-03-28T20:30:47.055026","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-04-01T10:41:39.642983Z","iopub.execute_input":"2024-04-01T10:41:39.643343Z","iopub.status.idle":"2024-04-01T10:41:57.950560Z","shell.execute_reply.started":"2024-04-01T10:41:39.643310Z","shell.execute_reply":"2024-04-01T10:41:57.949711Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def embed_images(\n    paths: list[Path],\n    model_name: str,\n    device: torch.device = torch.device(\"cpu\"),\n) -> T:\n    \"\"\"Computes image embeddings.\n    \n    Returns a tensor of shape [len(filenames), output_dim]\n    \"\"\"\n    processor = AutoImageProcessor.from_pretrained(model_name)\n    model = AutoModel.from_pretrained(model_name).eval().to(device)\n    \n    embeddings = []\n    \n    for i, path in tqdm(enumerate(paths), desc=\"Global descriptors\"):\n        image = load_torch_image(path)\n        \n        with torch.inference_mode():\n            inputs = processor(images=image, return_tensors=\"pt\", do_rescale=False).to(device)\n            outputs = model(**inputs) # last_hidden_state and pooled\n            \n            # Max pooling over all the hidden states but the first (starting token)\n            # To obtain a tensor of shape [1, output_dim]\n            # We normalize so that distances are computed in a better fashion later\n            embedding = F.normalize(outputs.last_hidden_state[:,1:].max(dim=1)[0], dim=-1, p=2)\n            \n        embeddings.append(embedding.detach().cpu())\n    return torch.cat(embeddings, dim=0)","metadata":{"papermill":{"duration":0.022669,"end_time":"2024-03-28T20:31:06.815770","exception":false,"start_time":"2024-03-28T20:31:06.793101","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-04-01T10:41:57.951564Z","iopub.execute_input":"2024-04-01T10:41:57.952085Z","iopub.status.idle":"2024-04-01T10:41:57.960407Z","shell.execute_reply.started":"2024-04-01T10:41:57.952061Z","shell.execute_reply":"2024-04-01T10:41:57.959422Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_pairs_exhaustive(lst: list[Any]) -> list[tuple[int, int]]:\n    \"\"\"Obtains all possible index pairs of a list\"\"\"\n    return list(itertools.combinations(range(len(lst)), 2))            \n    \ndef get_image_pairs(\n    paths: list[Path],\n    model_name: str,\n    similarity_threshold: float = 0.6,\n    tolerance: int = 1000,\n    min_matches: int = 20,\n    exhaustive_if_less: int = 20,\n    p: float = 2.0,\n    device: torch.device = torch.device(\"cpu\"),\n) -> list[tuple[int, int]]:\n    \"\"\"Obtains pairs of similar images\"\"\"\n    if len(paths) <= exhaustive_if_less:\n        return get_pairs_exhaustive(paths)\n    \n    matches = []\n    \n    # Embed images and compute distances for filtering\n    embeddings = embed_images(paths, model_name)\n    distances = torch.cdist(embeddings, embeddings, p=p)\n    \n    # Remove pairs above similarity threshold (if enough)\n    mask = distances <= similarity_threshold\n    image_indices = np.arange(len(paths))\n    \n    for current_image_index in range(len(paths)):\n        mask_row = mask[current_image_index]\n        indices_to_match = image_indices[mask_row]\n        \n        # We don't have enough matches below the threshold, we pick most similar ones\n        if len(indices_to_match) < min_matches:\n            indices_to_match = np.argsort(distances[current_image_index])[:min_matches]\n            \n        for other_image_index in indices_to_match:\n            # Skip an image matching itself\n            if other_image_index == current_image_index:\n                continue\n            \n            # We need to check if we are below a certain distance tolerance \n            # since for images that don't have enough matches, we picked\n            # the most similar ones (which could all still be very different \n            # to the image we are analyzing)\n            if distances[current_image_index, other_image_index] < tolerance:\n                # Add the pair in a sorted manner to avoid redundancy\n                matches.append(tuple(sorted((current_image_index, other_image_index.item()))))\n                \n    return sorted(list(set(matches)))","metadata":{"papermill":{"duration":0.025005,"end_time":"2024-03-28T20:31:06.851324","exception":false,"start_time":"2024-03-28T20:31:06.826319","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-04-01T10:41:57.963488Z","iopub.execute_input":"2024-04-01T10:41:57.964080Z","iopub.status.idle":"2024-04-01T10:41:57.998104Z","shell.execute_reply.started":"2024-04-01T10:41:57.964044Z","shell.execute_reply":"2024-04-01T10:41:57.997311Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if DEBUG:\n    images_list = list(Path(\"/kaggle/input/image-matching-challenge-2024/test/church/images/\").glob(\"*.png\"))[:10]\n    index_pairs = get_image_pairs(images_list, \"/kaggle/input/dinov2/pytorch/base/1\")\n    print(index_pairs)","metadata":{"papermill":{"duration":0.030083,"end_time":"2024-03-28T20:31:06.892093","exception":false,"start_time":"2024-03-28T20:31:06.862010","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-04-01T10:41:57.999212Z","iopub.execute_input":"2024-04-01T10:41:58.000141Z","iopub.status.idle":"2024-04-01T10:41:58.032543Z","shell.execute_reply.started":"2024-04-01T10:41:58.000109Z","shell.execute_reply":"2024-04-01T10:41:58.031754Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if DEBUG:\n    dtype = torch.float32 # ALIKED has issues with float16\n\n    extractor = ALIKED(\n            max_num_keypoints=4096, \n            detection_threshold=0.01, \n            resize=1024\n        ).eval().to(device, dtype)\n\n    path = images_list[0]\n    image = load_torch_image(path, device=device).to(dtype)\n    features = extractor.extract(image)\n\n    fig, ax = plt.subplots(1, 2, figsize=(10, 20))\n    ax[0].imshow(image[0, ...].permute(1,2,0).cpu())\n    ax[1].imshow(image[0, ...].permute(1,2,0).cpu())\n    ax[1].scatter(features[\"keypoints\"][0, :, 0].cpu(), features[\"keypoints\"][0, :, 1].cpu(), s=0.5, c=\"red\")\n\n    del extractor","metadata":{"papermill":{"duration":2.589737,"end_time":"2024-03-28T20:31:09.512841","exception":false,"start_time":"2024-03-28T20:31:06.923104","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-04-01T10:41:58.033634Z","iopub.execute_input":"2024-04-01T10:41:58.034096Z","iopub.status.idle":"2024-04-01T10:42:00.420494Z","shell.execute_reply.started":"2024-04-01T10:41:58.034064Z","shell.execute_reply":"2024-04-01T10:42:00.419603Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def detect_keypoints(\n    paths: list[Path],\n    feature_dir: Path,\n    num_features: int = 4096,\n    resize_to: int = 1024,\n    device: torch.device = torch.device(\"cpu\"),\n) -> None:\n    \"\"\"Detects the keypoints in a list of images with ALIKED\n    \n    Stores them in feature_dir/keypoints.h5 and feature_dir/descriptors.h5\n    to be used later with LightGlue\n    \"\"\"\n    dtype = torch.float32 # ALIKED has issues with float16\n    \n    extractor = ALIKED(\n        max_num_keypoints=num_features, \n        detection_threshold=0.01, \n        resize=resize_to\n    ).eval().to(device, dtype)\n    \n    feature_dir.mkdir(parents=True, exist_ok=True)\n    \n    with h5py.File(feature_dir / \"keypoints.h5\", mode=\"w\") as f_keypoints, \\\n         h5py.File(feature_dir / \"descriptors.h5\", mode=\"w\") as f_descriptors:\n        \n        for path in tqdm(paths, desc=\"Computing keypoints\"):\n            key = path.name\n            \n            with torch.inference_mode():\n                image = load_torch_image(path, device=device).to(dtype)\n                features = extractor.extract(image)\n                \n                f_keypoints[key] = features[\"keypoints\"].squeeze().detach().cpu().numpy()\n                f_descriptors[key] = features[\"descriptors\"].squeeze().detach().cpu().numpy()","metadata":{"papermill":{"duration":0.033659,"end_time":"2024-03-28T20:31:09.567815","exception":false,"start_time":"2024-03-28T20:31:09.534156","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-04-01T10:42:00.421694Z","iopub.execute_input":"2024-04-01T10:42:00.422002Z","iopub.status.idle":"2024-04-01T10:42:00.431351Z","shell.execute_reply.started":"2024-04-01T10:42:00.421976Z","shell.execute_reply":"2024-04-01T10:42:00.430409Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if DEBUG:\n    feature_dir = Path(\"./sample_test_features\")\n    detect_keypoints(images_list, feature_dir)","metadata":{"papermill":{"duration":29.434528,"end_time":"2024-03-28T20:31:39.022035","exception":false,"start_time":"2024-03-28T20:31:09.587507","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-04-01T10:42:00.432666Z","iopub.execute_input":"2024-04-01T10:42:00.432976Z","iopub.status.idle":"2024-04-01T10:42:28.302194Z","shell.execute_reply.started":"2024-04-01T10:42:00.432952Z","shell.execute_reply":"2024-04-01T10:42:28.301316Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if DEBUG:\n    matcher_params = {\n        \"width_confidence\": -1,\n        \"depth_confidence\": -1,\n        \"mp\": True if 'cuda' in str(device) else False,\n    }\n    matcher = KF.LightGlueMatcher(\"aliked\", matcher_params).eval().to(device)\n\n    with h5py.File(feature_dir / \"keypoints.h5\", mode=\"r\") as f_keypoints, \\\n         h5py.File(feature_dir / \"descriptors.h5\", mode=\"r\") as f_descriptors:\n            idx1, idx2 = index_pairs[0]\n            key1, key2 = images_list[idx1].name, images_list[idx2].name\n\n            keypoints1 = torch.from_numpy(f_keypoints[key1][...]).to(device)\n            keypoints2 = torch.from_numpy(f_keypoints[key2][...]).to(device)\n            print(\"Keypoints:\", keypoints1.shape, keypoints2.shape)\n            descriptors1 = torch.from_numpy(f_descriptors[key1][...]).to(device)\n            descriptors2 = torch.from_numpy(f_descriptors[key2][...]).to(device)\n            print(\"Descriptors:\", descriptors1.shape, descriptors2.shape)\n\n            with torch.inference_mode():\n                distances, indices = matcher(\n                    descriptors1, \n                    descriptors2, \n                    KF.laf_from_center_scale_ori(keypoints1[None]),\n                    KF.laf_from_center_scale_ori(keypoints2[None]),\n                )\n    print(distances, indices)","metadata":{"papermill":{"duration":0.691886,"end_time":"2024-03-28T20:31:39.775809","exception":false,"start_time":"2024-03-28T20:31:39.083923","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-04-01T10:42:28.303320Z","iopub.execute_input":"2024-04-01T10:42:28.303590Z","iopub.status.idle":"2024-04-01T10:42:28.938265Z","shell.execute_reply.started":"2024-04-01T10:42:28.303567Z","shell.execute_reply":"2024-04-01T10:42:28.937280Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def keypoint_distances(\n    paths: list[Path],\n    index_pairs: list[tuple[int, int]],\n    feature_dir: Path,\n    min_matches: int = 15,\n    verbose: bool = True,\n    device: torch.device = torch.device(\"cpu\"),\n) -> None:\n    \"\"\"Computes distances between keypoints of images.\n    \n    Stores output at feature_dir/matches.h5\n    \"\"\"\n    \n    matcher_params = {\n        \"width_confidence\": -1,\n        \"depth_confidence\": -1,\n        \"mp\": True if 'cuda' in str(device) else False,\n    }\n    matcher = KF.LightGlueMatcher(\"aliked\", matcher_params).eval().to(device)\n    \n    with h5py.File(feature_dir / \"keypoints.h5\", mode=\"r\") as f_keypoints, \\\n         h5py.File(feature_dir / \"descriptors.h5\", mode=\"r\") as f_descriptors, \\\n         h5py.File(feature_dir / \"matches.h5\", mode=\"w\") as f_matches:\n        \n            for idx1, idx2 in tqdm(index_pairs, desc=\"Computing keypoing distances\"):\n                key1, key2 = paths[idx1].name, paths[idx2].name\n\n                keypoints1 = torch.from_numpy(f_keypoints[key1][...]).to(device)\n                keypoints2 = torch.from_numpy(f_keypoints[key2][...]).to(device)\n                descriptors1 = torch.from_numpy(f_descriptors[key1][...]).to(device)\n                descriptors2 = torch.from_numpy(f_descriptors[key2][...]).to(device)\n\n                with torch.inference_mode():\n                    distances, indices = matcher(\n                        descriptors1, \n                        descriptors2, \n                        KF.laf_from_center_scale_ori(keypoints1[None]),\n                        KF.laf_from_center_scale_ori(keypoints2[None]),\n                    )\n\n                # We have matches to consider\n                n_matches = len(indices)\n                if n_matches:\n                    if verbose:\n                        print(f\"{key1}-{key2}: {n_matches} matches\")\n                    # Store the matches in the group of one image\n                    if n_matches >= min_matches:\n                        group  = f_matches.require_group(key1)\n                        group.create_dataset(key2, data=indices.detach().cpu().numpy().reshape(-1, 2))","metadata":{"papermill":{"duration":0.039144,"end_time":"2024-03-28T20:31:39.836992","exception":false,"start_time":"2024-03-28T20:31:39.797848","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-04-01T10:42:28.941358Z","iopub.execute_input":"2024-04-01T10:42:28.941664Z","iopub.status.idle":"2024-04-01T10:42:28.954683Z","shell.execute_reply.started":"2024-04-01T10:42:28.941623Z","shell.execute_reply":"2024-04-01T10:42:28.953708Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if DEBUG:\n    keypoint_distances(images_list, index_pairs, feature_dir, verbose=False)","metadata":{"papermill":{"duration":808.312267,"end_time":"2024-03-28T20:45:08.171096","exception":false,"start_time":"2024-03-28T20:31:39.858829","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-04-01T10:42:28.955787Z","iopub.execute_input":"2024-04-01T10:42:28.956099Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def import_into_colmap(\n    path: Path,\n    feature_dir: Path,\n    database_path: str = \"colmap.db\",\n) -> None:\n    \"\"\"Adds keypoints into colmap\"\"\"\n    db = COLMAPDatabase.connect(database_path)\n    db.create_tables()\n    single_camera = False\n    fname_to_id = add_keypoints(db, feature_dir, path, \"\", \"simple-pinhole\", single_camera)\n    add_matches(\n        db,\n        feature_dir,\n        fname_to_id,\n    )\n    db.commit()","metadata":{"papermill":{"duration":0.034466,"end_time":"2024-03-28T20:45:08.281211","exception":false,"start_time":"2024-03-28T20:45:08.246745","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if DEBUG:\n    database_path = \"colmap.db\"\n    images_dir = images_list[0].parent\n    import_into_colmap(\n        images_dir, \n        feature_dir, \n        database_path,\n    )\n\n    # This does RANSAC\n    pycolmap.match_exhaustive(database_path)","metadata":{"papermill":{"duration":1.406543,"end_time":"2024-03-28T20:45:09.763541","exception":false,"start_time":"2024-03-28T20:45:08.356998","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if DEBUG:\n    mapper_options = pycolmap.IncrementalPipelineOptions()\n    mapper_options.min_model_size = 3\n    mapper_options.max_num_models = 2\n\n    maps = pycolmap.incremental_mapping(\n        database_path=database_path, \n        image_path=images_dir,\n        output_path=Path.cwd() / \"incremental_pipeline_outputs\", \n        options=mapper_options,\n    )","metadata":{"_kg_hide-output":true,"papermill":{"duration":2.81272,"end_time":"2024-03-28T20:45:12.653308","exception":false,"start_time":"2024-03-28T20:45:09.840588","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if DEBUG:\n    print(maps[0].summary())\n    for k, im in maps[0].images.items():\n        print(\"Rotation\", im.cam_from_world.rotation.matrix(), \"Translation:\", im.cam_from_world.translation, sep=\"\\n\")\n        print()","metadata":{"papermill":{"duration":0.040927,"end_time":"2024-03-28T20:45:12.723814","exception":false,"start_time":"2024-03-28T20:45:12.682887","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Running everything","metadata":{"papermill":{"duration":0.028391,"end_time":"2024-03-28T20:45:12.782751","exception":false,"start_time":"2024-03-28T20:45:12.754360","status":"completed"},"tags":[]}},{"cell_type":"code","source":"def parse_sample_submission(\n    base_path: Path,\n) -> dict[dict[str, list[Path]]]:\n    \"\"\"Construct a dict describing the test data as \n    \n    {\"dataset\": {\"scene\": [<image paths>]}}\n    \"\"\"\n    data_dict = {}\n    with open(base_path / \"sample_submission.csv\", \"r\") as f:\n        for i, l in enumerate(f):\n            # Skip header\n            if i == 0:\n                print(\"header:\", l)\n\n            if l and i > 0:\n                image_path, 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(Path(base_path / image_path))\n\n    for dataset in data_dict:\n        for scene in data_dict[dataset]:\n            print(f\"{dataset} / {scene} -> {len(data_dict[dataset][scene])} images\")\n\n    return data_dict","metadata":{"_kg_hide-input":true,"papermill":{"duration":0.041459,"end_time":"2024-03-28T20:45:12.853684","exception":false,"start_time":"2024-03-28T20:45:12.812225","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def create_submission(\n    results: dict,\n    data_dict: dict[dict[str, list[Path]]],\n    base_path: Path,\n) -> None:\n    \"\"\"Prepares a submission file.\"\"\"\n    \n    with open(\"submission.csv\", \"w\") as f:\n        f.write(\"image_path,dataset,scene,rotation_matrix,translation_vector\\n\")\n        \n        for dataset in data_dict:\n            # Only write results for datasets with images that have results \n            if dataset in results:\n                res = results[dataset]\n            else:\n                res = {}\n            \n            # Same for scenes\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                    \n                # Write the row with rotation and translation matrices\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                    image_path = str(image.relative_to(base_path))\n                    f.write(f\"{image_path},{dataset},{scene},{arr_to_str(R)},{arr_to_str(T)}\\n\")","metadata":{"_kg_hide-input":true,"papermill":{"duration":0.042091,"end_time":"2024-03-28T20:45:12.925424","exception":false,"start_time":"2024-03-28T20:45:12.883333","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Config:\n    base_path: Path = Path(\"/kaggle/input/image-matching-challenge-2024\")\n    feature_dir: Path = Path.cwd() / \".feature_outputs\"\n        \n    device: torch.device = K.utils.get_cuda_device_if_available(0)\n    \n    pair_matching_args = {\n        \"model_name\": \"/kaggle/input/dinov2/pytorch/base/1\",\n        \"similarity_threshold\": 0.3,\n        \"tolerance\": 500,\n        \"min_matches\": 50,\n        \"exhaustive_if_less\": 50,\n        \"p\": 2.0,\n    }\n    \n    keypoint_detection_args = {\n        \"num_features\": 4096,\n        \"resize_to\": 1024,\n    }\n    \n    keypoint_distances_args = {\n        \"min_matches\": 15,\n        \"verbose\": False,\n    }\n    \n    colmap_mapper_options = {\n        \"min_model_size\": 3, # By default colmap does not generate a reconstruction if less than 10 images are registered. Lower it to 3.\n        \"max_num_models\": 2,\n    }","metadata":{"papermill":{"duration":0.04094,"end_time":"2024-03-28T20:45:12.995625","exception":false,"start_time":"2024-03-28T20:45:12.954685","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def run_from_config(config: Config) -> None:\n    results = {}\n    \n    data_dict = parse_sample_submission(config.base_path)\n    datasets = list(data_dict.keys())\n    \n    for dataset in datasets:\n        if dataset not in results:\n            results[dataset] = {}\n            \n        for scene in data_dict[dataset]:\n            images_dir = data_dict[dataset][scene][0].parent\n            results[dataset][scene] = {}\n            image_paths = data_dict[dataset][scene]\n            print (f\"Got {len(image_paths)} images\")\n            \n            try:\n                feature_dir = config.feature_dir / f\"{dataset}_{scene}\"\n                feature_dir.mkdir(parents=True, exist_ok=True)\n                database_path = feature_dir / \"colmap.db\"\n                if database_path.exists():\n                    database_path.unlink()\n                \n                # 1. Get the pairs of images that are somewhat similar\n                index_pairs = get_image_pairs(\n                    image_paths,\n                    **config.pair_matching_args,\n                    device=config.device,\n                )\n                gc.collect()\n                \n                # 2. Detect keypoints of all images\n                detect_keypoints(\n                    image_paths,\n                    feature_dir,\n                    **config.keypoint_detection_args,\n                    device=device,\n                )\n                gc.collect()\n                \n                # 3. Match  keypoints of pairs of similar images\n                keypoint_distances(\n                    image_paths, \n                    index_pairs, \n                    feature_dir,\n                    **config.keypoint_distances_args,\n                    device=device,\n                )\n                gc.collect()\n                \n                sleep(1)\n                \n                # 4.1. Import keypoint distances of matches into colmap for RANSAC \n                import_into_colmap(\n                    images_dir, \n                    feature_dir, \n                    database_path,\n                )\n                \n                output_path = feature_dir / \"colmap_rec_aliked\"\n                output_path.mkdir(parents=True, exist_ok=True)\n                \n                # 4.2. Compute RANSAC (detect match outliers)\n                # By doing it exhaustively we guarantee we will find the best possible configuration\n                pycolmap.match_exhaustive(database_path)\n                \n                mapper_options = pycolmap.IncrementalPipelineOptions(**config.colmap_mapper_options)\n                \n                # 5.1 Incrementally start reconstructing the scene (sparse reconstruction)\n                # The process starts from a random pair of images and is incrementally extended by \n                # registering new images and triangulating new points.\n                maps = pycolmap.incremental_mapping(\n                    database_path=database_path, \n                    image_path=images_dir,\n                    output_path=output_path, \n                    options=mapper_options,\n                )\n                \n                print(maps)\n                clear_output(wait=False)\n                \n                # 5.2. Look for the best reconstruction: The incremental mapping offered by \n                # pycolmap attempts to reconstruct multiple models, we must pick the best one\n                images_registered  = 0\n                best_idx = None\n                \n                print (\"Looking for the best reconstruction\")\n            \n                if isinstance(maps, dict):\n                    for idx1, rec in maps.items():\n                        print(idx1, rec.summary())\n                        try:\n                            if len(rec.images) > images_registered:\n                                images_registered = len(rec.images)\n                                best_idx = idx1\n                        except Exception:\n                            continue\n                \n                # Parse the reconstruction object to get the rotation matrix and translation vector\n                # obtained for each image in the reconstruction\n                if best_idx is not None:\n                    for k, im in maps[best_idx].images.items():\n                        key = config.base_path / \"test\" / scene / \"images\" / im.name\n                        results[dataset][scene][key] = {}\n                        results[dataset][scene][key][\"R\"] = deepcopy(im.cam_from_world.rotation.matrix())\n                        results[dataset][scene][key][\"t\"] = deepcopy(np.array(im.cam_from_world.translation))\n                        \n                print(f\"Registered: {dataset} / {scene} -> {len(results[dataset][scene])} images\")\n                print(f\"Total: {dataset} / {scene} -> {len(data_dict[dataset][scene])} images\")\n                create_submission(results, data_dict, config.base_path)\n                gc.collect()\n            \n            except Exception as e:\n                print(e)","metadata":{"papermill":{"duration":0.052701,"end_time":"2024-03-28T20:45:13.077146","exception":false,"start_time":"2024-03-28T20:45:13.024445","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"run_from_config(Config)","metadata":{"papermill":{"duration":421.77681,"end_time":"2024-03-28T20:52:14.883079","exception":false,"start_time":"2024-03-28T20:45:13.106269","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!cat submission.csv","metadata":{"papermill":{"duration":1.027393,"end_time":"2024-03-28T20:52:15.940515","exception":false,"start_time":"2024-03-28T20:52:14.913122","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]}]}