{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":71885,"databundleVersionId":8143495,"sourceType":"competition"},{"sourceId":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":30665,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"#
📸 Image Matching Challenge - 📊 Understanding the baseline
\n

Reconstruct 3D scenes from 2D images over six different domains

\n\n***\n\nIn this notebook I explain the baseline solution provided by the organizers in [this notebook](https://www.kaggle.com/code/oldufo/imc-2024-submission-example).\n\nI have made the code a bit easier to read, adding comments and type annotations to make it easier for you to understand what is going on.\n\nHope you enjoy ❤️","metadata":{}},{"cell_type":"markdown","source":"# Structure from Motion","metadata":{}},{"cell_type":"markdown","source":"Structure from Motion (SfM) is the name given to the procedure of **reconstructing a 3D scene and simultaneously obtaining the camera poses of a camera w.r.t. the given scene**. This means that, as the name suggests, we are creating the entire rigid structure from a set of images with different view points (or equivalently a camera in motion).\n\nIn this competition, the important aspect of SfM we are interested in is *obtaining the camera poses* of where each image was taken, described by a rotation matrix and translation vector from the origin. These are the objects that will be scored in our submission!","metadata":{}},{"cell_type":"markdown","source":"
\"Structure
","metadata":{}},{"cell_type":"markdown","source":"# Baseline solution steps\nIn order to be able to estimate the camera poses, the solution provided by the organizers consists in the following steps:\n\n* [1. Find pairs of images that are similar](#1)\n* [2. Compute image keypoints](#2)\n* [3. Match keypoints between images](#3)\n* [4. Outlier detection with RANSAC](#4)\n* [5. Sparse reconstruction](#5)\n\nLet's understand how these steps are carried out","metadata":{}},{"cell_type":"markdown","source":"# Installing & importing relevant packages and models","metadata":{}},{"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":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2024-03-27T20:44:43.769283Z","iopub.execute_input":"2024-03-27T20:44:43.770095Z","iopub.status.idle":"2024-03-27T20:44:54.026101Z","shell.execute_reply.started":"2024-03-27T20:44:43.770064Z","shell.execute_reply":"2024-03-27T20:44:54.025Z"},"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,"execution":{"iopub.status.busy":"2024-03-27T20:44:54.028186Z","iopub.execute_input":"2024-03-27T20:44:54.0285Z","iopub.status.idle":"2024-03-27T20:45:12.12253Z","shell.execute_reply.started":"2024-03-27T20:44:54.02847Z","shell.execute_reply":"2024-03-27T20:45:12.121485Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"\n# Finding image pairs\n\nTo find pairs of similar images, we use [DINOv2](https://arxiv.org/pdf/2304.07193.pdf) to obtain normalized image embeddings.\n\n
\"DINOv2
\nThen, we calculate the distances between all the embeddings, and only keep those below a given distance threshold. For images with less than a set minimum number of pairs, the closest ones are kept instead.","metadata":{}},{"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":{"execution":{"iopub.status.busy":"2024-03-27T19:26:19.965897Z","iopub.execute_input":"2024-03-27T19:26:19.966233Z","iopub.status.idle":"2024-03-27T19:26:20.278034Z","shell.execute_reply.started":"2024-03-27T19:26:19.9662Z","shell.execute_reply":"2024-03-27T19:26:20.276792Z"},"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":{"execution":{"iopub.status.busy":"2024-03-27T19:26:20.279657Z","iopub.execute_input":"2024-03-27T19:26:20.280099Z","iopub.status.idle":"2024-03-27T19:26:20.295316Z","shell.execute_reply.started":"2024-03-27T19:26:20.280059Z","shell.execute_reply":"2024-03-27T19:26:20.294217Z"},"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":{"execution":{"iopub.status.busy":"2024-03-27T19:26:20.297963Z","iopub.execute_input":"2024-03-27T19:26:20.298297Z","iopub.status.idle":"2024-03-27T19:26:20.311085Z","shell.execute_reply.started":"2024-03-27T19:26:20.298262Z","shell.execute_reply":"2024-03-27T19:26:20.310156Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"\n# Computing keypoints\n\nIn order to be able to know the position of each camera, we must be able to relate images to each other. For this, we extract relevant keypoints and compare pairs of image keypoints against each other. There are many ways to extract relevant keypoints, the most traditional one being [SIFT](https://en.wikipedia.org/wiki/Scale-invariant_feature_transform). However, newer and improved methods exist now, one of which is [ALIKED](https://arxiv.org/abs/2304.03608), the keypoint extraction method used in the solution.\n\n
\"ALIKED
\n\n\nLet's take a closer look at the keypoints that ALIKED extracts.","metadata":{}},{"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.025, \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":{"execution":{"iopub.status.busy":"2024-03-27T19:26:20.312401Z","iopub.execute_input":"2024-03-27T19:26:20.312775Z","iopub.status.idle":"2024-03-27T19:26:20.323257Z","shell.execute_reply.started":"2024-03-27T19:26:20.312742Z","shell.execute_reply":"2024-03-27T19:26:20.322392Z"},"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":{"execution":{"iopub.status.busy":"2024-03-27T19:26:20.324258Z","iopub.execute_input":"2024-03-27T19:26:20.324524Z","iopub.status.idle":"2024-03-27T19:26:20.338576Z","shell.execute_reply.started":"2024-03-27T19:26:20.3245Z","shell.execute_reply":"2024-03-27T19:26:20.337641Z"},"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":{"execution":{"iopub.status.busy":"2024-03-27T19:26:20.339756Z","iopub.execute_input":"2024-03-27T19:26:20.340111Z","iopub.status.idle":"2024-03-27T19:26:20.34904Z","shell.execute_reply.started":"2024-03-27T19:26:20.340086Z","shell.execute_reply":"2024-03-27T19:26:20.3482Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"\n# Match and compute keypoint distances\n\nNow that we have the relevant image pairs and keypoints, we can go ahead and compare the keypoints of the images in a pair to find a good relationship between them. This is done with [LightGlue](https://arxiv.org/abs/2306.13643), which matches the keypoints and their descriptors between two images.\n\n
\"LightGlue
","metadata":{}},{"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":{"execution":{"iopub.status.busy":"2024-03-27T19:26:20.350103Z","iopub.execute_input":"2024-03-27T19:26:20.350361Z","iopub.status.idle":"2024-03-27T19:26:20.363754Z","shell.execute_reply.started":"2024-03-27T19:26:20.350339Z","shell.execute_reply":"2024-03-27T19:26:20.362738Z"},"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":{"execution":{"iopub.status.busy":"2024-03-27T19:26:20.364842Z","iopub.execute_input":"2024-03-27T19:26:20.365126Z","iopub.status.idle":"2024-03-27T19:26:20.378255Z","shell.execute_reply.started":"2024-03-27T19:26:20.365102Z","shell.execute_reply":"2024-03-27T19:26:20.377301Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if DEBUG:\n keypoint_distances(images_list, index_pairs, feature_dir, verbose=False)","metadata":{"execution":{"iopub.status.busy":"2024-03-27T19:26:20.379532Z","iopub.execute_input":"2024-03-27T19:26:20.379881Z","iopub.status.idle":"2024-03-27T19:26:20.391846Z","shell.execute_reply.started":"2024-03-27T19:26:20.379837Z","shell.execute_reply":"2024-03-27T19:26:20.390916Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"\n# RANSAC\nUp to now, we have matched keypoints and their descriptors extracted from pairs of images. This is described by a [fundamental matrix](https://en.wikipedia.org/wiki/Fundamental_matrix_(computer_vision)) denoted as $F$. In epipolar geometry, with homogeneous image coordinates, $x$ and $x′$, of corresponding points in a stereo image pair, $Fx$ describes a line (an epipolar line) on which the corresponding point $x′$ on the other image must lie. That means, for all pairs of corresponding points, $x'Fx = 0$ holds. This is known as epipolar constraint or correspondance condition (or Longuet-Higgins equation), and is solved via the [eight-point algorithm](https://en.wikipedia.org/wiki/Eight-point_algorithm).\n\n
\"Fundamental
\n\nSince the keypoint correspondences are computed using feature descriptors, the data is bound to be noisy and (in general) contains several outliers. Thus, to remove these outliers, we use a [RANSAC](https://en.wikipedia.org/wiki/Random_sample_consensus) algorithm to find the best possible fundamental matrix. So, out of all possibilities, the $F$ matrix with maximum number of inliers is chosen.\n\n
\"RANSAC\"
","metadata":{}},{"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":{"execution":{"iopub.status.busy":"2024-03-27T19:26:20.393044Z","iopub.execute_input":"2024-03-27T19:26:20.393344Z","iopub.status.idle":"2024-03-27T19:26:20.402371Z","shell.execute_reply.started":"2024-03-27T19:26:20.393321Z","shell.execute_reply":"2024-03-27T19:26:20.401664Z"},"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":{"execution":{"iopub.status.busy":"2024-03-27T19:26:20.40346Z","iopub.execute_input":"2024-03-27T19:26:20.403745Z","iopub.status.idle":"2024-03-27T19:26:20.415434Z","shell.execute_reply.started":"2024-03-27T19:26:20.403721Z","shell.execute_reply":"2024-03-27T19:26:20.414517Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"\n# Sparse Reconstruction\n\nNow we have similar image pairs, with matched keypoint descriptors, without outliers! All that is left is to construct the scene and obtain the camera positions. We do this with pycolmap, which offers an incremental reconstruction algorithm that starts from two pairs of images and continually adds more and more images to the scene, resulting in a reconstructed scene with camera information. We can then use the camera rotation and translation as our submission!","metadata":{}},{"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,"execution":{"iopub.status.busy":"2024-03-27T19:26:20.418947Z","iopub.execute_input":"2024-03-27T19:26:20.419258Z","iopub.status.idle":"2024-03-27T19:26:20.426206Z","shell.execute_reply.started":"2024-03-27T19:26:20.419234Z","shell.execute_reply":"2024-03-27T19:26:20.425348Z"},"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":{"execution":{"iopub.status.busy":"2024-03-27T19:26:20.427189Z","iopub.execute_input":"2024-03-27T19:26:20.427423Z","iopub.status.idle":"2024-03-27T19:26:20.437259Z","shell.execute_reply.started":"2024-03-27T19:26:20.427402Z","shell.execute_reply":"2024-03-27T19:26:20.436438Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Running everything","metadata":{}},{"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\": []}}\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,"execution":{"iopub.status.busy":"2024-03-27T19:26:20.438391Z","iopub.execute_input":"2024-03-27T19:26:20.438648Z","iopub.status.idle":"2024-03-27T19:26:20.44954Z","shell.execute_reply.started":"2024-03-27T19:26:20.438626Z","shell.execute_reply":"2024-03-27T19:26:20.44869Z"},"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,"execution":{"iopub.status.busy":"2024-03-27T19:26:20.450681Z","iopub.execute_input":"2024-03-27T19:26:20.450974Z","iopub.status.idle":"2024-03-27T19:26:20.464365Z","shell.execute_reply.started":"2024-03-27T19:26:20.450951Z","shell.execute_reply":"2024-03-27T19:26:20.463622Z"},"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":{"execution":{"iopub.status.busy":"2024-03-27T19:26:20.465376Z","iopub.execute_input":"2024-03-27T19:26:20.466245Z","iopub.status.idle":"2024-03-27T19:26:20.478884Z","shell.execute_reply.started":"2024-03-27T19:26:20.466214Z","shell.execute_reply":"2024-03-27T19:26:20.478113Z"},"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":{"execution":{"iopub.status.busy":"2024-03-27T19:26:20.480193Z","iopub.execute_input":"2024-03-27T19:26:20.480514Z","iopub.status.idle":"2024-03-27T19:26:20.499335Z","shell.execute_reply.started":"2024-03-27T19:26:20.480486Z","shell.execute_reply":"2024-03-27T19:26:20.498449Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"run_from_config(Config)","metadata":{"execution":{"iopub.status.busy":"2024-03-27T19:26:20.500314Z","iopub.execute_input":"2024-03-27T19:26:20.500575Z","iopub.status.idle":"2024-03-27T19:31:55.75141Z","shell.execute_reply.started":"2024-03-27T19:26:20.500542Z","shell.execute_reply":"2024-03-27T19:31:55.750638Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!cat submission.csv","metadata":{"execution":{"iopub.status.busy":"2024-03-27T19:31:55.752732Z","iopub.execute_input":"2024-03-27T19:31:55.753102Z","iopub.status.idle":"2024-03-27T19:31:56.760105Z","shell.execute_reply.started":"2024-03-27T19:31:55.753069Z","shell.execute_reply":"2024-03-27T19:31:56.758889Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# What to do next?\n\nHere are some ways in which you can explore potential improvements:\n\n- Using a different image embedding model to obtain the image pairs\n- Trying other approaches for keypoint extraction, such as SIFT or DISK\n- Leveraging the training data to train a better models for each dataset","metadata":{}}]}