{"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":"nvidiaTeslaT4","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":30698,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install --no-index /kaggle/input/imc2024-packages-lightglue-rerun-kornia/* --no-deps\n!mkdir -p /root/.cache/torch/hub/checkpoints\n!cp /kaggle/input/aliked/pytorch/aliked-n16/1/* /root/.cache/torch/hub/checkpoints/\n!cp /kaggle/input/lightglue/pytorch/aliked/1/* /root/.cache/torch/hub/checkpoints/\n!cp /kaggle/input/lightglue/pytorch/aliked/1/aliked_lightglue.pth /root/.cache/torch/hub/checkpoints/aliked_lightglue_v0-1_arxiv-pth","metadata":{"execution":{"iopub.status.busy":"2024-05-22T17:22:31.919418Z","iopub.execute_input":"2024-05-22T17:22:31.920433Z","iopub.status.idle":"2024-05-22T17:22:41.573383Z","shell.execute_reply.started":"2024-05-22T17:22:31.920400Z","shell.execute_reply":"2024-05-22T17:22:41.572321Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import 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\n\ndef embed_images(\n    paths: list[Path],\n    model_name: str,\n    device: torch.device = torch.device(\"cpu\"),\n) :\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)\n\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\ndef get_image_pairs(paths: list[Path],model_name: str,similarity_threshold: float = 0.6,\n                    tolerance: int = 1000,min_matches: int = 20,exhaustive_if_less: int = 20,\n                    p: float = 2.0,device: torch.device = torch.device(\"cpu\")):\n    # We try to minimize the camera to reconstruct the scene due to saving computing cost and minimizing the error.\n    if len(paths) <= exhaustive_if_less:\n        return get_pairs_exhaustive(paths)\n    matches = []\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)))\n    \n    ","metadata":{"execution":{"iopub.status.busy":"2024-05-22T17:22:47.438489Z","iopub.execute_input":"2024-05-22T17:22:47.439372Z","iopub.status.idle":"2024-05-22T17:23:15.850836Z","shell.execute_reply.started":"2024-05-22T17:22:47.439330Z","shell.execute_reply":"2024-05-22T17:23:15.849841Z"},"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-22T17:23:29.279523Z","iopub.execute_input":"2024-05-22T17:23:29.280439Z","iopub.status.idle":"2024-05-22T17:23:29.291246Z","shell.execute_reply.started":"2024-05-22T17:23:29.280385Z","shell.execute_reply":"2024-05-22T17:23:29.290339Z"},"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                group  = f_matches.require_group(key1)\n                if n_matches >= min_matches:\n                     group.create_dataset(key2, data=indices.detach().cpu().numpy().reshape(-1, 2))","metadata":{"execution":{"iopub.status.busy":"2024-05-22T17:23:32.585775Z","iopub.execute_input":"2024-05-22T17:23:32.586603Z","iopub.status.idle":"2024-05-22T17:23:32.598109Z","shell.execute_reply.started":"2024-05-22T17:23:32.586569Z","shell.execute_reply":"2024-05-22T17:23:32.597204Z"},"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    device: torch.device= K.utils.get_cuda_device_if_available(0)\n    pair_matching_args = {\n        \"model_name\": \"/kaggle/input/dinov2/pytorch/base/1\",\n        \"similarity_threshold\": 0.3,\n        \"tolerance\": 1000,\n        \"min_matches\": 20,\n        \"exhaustive_if_less\": 20,\n        \"p\": 2.0\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    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":{"scrolled":true,"_kg_hide-input":false,"execution":{"iopub.status.busy":"2024-05-22T17:23:35.884612Z","iopub.execute_input":"2024-05-22T17:23:35.885318Z","iopub.status.idle":"2024-05-22T17:23:35.945100Z","shell.execute_reply.started":"2024-05-22T17:23:35.885266Z","shell.execute_reply":"2024-05-22T17:23:35.943958Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def parse_sample_submission(\n    base_path: Path,\n) -> dict[dict[str, list[Path]]]:\n    \"\"\"Construct a dict describing the test data as \n    \n    {\"dataset\": {\"scene\": [<image paths>]}}\n    \"\"\"\n    data_dict = {}\n    with open(base_path / \"sample_submission.csv\", \"r\") as f:\n        for i, l in enumerate(f):\n            # Skip header\n            if i == 0:\n                print(\"header:\", l)\n\n            if l and i > 0:\n                image_path, dataset, scene, _, _ = l.strip().split(',')\n                if dataset not in data_dict:\n                    data_dict[dataset] = {}\n                if scene not in data_dict[dataset]:\n                    data_dict[dataset][scene] = []\n                data_dict[dataset][scene].append(Path(base_path / image_path))\n\n    for dataset in data_dict:\n        for scene in data_dict[dataset]:\n            print(f\"{dataset} / {scene} -> {len(data_dict[dataset][scene])} images\")\n\n    return data_dict","metadata":{"execution":{"iopub.status.busy":"2024-05-22T17:23:37.948963Z","iopub.execute_input":"2024-05-22T17:23:37.949329Z","iopub.status.idle":"2024-05-22T17:23:37.958069Z","shell.execute_reply.started":"2024-05-22T17:23:37.949289Z","shell.execute_reply":"2024-05-22T17:23:37.957082Z"},"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":{"execution":{"iopub.status.busy":"2024-05-22T17:23:40.622123Z","iopub.execute_input":"2024-05-22T17:23:40.623022Z","iopub.status.idle":"2024-05-22T17:23:40.634374Z","shell.execute_reply.started":"2024-05-22T17:23:40.622980Z","shell.execute_reply":"2024-05-22T17:23:40.633419Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# results = {}\n# config = Config()\n# data_dict = parse_sample_submission(config.base_path)\n# datasets = list(data_dict.keys())\n\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=config.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=config.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":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-05-03T06:43:49.013134Z","iopub.execute_input":"2024-05-03T06:43:49.014042Z","iopub.status.idle":"2024-05-03T06:48:38.264876Z","shell.execute_reply.started":"2024-05-03T06:43:49.014008Z","shell.execute_reply":"2024-05-03T06:48:38.263773Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport pickle\nfrom tqdm import tqdm\n\n\nresults = {}\nconfig = Config()\ndata_dict = parse_sample_submission(config.base_path)\ndatasets = list(data_dict.keys())\nroot_dir = '/kaggle/input/image-matching-challenge-2024/train'\nscene_path_lists = [ os.path.join(root_dir,scene_name) for scene_name in os.listdir(root_dir) if('.csv' not in scene_name) ]\ndtype = torch.float32 # ALIKED has issues with float16\ndevice = 'cuda'\n\nextractor = ALIKED(\n    max_num_keypoints=4096, \n    detection_threshold=0.01, \n    resize=1024\n).eval().to('cuda', dtype)\noutput_dir_root: Path = Path.cwd() / \"results_new\"\n\nfor scene_path in scene_path_lists:\n    output_dir = output_dir_root / scene_path.split(\"/\")[-1]\n    output_dir.mkdir(parents=True, exist_ok=True)\n    scene_path = os.path.join(scene_path,'images')\n    img_path_lists = [ os.path.join(scene_path,img_path) for img_path in os.listdir(scene_path)  ]\n    img_path_lists.sort()\n    for idx,img_path in enumerate(tqdm(img_path_lists)):\n        img_num = img_path.split(\"/\")[-1]\n        img_num = img_num.split(\".\")[0]\n        output_path = output_dir / img_num\n        output_path.mkdir(parents=True, exist_ok=True)\n        image = load_torch_image(img_path, device=device).to(dtype)\n        features = extractor.extract(image)\n        with open(f'{output_path}/keypoints.pkl', 'wb') as f:\n            pickle.dump(features['keypoints'].detach().cpu().numpy(), f)\n        with open(f'{output_path}/descriptors.pkl', 'wb') as f:\n            pickle.dump(features['descriptors'].detach().cpu().numpy(), f)\n        \n#     asd","metadata":{"execution":{"iopub.status.busy":"2024-05-22T17:23:45.266913Z","iopub.execute_input":"2024-05-22T17:23:45.267760Z","iopub.status.idle":"2024-05-22T17:27:26.249537Z","shell.execute_reply.started":"2024-05-22T17:23:45.267729Z","shell.execute_reply":"2024-05-22T17:27:26.248670Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!zip -r result.zip /kaggle/working/results_new","metadata":{"execution":{"iopub.status.busy":"2024-05-22T17:28:04.138778Z","iopub.execute_input":"2024-05-22T17:28:04.139157Z","iopub.status.idle":"2024-05-22T17:32:13.181003Z","shell.execute_reply.started":"2024-05-22T17:28:04.139131Z","shell.execute_reply":"2024-05-22T17:32:13.179872Z"},"trusted":true},"execution_count":null,"outputs":[]}]}