{"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":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":30665,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Exploratory Data Analysis\n**🏛️ Phototourism and historical preservation:** different viewpoints, sensor types, time of day/year, and occlusions. Ancient historical sites add a unique set of challenges\n\n**☀️ Night vs day and temporal changes:** combination of day and night photographs, including poor lighting, or photographs taken months or years apart, in different weather\n\n**✈️ Aerial and mixed aerial-ground:** images from drones, featuring arbitrary in-plane rotations, matched against similar images and also images taken from the ground\n\n**♻️ Repeated structures:** symmetrical objects require details to disambiguate perspective\n\n**🌲 Natural environments:** highly non-regular structures such as trees and foliage\n\n**🪞 Transparencies and reflections:** objects like glassware are lacking in texture and create reflections and specularities which pose a different set of problems","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-04-01T23:18:11.151217Z","iopub.execute_input":"2024-04-01T23:18:11.151607Z","iopub.status.idle":"2024-04-01T23:18:21.273213Z","shell.execute_reply.started":"2024-04-01T23:18:11.151574Z","shell.execute_reply":"2024-04-01T23:18:21.272012Z"},"collapsed":true,"jupyter":{"outputs_hidden":true},"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-04-01T23:18:27.868906Z","iopub.execute_input":"2024-04-01T23:18:27.869251Z","iopub.status.idle":"2024-04-01T23:18:45.785797Z","shell.execute_reply.started":"2024-04-01T23:18:27.869224Z","shell.execute_reply":"2024-04-01T23:18:45.784851Z"},"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.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-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":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}