{"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":8367624,"sourceType":"datasetVersion","datasetId":4974117},{"sourceId":8367334,"sourceType":"datasetVersion","datasetId":4973903},{"sourceId":176805602,"sourceType":"kernelVersion"},{"sourceId":4534,"sourceType":"modelInstanceVersion","modelInstanceId":3326},{"sourceId":17191,"sourceType":"modelInstanceVersion","modelInstanceId":14317},{"sourceId":17555,"sourceType":"modelInstanceVersion","modelInstanceId":14611},{"sourceId":45794,"sourceType":"modelInstanceVersion","modelInstanceId":38394}],"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\n\n##############ADDED######################\n!pip install --no-index --find-links=/kaggle/input/mediapy/mediapy mediapy\n!pip install --no-index /kaggle/input/check-orientation/check-orientation/iglovikov_helper_functions-0.0.53-py2.py3-none-any.whl --no-deps\n#######################################","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2024-05-10T09:19:40.500147Z","iopub.execute_input":"2024-05-10T09:19:40.500848Z","iopub.status.idle":"2024-05-10T09:20:05.909502Z","shell.execute_reply.started":"2024-05-10T09:19:40.500799Z","shell.execute_reply":"2024-05-10T09:20:05.908463Z"},"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#########ADDED###########\nimport mediapy as media\nimport cv2\nfrom glob import glob\nfrom pprint import pprint\nimport shutil\n\nfrom collections import namedtuple\nfrom typing import Optional\nfrom iglovikov_helper_functions.dl.pytorch.utils import rename_layers\nfrom timm import create_model as timm_create_model\nfrom torch import nn\n###################\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\nDEBUG = False\nprint(\"DEBUG:\", DEBUG)","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2024-05-10T09:20:05.911453Z","iopub.execute_input":"2024-05-10T09:20:05.911783Z","iopub.status.idle":"2024-05-10T09:20:23.121000Z","shell.execute_reply.started":"2024-05-10T09:20:05.911754Z","shell.execute_reply":"2024-05-10T09:20:23.119978Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"JEANS = True","metadata":{"execution":{"iopub.status.busy":"2024-05-10T09:20:23.121948Z","iopub.execute_input":"2024-05-10T09:20:23.122496Z","iopub.status.idle":"2024-05-10T09:20:23.126846Z","shell.execute_reply.started":"2024-05-10T09:20:23.122470Z","shell.execute_reply":"2024-05-10T09:20:23.125762Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Import Check Orientation model\nInput an image \n\nOutput a prediction in 4 rotation class [0deg 90deg 180deg 270deg]\n\nusing the MODELS at /kaggle/input/swsl_resnext50_32x4d/pytorch/v1/1","metadata":{}},{"cell_type":"code","source":"def create_model(model_name: str, activation: Optional[str] = \"softmax\") -> nn.Module:\n # Define the 'model' namedtuple\n model = namedtuple(\"model\", [\"url\", \"model\"])\n\n # Define the models dictionary\n models = {\n \"swsl_resnext50_32x4d\": model(\n model=timm_create_model(\"swsl_resnext50_32x4d\", pretrained=False, num_classes=4),\n url=\"/kaggle/input/swsl_resnext50_32x4d/pytorch/v1/1/2020-11-16_resnext50_32x4d.pth\",\n ),\n }\n\n # Load the model from models dictionary\n model = models[model_name].model\n \n # Load the model state dictionary from local path\n state_dict = torch.load(models[model_name].url, map_location=\"cpu\")[\"state_dict\"]\n state_dict = rename_layers(state_dict, {\"model.\": \"\"})\n # Load the state dictionary into the model\n model.load_state_dict(state_dict)\n\n if activation == \"softmax\":\n return nn.Sequential(model, nn.Softmax(dim=1))\n\n return model","metadata":{"execution":{"iopub.status.busy":"2024-05-10T09:20:23.129456Z","iopub.execute_input":"2024-05-10T09:20:23.129767Z","iopub.status.idle":"2024-05-10T09:20:23.270866Z","shell.execute_reply.started":"2024-05-10T09:20:23.129733Z","shell.execute_reply":"2024-05-10T09:20:23.269733Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"orientation_model = create_model(\"swsl_resnext50_32x4d\")\norientation_model.eval()","metadata":{"execution":{"iopub.status.busy":"2024-05-10T09:20:23.272067Z","iopub.execute_input":"2024-05-10T09:20:23.272927Z","iopub.status.idle":"2024-05-10T09:20:25.061996Z","shell.execute_reply.started":"2024-05-10T09:20:23.272889Z","shell.execute_reply":"2024-05-10T09:20:25.061082Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def rotate_image(image, degrees):\n \"\"\"\n Rotate a Torch tensor representing an image by a specified angle along dimensions 2 and 3.\n \n Args:\n image (torch.Tensor): Input image tensor.\n degrees (int): Rotation angle in degrees (90, 180, or 270).\n \n Returns:\n torch.Tensor: Rotated image tensor.\n \"\"\"\n if degrees not in [90, 180, 270]:\n raise ValueError(\"Rotation angle must be 90, 180, or 270 degrees.\")\n \n # Calculate the number of times to rotate the tensor by 90 degrees\n k = degrees // 90\n \n # Rotate the image tensor along dimensions 2 and 3\n rotated_image = torch.rot90(image, k=k, dims=(2, 3))\n \n return rotated_image","metadata":{"execution":{"iopub.status.busy":"2024-05-10T09:28:00.960197Z","iopub.execute_input":"2024-05-10T09:28:00.960962Z","iopub.status.idle":"2024-05-10T09:28:00.967596Z","shell.execute_reply.started":"2024-05-10T09:28:00.960928Z","shell.execute_reply":"2024-05-10T09:28:00.966576Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def correct_orientation(image,path):\n \n if DEBUG:\n plt.imshow(image.squeeze(0).permute(1, 2, 0).numpy())\n plt.axis('off') # Hide axis\n plt.show()\n \n resized_image = F.interpolate(image, size=(224, 224), mode='bilinear', align_corners=False)\n \n ####################\n with torch.no_grad():\n prediction = orientation_model(resized_image).numpy()[0]\n \n predictedRotationDeg = 360 - (prediction.argmax() * 90)\n \n if predictedRotationDeg == 360:\n print(f'{path}`s orientation is correct.')\n else:\n print(f'{path}`s ROTATED.')\n image = rotate_image(image, predictedRotationDeg)\n ####################\n \n if DEBUG:\n plt.imshow(image.squeeze(0).permute(1, 2, 0).numpy())\n plt.axis('off') # Hide axis\n plt.show()\n \n return image","metadata":{"execution":{"iopub.status.busy":"2024-05-10T09:28:01.101318Z","iopub.execute_input":"2024-05-10T09:28:01.101688Z","iopub.status.idle":"2024-05-10T09:28:01.110324Z","shell.execute_reply.started":"2024-05-10T09:28:01.101658Z","shell.execute_reply":"2024-05-10T09:28:01.109351Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Rotation Correction\nIn this section, we correct the rotation of the images.","metadata":{}},{"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 print(image.shape)\n ########################################################################################################################\n image = correct_orientation(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-05-10T09:28:02.214406Z","iopub.execute_input":"2024-05-10T09:28:02.215367Z","iopub.status.idle":"2024-05-10T09:28:02.226408Z","shell.execute_reply.started":"2024-05-10T09:28:02.215328Z","shell.execute_reply":"2024-05-10T09:28:02.225240Z"},"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 print(\"Do exhaustive pair\")\n return get_pairs_exhaustive(paths)\n \n print(f'Use image embed dist pair with similarity_threshold = {similarity_threshold}')\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-05-10T09:28:02.671058Z","iopub.execute_input":"2024-05-10T09:28:02.671792Z","iopub.status.idle":"2024-05-10T09:28:02.682975Z","shell.execute_reply.started":"2024-05-10T09:28:02.671753Z","shell.execute_reply":"2024-05-10T09:28:02.681873Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# if JEANS:\n# images_list = list(Path(\"/kaggle/input/image-matching-challenge-2024/test/church/images/\").glob(\"*.png\"))[:30]\n# index_pairs = get_image_pairs(images_list, \"/kaggle/input/dinov2/pytorch/base/1\")\n# print(index_pairs)\n\nif DEBUG:\n images_list = list(Path(\"/kaggle/input/image-matching-challenge-2024/train/dioscuri/images\").glob(\"*.png\"))[:30]\n index_pairs = get_image_pairs(images_list, \"/kaggle/input/dinov2/pytorch/base/1\")\n print(index_pairs)","metadata":{"execution":{"iopub.status.busy":"2024-05-10T09:28:03.955061Z","iopub.execute_input":"2024-05-10T09:28:03.955403Z","iopub.status.idle":"2024-05-10T09:28:03.960564Z","shell.execute_reply.started":"2024-05-10T09:28:03.955379Z","shell.execute_reply":"2024-05-10T09:28:03.959583Z"},"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.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":{"execution":{"iopub.status.busy":"2024-05-10T09:20:25.122607Z","iopub.execute_input":"2024-05-10T09:20:25.122880Z","iopub.status.idle":"2024-05-10T09:20:25.131413Z","shell.execute_reply.started":"2024-05-10T09:20:25.122855Z","shell.execute_reply":"2024-05-10T09:20:25.130604Z"},"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-05-10T09:20:25.132402Z","iopub.execute_input":"2024-05-10T09:20:25.132670Z","iopub.status.idle":"2024-05-10T09:20:25.143666Z","shell.execute_reply.started":"2024-05-10T09:20:25.132648Z","shell.execute_reply":"2024-05-10T09:20:25.142859Z"},"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-05-10T09:20:25.144766Z","iopub.execute_input":"2024-05-10T09:20:25.145150Z","iopub.status.idle":"2024-05-10T09:20:25.152526Z","shell.execute_reply.started":"2024-05-10T09:20:25.145128Z","shell.execute_reply":"2024-05-10T09:20:25.151781Z"},"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-05-10T09:20:25.155433Z","iopub.execute_input":"2024-05-10T09:20:25.156191Z","iopub.status.idle":"2024-05-10T09:20:25.164616Z","shell.execute_reply.started":"2024-05-10T09:20:25.156168Z","shell.execute_reply":"2024-05-10T09:20:25.163774Z"},"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-05-10T09:20:25.165671Z","iopub.execute_input":"2024-05-10T09:20:25.166005Z","iopub.status.idle":"2024-05-10T09:20:25.177818Z","shell.execute_reply.started":"2024-05-10T09:20:25.165974Z","shell.execute_reply":"2024-05-10T09:20:25.177091Z"},"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-05-10T09:20:25.178951Z","iopub.execute_input":"2024-05-10T09:20:25.179199Z","iopub.status.idle":"2024-05-10T09:20:25.189610Z","shell.execute_reply.started":"2024-05-10T09:20:25.179178Z","shell.execute_reply":"2024-05-10T09:20:25.188884Z"},"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-05-10T09:20:25.190653Z","iopub.execute_input":"2024-05-10T09:20:25.191014Z","iopub.status.idle":"2024-05-10T09:20:25.198861Z","shell.execute_reply.started":"2024-05-10T09:20:25.190991Z","shell.execute_reply":"2024-05-10T09:20:25.198146Z"},"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-05-10T09:20:25.199747Z","iopub.execute_input":"2024-05-10T09:20:25.200038Z","iopub.status.idle":"2024-05-10T09:20:25.211151Z","shell.execute_reply.started":"2024-05-10T09:20:25.200016Z","shell.execute_reply":"2024-05-10T09:20:25.210246Z"},"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-05-10T09:20:25.212213Z","iopub.execute_input":"2024-05-10T09:20:25.212565Z","iopub.status.idle":"2024-05-10T09:20:25.219267Z","shell.execute_reply.started":"2024-05-10T09:20:25.212534Z","shell.execute_reply":"2024-05-10T09:20:25.218442Z"},"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-05-10T09:20:25.222271Z","iopub.execute_input":"2024-05-10T09:20:25.222775Z","iopub.status.idle":"2024-05-10T09:20:25.228625Z","shell.execute_reply.started":"2024-05-10T09:20:25.222740Z","shell.execute_reply":"2024-05-10T09:20:25.227768Z"},"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-05-10T09:20:25.229672Z","iopub.execute_input":"2024-05-10T09:20:25.230021Z","iopub.status.idle":"2024-05-10T09:20:25.239098Z","shell.execute_reply.started":"2024-05-10T09:20:25.229990Z","shell.execute_reply":"2024-05-10T09:20:25.238197Z"},"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-05-10T09:28:10.817775Z","iopub.execute_input":"2024-05-10T09:28:10.818620Z","iopub.status.idle":"2024-05-10T09:28:10.827635Z","shell.execute_reply.started":"2024-05-10T09:28:10.818593Z","shell.execute_reply":"2024-05-10T09:28:10.826644Z"},"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-05-10T09:28:11.598026Z","iopub.execute_input":"2024-05-10T09:28:11.598445Z","iopub.status.idle":"2024-05-10T09:28:11.605634Z","shell.execute_reply.started":"2024-05-10T09:28:11.598414Z","shell.execute_reply":"2024-05-10T09:28:11.604654Z"},"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-05-10T09:28:12.602742Z","iopub.execute_input":"2024-05-10T09:28:12.603592Z","iopub.status.idle":"2024-05-10T09:28:12.622616Z","shell.execute_reply.started":"2024-05-10T09:28:12.603556Z","shell.execute_reply":"2024-05-10T09:28:12.621585Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"run_from_config(Config)","metadata":{"execution":{"iopub.status.busy":"2024-05-10T09:28:20.184765Z","iopub.execute_input":"2024-05-10T09:28:20.185156Z","iopub.status.idle":"2024-05-10T09:28:41.945321Z","shell.execute_reply.started":"2024-05-10T09:28:20.185127Z","shell.execute_reply":"2024-05-10T09:28:41.944112Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!cat submission.csv","metadata":{"execution":{"iopub.status.busy":"2024-05-10T09:26:26.318474Z","iopub.execute_input":"2024-05-10T09:26:26.318773Z","iopub.status.idle":"2024-05-10T09:26:27.303026Z","shell.execute_reply.started":"2024-05-10T09:26:26.318747Z","shell.execute_reply":"2024-05-10T09:26:27.301819Z"},"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":{}}]}