{"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":8347846,"sourceType":"datasetVersion","datasetId":4951592},{"sourceId":4534,"sourceType":"modelInstanceVersion","modelInstanceId":3326},{"sourceId":4535,"sourceType":"modelInstanceVersion","modelInstanceId":3327},{"sourceId":17191,"sourceType":"modelInstanceVersion","modelInstanceId":14317},{"sourceId":17555,"sourceType":"modelInstanceVersion","modelInstanceId":14611}],"dockerImageVersionId":30699,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import shutil, os\nfrom pathlib import Path\nimport pandas as pd\nimport numpy as np\nhome_folder = \"/kaggle/working/\"\nnew_test_fold = home_folder + \"test/\"\nold_test_fold = '/kaggle/input/image-matching-challenge-2024/test/'\nif os.path.exists(new_test_fold):\n    shutil.rmtree(new_test_fold)\nos.makedirs(new_test_fold)\nlist_scene_test = [p for p in Path(old_test_fold).iterdir() if p.is_dir()]\n\nuser_df = pd.read_csv('/kaggle/input/image-matching-challenge-2024/sample_submission.csv')\norin_image_path = user_df['image_path'].to_list()\nprint(\"len sbmission: \", len(user_df))\nfor scene in list_scene_test:\n    scene = str(scene)\n    scene_train = os.path.join(\"/kaggle/input/custom-dataset/train/train/\", os.path.basename(scene))\n    save_new_test = os.path.join(new_test_fold, os.path.basename(scene), \"images\")\n    if os.path.exists(scene_train):\n        os.makedirs(save_new_test, exist_ok = True)\n        images_training = os.listdir(scene_train+\"/images\")\n        images_test = os.listdir(scene+\"/images\")\n        num_add = min(50, len(images_test))\n        add_imgs = np.random.choice(images_training, num_add)\n    #     add_imgs = np.random.choice(images_training, 2)\n\n        for images in add_imgs:\n            src = os.path.join(scene_train, \"images\", images)\n            dst = os.path.join(save_new_test, images)\n            image_path = dst.replace(home_folder, \"\")\n            if image_path in user_df['image_path'].to_list():\n                continue\n            try:\n                shutil.copy(src, dst)\n                new_row = pd.Series({'image_path': image_path, 'dataset': os.path.basename(scene), 'scene': os.path.basename(scene), \"rotation_matrix\":\"\", \"translation_vector\":\"\"})\n                user_df.loc[len(user_df)] = new_row\n            except:\n                print(\"cannot create new\")\nprint(\"len after\", len(user_df))\nfor dirpath, dirnames, filenames in os.walk(new_test_fold):\n    # Skip the root directory\n    if dirpath == new_test_fold:\n        continue\n\n    num_files = len(filenames)\n    print(f\"Folder: {os.path.relpath(dirpath, new_test_fold)} - Number of files: {num_files}\")\nuser_df.to_csv(\"/kaggle/working/sample_submission.csv\", index=False, index_label=False)","metadata":{"execution":{"iopub.status.busy":"2024-05-08T16:47:55.572291Z","iopub.execute_input":"2024-05-08T16:47:55.572669Z","iopub.status.idle":"2024-05-08T16:47:56.324163Z","shell.execute_reply.started":"2024-05-08T16:47:55.572639Z","shell.execute_reply":"2024-05-08T16:47:56.323218Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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/aliked_lightglue.pth /root/.cache/torch/hub/checkpoints/aliked_lightglue_v0-1_arxiv-pth","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-05-08T16:47:56.326414Z","iopub.execute_input":"2024-05-08T16:47:56.327086Z","iopub.status.idle":"2024-05-08T16:48:01.516127Z","shell.execute_reply.started":"2024-05-08T16:47:56.327051Z","shell.execute_reply":"2024-05-08T16:48:01.514883Z"},"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\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\n\n# Provided by organizers\nimport sys\nsys.path.append(\"/kaggle/input/custom-dataset\")\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, 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)\nlist_scene_test = [p for p in Path(\"/kaggle/input/image-matching-challenge-2024/test/\").iterdir() if p.is_dir()]\nDEBUG = len(list_scene_test) == 1\nprint(\"DEBUG:\", DEBUG)\nprint(list_scene_test)","metadata":{"execution":{"iopub.status.busy":"2024-05-08T16:48:01.517454Z","iopub.execute_input":"2024-05-08T16:48:01.517727Z","iopub.status.idle":"2024-05-08T16:48:08.186294Z","shell.execute_reply.started":"2024-05-08T16:48:01.517701Z","shell.execute_reply":"2024-05-08T16:48:08.185338Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def embed_images(\n    paths: list[Path],\n    model_name: str,\n    device: torch.device = torch.device(\"cpu\"),\n) -> T:\n    \"\"\"Computes image embeddings.\n    \n    Returns a tensor of shape [len(filenames), output_dim]\n    \"\"\"\n    processor = AutoImageProcessor.from_pretrained(model_name)\n    model = AutoModel.from_pretrained(model_name).eval().to(device)\n    \n    embeddings = []\n    \n    for i, path in tqdm(enumerate(paths), desc=\"Global descriptors\"):\n        image = load_torch_image(path)\n        \n        with torch.inference_mode():\n            inputs = processor(images=image, return_tensors=\"pt\", do_rescale=False).to(device)\n            outputs = model(**inputs) # last_hidden_state and pooled\n            \n            # Max pooling over all the hidden states but the first (starting token)\n            # To obtain a tensor of shape [1, output_dim]\n            # We normalize so that distances are computed in a better fashion later\n            embedding = F.normalize(outputs.last_hidden_state[:,1:].max(dim=1)[0], dim=-1, p=2)\n            \n        embeddings.append(embedding.detach().cpu())\n    return torch.cat(embeddings, dim=0)\ndef 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, device)\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-08T16:48:08.187928Z","iopub.execute_input":"2024-05-08T16:48:08.188752Z","iopub.status.idle":"2024-05-08T16:48:08.209066Z","shell.execute_reply.started":"2024-05-08T16:48:08.188714Z","shell.execute_reply":"2024-05-08T16:48:08.208079Z"},"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    print(\"computing detect_keypoints\", len(paths))\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=\"detect_keypoints\"):\n#         for path in paths:\n            \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-08T16:48:08.212385Z","iopub.execute_input":"2024-05-08T16:48:08.212667Z","iopub.status.idle":"2024-05-08T16:48:08.227768Z","shell.execute_reply.started":"2024-05-08T16:48:08.212643Z","shell.execute_reply":"2024-05-08T16:48:08.226835Z"},"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    print(\"computing keypoint_distances\", len(paths), \"len pair\", len(index_pairs))\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=\"keypoint_distances\"):\n#             for idx1, idx2 in index_pairs:\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-08T16:48:08.228956Z","iopub.execute_input":"2024-05-08T16:48:08.229218Z","iopub.status.idle":"2024-05-08T16:48:08.243838Z","shell.execute_reply.started":"2024-05-08T16:48:08.229197Z","shell.execute_reply":"2024-05-08T16:48:08.242899Z"},"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, \"dinov2\")\n#     print(index_pairs)\n#     feature_dir = Path(\"./sample_test_features\")\n#     detect_keypoints(images_list, feature_dir)\n#     keypoint_distances(images_list, index_pairs, feature_dir, verbose=False, device=device)\n#     for idx1, idx2 in index_pairs:\n#         key1, key2 = images_list[idx1], images_list[idx2]\n#         img1 = load_torch_image(key1)\n#         img2 = load_torch_image(key2)\n#         fig, ax = plt.subplots(1, 2, figsize=(10, 20))\n#         ax[0].imshow(img1[0, ...].permute(1,2,0).cpu())\n#         ax[1].imshow(img2[0, ...].permute(1,2,0).cpu())\n#         plt.show()","metadata":{"execution":{"iopub.status.busy":"2024-05-08T16:48:08.245032Z","iopub.execute_input":"2024-05-08T16:48:08.245338Z","iopub.status.idle":"2024-05-08T16:48:08.258065Z","shell.execute_reply.started":"2024-05-08T16:48:08.245302Z","shell.execute_reply":"2024-05-08T16:48:08.257103Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def import_into_colmap(\n    path: Path,\n    feature_dir: Path,\n    database_path: str = \"colmap.db\",\n) -> None:\n    \"\"\"Adds keypoints into colmap\"\"\"\n    db = COLMAPDatabase.connect(database_path)\n    db.create_tables()\n    single_camera = False\n    path1 = path\n    path2 = Path(str(path1).replace(\"input/image-matching-challenge-2024\", \"working\"))\n    orin_image_name = os.listdir(path1)\n    fname_to_id = add_keypoints(db, feature_dir, path1, path2, orin_image_name, \"\", \"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-08T16:48:08.259474Z","iopub.execute_input":"2024-05-08T16:48:08.259736Z","iopub.status.idle":"2024-05-08T16:48:08.272976Z","shell.execute_reply.started":"2024-05-08T16:48:08.259714Z","shell.execute_reply":"2024-05-08T16:48:08.272131Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Config:\n    base_path: Path = Path(\"/kaggle/input/image-matching-challenge-2024/\")\n#     base_path: Path = Path(\"/kaggle/input/custom-dataset\")\n        \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/large/1\",\n        \"similarity_threshold\": 0.6,\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-08T17:07:51.811987Z","iopub.execute_input":"2024-05-08T17:07:51.812748Z","iopub.status.idle":"2024-05-08T17:07:51.819923Z","shell.execute_reply.started":"2024-05-08T17:07:51.812715Z","shell.execute_reply":"2024-05-08T17:07:51.818861Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def parse_sample_submission() -> 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(\"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                if image_path in orin_image_path:\n                    base_path = Path(\"/kaggle/input/image-matching-challenge-2024/\")\n                else:\n                    base_path = Path(\"/kaggle/working/\")\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\ndef create_submission(\n    results: dict,\n    data_dict: dict[dict[str, list[Path]]],\n    base_path: Path,\n    orin_image_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                scene_res_2 = {}\n                for skey in scene_res.keys():\n                    idx = str(skey).find(\"test\")\n                \n                    skey2 = str(skey.relative_to(base_path))\n                    scene_res_2[skey2] = scene_res[skey]\n                    \n                # Write the row with rotation and translation matrices\n                for image in data_dict[dataset][scene]:\n                    image = str(image)\n                    idx = image.find(\"test\")\n                    image = image[idx:]\n                    if image in scene_res_2:\n                        print(image)\n                        R = scene_res_2[image][\"R\"].reshape(-1)\n                        T = scene_res_2[image][\"t\"].reshape(-1)\n                    else:\n                        R = np.eye(3).reshape(-1)\n                        T = np.zeros((3))\n                    if image in orin_image_path:\n                        f.write(f\"{image},{dataset},{scene},{arr_to_str(R)},{arr_to_str(T)}\\n\")","metadata":{"execution":{"iopub.status.busy":"2024-05-08T16:48:08.286650Z","iopub.execute_input":"2024-05-08T16:48:08.286909Z","iopub.status.idle":"2024-05-08T16:48:08.302999Z","shell.execute_reply.started":"2024-05-08T16:48:08.286887Z","shell.execute_reply":"2024-05-08T16:48:08.301967Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# DEBUG = False","metadata":{"execution":{"iopub.status.busy":"2024-05-08T16:48:08.304378Z","iopub.execute_input":"2024-05-08T16:48:08.304939Z","iopub.status.idle":"2024-05-08T16:48:08.316013Z","shell.execute_reply.started":"2024-05-08T16:48:08.304908Z","shell.execute_reply":"2024-05-08T16:48:08.314962Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"results = {}\n\ndata_dict = parse_sample_submission()\ndatasets = list(data_dict.keys())\n\nfor 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\"{scene}: 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, orin_image_path)\n            gc.collect()\n\n        except Exception as e:\n            print(e)","metadata":{"execution":{"iopub.status.busy":"2024-05-08T16:48:08.317353Z","iopub.execute_input":"2024-05-08T16:48:08.317945Z","iopub.status.idle":"2024-05-08T16:56:22.495238Z","shell.execute_reply.started":"2024-05-08T16:48:08.317919Z","shell.execute_reply":"2024-05-08T16:56:22.494265Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Evaluate","metadata":{}},{"cell_type":"code","source":"import time\nimport math\nimport numpy as np\nimport pandas as pd\nimport pandas.api.types\n\n_EPS = np.finfo(float).eps * 4.0\n\n# mAA evaluation thresholds per scene, different accoring to the scene\ntranslation_thresholds_meters_dict = {\n 'multi-temporal-temple-baalshamin':  np.array([0.025,  0.05,  0.1,  0.2,  0.5,  1.0]),\n 'pond':                              np.array([0.025,  0.05,  0.1,  0.2,  0.5,  1.0]),\n 'transp_obj_glass_cylinder':         np.array([0.0025, 0.005, 0.01, 0.02, 0.05, 0.1]),\n 'transp_obj_glass_cup':              np.array([0.0025, 0.005, 0.01, 0.02, 0.05, 0.1]),\n 'church':                            np.array([0.025,  0.05,  0.1,  0.2,  0.5,  1.0]),\n 'lizard':                            np.array([0.025,  0.05,  0.1,  0.2,  0.5,  1.0]),\n 'dioscuri':                          np.array([0.025,  0.05,  0.1,  0.2,  0.5,  1.0]), \n}\n\n\ndef vector_norm(data, axis=None, out=None):\n    '''Return length, i.e. Euclidean norm, of ndarray along axis.'''\n    data = np.array(data, dtype=np.float64, copy=True)\n    if out is None:\n        if data.ndim == 1:\n            return math.sqrt(np.dot(data, data))\n        data *= data\n        out = np.atleast_1d(np.sum(data, axis=axis))\n        np.sqrt(out, out)\n        return out\n    data *= data\n    np.sum(data, axis=axis, out=out)\n    np.sqrt(out, out)\n    return None\n\n\ndef quaternion_matrix(quaternion):\n    '''Return homogeneous rotation matrix from quaternion.'''\n    q = np.array(quaternion, dtype=np.float64, copy=True)\n    n = np.dot(q, q)\n    if n < _EPS:\n        # print(\"special case\")\n        return np.identity(4)\n    q *= math.sqrt(2.0 / n)\n    q = np.outer(q, q)\n    return np.array(\n        [\n            [\n                1.0 - q[2, 2] - q[3, 3],\n                q[1, 2] - q[3, 0],\n                q[1, 3] + q[2, 0],\n                0.0,\n            ],\n            [\n                q[1, 2] + q[3, 0],\n                1.0 - q[1, 1] - q[3, 3],\n                q[2, 3] - q[1, 0],\n                0.0,\n            ],\n            [\n                q[1, 3] - q[2, 0],\n                q[2, 3] + q[1, 0],\n                1.0 - q[1, 1] - q[2, 2],\n                0.0,\n            ],\n            [0.0, 0.0, 0.0, 1.0],\n        ]\n    )\n\n\n# based on the 3D registration from https://github.com/cgohlke/transformations\ndef affine_matrix_from_points(v0, v1, shear=False, scale=True, usesvd=True):\n    '''Return affine transform matrix to register two point sets.\n    v0 and v1 are shape (ndims, -1) arrays of at least ndims non-homogeneous\n    coordinates, where ndims is the dimensionality of the coordinate space.\n    If shear is False, a similarity transformation matrix is returned.\n    If also scale is False, a rigid/Euclidean traffansformation matrix\n    is returned.\n    By default the algorithm by Hartley and Zissermann [15] is used.\n    If usesvd is True, similarity and Euclidean transformation matrices\n    are calculated by minimizing the weighted sum of squared deviations\n    (RMSD) according to the algorithm by Kabsch [8].\n    Otherwise, and if ndims is 3, the quaternion based algorithm by Horn [9]\n    is used, which is slower when using this Python implementation.\n    The returned matrix performs rotation, translation and uniform scaling\n    (if specified).'''\n    \n    v0 = np.array(v0, dtype=np.float64, copy=True)\n    v1 = np.array(v1, dtype=np.float64, copy=True)\n\n    ndims = v0.shape[0]\n    if ndims < 2 or v0.shape[1] < ndims or v0.shape != v1.shape:\n        raise ValueError(\"input arrays are of wrong shape or type\")\n\n    # move centroids to origin\n    t0 = -np.mean(v0, axis=1)\n    M0 = np.identity(ndims + 1)\n    M0[:ndims, ndims] = t0\n    v0 += t0.reshape(ndims, 1)\n    t1 = -np.mean(v1, axis=1)\n    M1 = np.identity(ndims + 1)\n    M1[:ndims, ndims] = t1\n    v1 += t1.reshape(ndims, 1)\n\n    if shear:\n        # Affine transformation\n        A = np.concatenate((v0, v1), axis=0)\n        u, s, vh = np.linalg.svd(A.T)\n        vh = vh[:ndims].T\n        B = vh[:ndims]\n        C = vh[ndims: 2 * ndims]\n        t = np.dot(C, np.linalg.pinv(B))\n        t = np.concatenate((t, np.zeros((ndims, 1))), axis=1)\n        M = np.vstack((t, ((0.0,) * ndims) + (1.0,)))\n    elif usesvd or ndims != 3:\n        # Rigid transformation via SVD of covariance matrix\n        u, s, vh = np.linalg.svd(np.dot(v1, v0.T))\n        # rotation matrix from SVD orthonormal bases\n        R = np.dot(u, vh)\n        if np.linalg.det(R) < 0.0:\n            # R does not constitute right handed system\n            R -= np.outer(u[:, ndims - 1], vh[ndims - 1, :] * 2.0)\n            s[-1] *= -1.0\n        # homogeneous transformation matrix\n        M = np.identity(ndims + 1)\n        M[:ndims, :ndims] = R\n    else:\n        # Rigid transformation matrix via quaternion\n        # compute symmetric matrix N\n        xx, yy, zz = np.sum(v0 * v1, axis=1)\n        xy, yz, zx = np.sum(v0 * np.roll(v1, -1, axis=0), axis=1)\n        xz, yx, zy = np.sum(v0 * np.roll(v1, -2, axis=0), axis=1)\n        N = [\n            [xx + yy + zz, 0.0, 0.0, 0.0],\n            [yz - zy, xx - yy - zz, 0.0, 0.0],\n            [zx - xz, xy + yx, yy - xx - zz, 0.0],\n            [xy - yx, zx + xz, yz + zy, zz - xx - yy],\n        ]\n        # quaternion: eigenvector corresponding to most positive eigenvalue\n        w, V = np.linalg.eigh(N)\n        q = V[:, np.argmax(w)]\n        # print (vector_norm(q), np.linalg.norm(q))\n        q /= vector_norm(q)  # unit quaternion\n        # homogeneous transformation matrix\n        M = quaternion_matrix(q)\n\n    if scale and not shear:\n        # Affine transformation; scale is ratio of RMS deviations from centroid\n        v0 *= v0\n        v1 *= v1\n        M[:ndims, :ndims] *= math.sqrt(np.sum(v1) / np.sum(v0))\n\n    # move centroids back\n    M = np.dot(np.linalg.inv(M1), np.dot(M, M0))\n    M /= M[ndims, ndims]\n\n    # print(\"transformation matrix Python Script: \", M)\n\n    return M\n\n\n# This is the IMC 3D error metric code\ndef register_by_Horn(ev_coord, gt_coord, ransac_threshold, inl_cf, strict_cf):\n    '''Return the best similarity transforms T that registers 3D points pt_ev in <ev_coord> to\n    the corresponding ones pt_gt in <gt_coord> according to a RANSAC-like approach for each\n    threshold value th in <ransac_threshold>.\n    \n    Given th, each triplet of 3D correspondences is examined if not already present as strict inlier,\n    a correspondence is a strict inlier if <strict_cf> * err_best < th, where err_best is the registration\n    error for the best model so far.\n    The minimal model given by the triplet is then refined using also its inliers if their total is greater\n    than <inl_cf> * ninl_best, where ninl_best is th number of inliers for the best model so far. Inliers\n    are 3D correspondences (pt_ev, pt_gt) for which the Euclidean distance |pt_gt-T*pt_ev| is less than th.'''\n    \n    # remove invalid cameras, the index is returned\n    idx_cams = np.all(np.isfinite(ev_coord), axis=0)\n    ev_coord = ev_coord[:, idx_cams]\n    gt_coord = gt_coord[:, idx_cams]\n\n    # initialization\n    n = ev_coord.shape[1]\n    r = ransac_threshold.shape[0]\n    ransac_threshold = np.expand_dims(ransac_threshold, axis=0)\n    ransac_threshold2 = ransac_threshold**2\n    ev_coord_1 = np.vstack((ev_coord, np.ones(n)))\n\n    max_no_inl = np.zeros((1, r))\n    best_inl_err = np.full(r, np.Inf)\n    best_transf_matrix = np.zeros((r, 4, 4))\n    best_err = np.full((n, r), np.Inf)\n    strict_inl = np.full((n, r), False)\n    triplets_used = np.zeros((3, r))\n\n    # run on camera triplets\n    for ii in range(n-2):\n        for jj in range(ii+1, n-1):\n            for kk in range(jj+1, n):\n                i = [ii, jj, kk]\n                triplets_used_now = np.full((n), False)\n                triplets_used_now[i] = True\n                # if both ii, jj, kk are strict inliers for the best current model just skip\n                if np.all(strict_inl[i]):\n                    continue\n                # get transformation T by Horn on the triplet camera center correspondences\n                transf_matrix = affine_matrix_from_points(ev_coord[:, i], gt_coord[:, i], usesvd=False)\n                # apply transformation T to test camera centres\n                rotranslated = np.matmul(transf_matrix[:3], ev_coord_1)\n                # compute error and inliers\n                err = np.sum((rotranslated - gt_coord)**2, axis=0)\n                inl = np.expand_dims(err, axis=1) < ransac_threshold2\n                no_inl = np.sum(inl, axis=0)\n                # if the number of inliers is close to that of the best model so far, go for refinement\n                to_ref = np.squeeze(((no_inl > 2) & (no_inl > max_no_inl * inl_cf)), axis=0)\n                for q in np.argwhere(to_ref):                        \n                    qq = q[0]\n                    if np.any(np.all((np.expand_dims(inl[:, qq], axis=1) == inl[:, :qq]), axis=0)):\n                        # already done for this set of inliers\n                        continue\n                    # get transformation T by Horn on the inlier camera center correspondences\n                    transf_matrix = affine_matrix_from_points(ev_coord[:, inl[:, qq]], gt_coord[:, inl[:, qq]])\n                    # apply transformation T to test camera centres\n                    rotranslated = np.matmul(transf_matrix[:3], ev_coord_1)\n                    # compute error and inliers\n                    err_ref = np.sum((rotranslated - gt_coord)**2, axis=0)\n                    err_ref_sum = np.sum(err_ref, axis=0)\n                    err_ref = np.expand_dims(err_ref, axis=1)\n                    inl_ref = err_ref < ransac_threshold2\n                    no_inl_ref = np.sum(inl_ref, axis=0)\n                    # update the model if better for each threshold\n                    to_update = np.squeeze((no_inl_ref > max_no_inl) | ((no_inl_ref == max_no_inl) & (err_ref_sum < best_inl_err)), axis=0)\n                    if np.any(to_update):\n                        triplets_used[0, to_update] = ii\n                        triplets_used[1, to_update] = jj\n                        triplets_used[2, to_update] = kk\n                        max_no_inl[:, to_update] = no_inl_ref[to_update]\n                        best_err[:, to_update] = np.sqrt(err_ref)\n                        best_inl_err[to_update] = err_ref_sum\n                        strict_inl[:, to_update] = (best_err[:, to_update] < strict_cf * ransac_threshold[:, to_update])\n                        best_transf_matrix[to_update] = transf_matrix\n\n    for i in range(r):\n       print(f'Registered cameras {int(max_no_inl[0, i])}/{n} for threshold {ransac_threshold[0, i]}')\n\n    best_model = {\n        \"valid_cams\": idx_cams,        \n        \"no_inl\": max_no_inl,\n        \"err\": best_err,\n        \"triplets_used\": triplets_used,\n        \"transf_matrix\": best_transf_matrix}\n    return best_model\n\n\n# mAA computation\ndef mAA_on_cameras(err, thresholds, n, skip_top_thresholds, to_dec=3):\n    '''mAA is the mean of mAA_i, where for each threshold th_i in <thresholds>, excluding the first <skip_top_thresholds values>,\n    mAA_i = max(0, sum(err_i < th_i) - <to_dec>) / (n - <to_dec>)\n    where <n> is the number of ground-truth cameras and err_i is the camera registration error for the best \n    registration corresponding to threshold th_i'''\n    \n    aux = err[:, skip_top_thresholds:] < np.expand_dims(np.asarray(thresholds[skip_top_thresholds:]), axis=0)\n    return np.sum(np.maximum(np.sum(aux, axis=0) - to_dec, 0)) / (len(thresholds[skip_top_thresholds:]) * (n - to_dec))\n\n\n# import data - no error handling in case float(x) fails\ndef get_camera_centers_from_df(df):\n    out = {}\n    for row in df.iterrows():\n        row = row[1]\n        fname = row['image_path']\n        R = np.array([float(x) for x in (row['rotation_matrix'].split(';'))]).reshape(3,3)\n        t = np.array([float(x) for x in (row['translation_vector'].split(';'))]).reshape(3)\n        center = -R.T @ t\n        out[fname] = center\n    return out\n\n\ndef evaluate_rec(gt_df, user_df, inl_cf = 0.8, strict_cf=0.5, skip_top_thresholds=2, to_dec=3,\n                 thresholds=[0.005, 0.01, 0.02, 0.03, 0.04, 0.05, 0.1, 0.15, 0.2]):\n    ''' Register the <user_df> camera centers to the ground-truth <gt_df> camera centers and\n    return the corresponding mAA as the average percentage of registered camera threshold.\n    \n    For each threshold value in <thresholds>, the best similarity transformation found which\n    maximizes the number of registered cameras is employed. A camera is marked as registered\n    if after the transformation its Euclidean distance to the corresponding ground-truth camera\n    center is less than the mentioned threshold. Current measurements are in meter.\n    \n    Registration parameters:\n    <inl_cf> coefficient to activate registration refinement, set to 1 to refine a new model\n    only when it gives more inliers, to 0 to refine a new model always; high values increase\n    speed but decrease precision.\n    <strict_cf> threshold coefficient to define strict inliers for the best registration so far,\n    new minimal models made up of strict inliers are skipped. It can vary from 0 (slower) to\n    1 (faster); set to -1 to check exhaustively all the minimal model triplets.\n\n    mAA parameters:\n    <skip_top_thresholds> excluded lower thresholds in the mAA computation; in case of using\n    heuristics for the registration, i.e. inl_cf!=0 and strict_cf!=-1, best model for lower\n    threshold can be not the optimal, so skip them in the mAA computation.\n    <to_dec> excludes the minimal model cameras from the computation of the mAA. Given the\n    minimal model, i.e. three pairs of 3D correspondences, there is a high chance to register by\n    a similarity transformation at any threshold, so do not account for mAA'''\n    \n    # get camera centers\n    ucameras = get_camera_centers_from_df(user_df)\n    gcameras = get_camera_centers_from_df(gt_df)    \n\n    # the denominator for mAA ratio\n    m = gt_df.shape[0]\n    \n    # get the image list to use\n    good_cams = []\n    for image_path in gcameras.keys():\n        if image_path in ucameras.keys():\n            good_cams.append(image_path)\n    print(\"good_cams\", good_cams)\n        \n    # put corresponding camera centers into matrices\n    n = len(good_cams)\n    u_cameras = np.zeros((3, n))\n    g_cameras = np.zeros((3, n))\n    \n    ii = 0\n    for i in good_cams:\n        u_cameras[:, ii] = ucameras[i]\n        g_cameras[:, ii] = gcameras[i]\n        ii += 1\n        \n    # Horn camera centers registration, a different best model for each camera threshold\n    model = register_by_Horn(u_cameras, g_cameras, np.asarray(thresholds), inl_cf, strict_cf)\n    \n    # transformation matrix\n    print(\"\\nTransformation matrix for maximum threshold\")\n    T = np.squeeze(model['transf_matrix'][-1])\n    print(T)\n    \n    # mAA\n    mAA = mAA_on_cameras(model[\"err\"], thresholds, m, skip_top_thresholds, to_dec)\n    # print(f'mAA = {mAA * 100 : .2f}% considering {m} input cameras - {to_dec}')\n    return mAA\n\n\ndef score(solution: pd.DataFrame, submission: pd.DataFrame) -> float:\n    '''The metric is an mean average accuracy between solution and submission camera centers.\n    Prior to calculate the metric, a function performs exhaustive registration (like RANSAC, but\n    not random, considering all possible configurations) to align the user camera system to the GT'''\n    \n    scenes = list(set(solution['dataset'].tolist()))\n    results_per_dataset = []\n    for dataset in scenes:\n        print(f\"\\n*** {dataset} ***\")\n        start = time.time()\n        gt_ds = solution[solution['dataset'] == dataset]\n        user_ds = submission[submission['dataset'] == dataset]\n        gt_ds = gt_ds.sort_values(by=['image_path'], ascending = True)\n        user_ds = user_ds.sort_values(by=['image_path'], ascending = True)\n        result = evaluate_rec(gt_ds, user_ds, inl_cf=0, strict_cf=-1, skip_top_thresholds=0, to_dec=3,\n                 thresholds=translation_thresholds_meters_dict[dataset])\n        end = time.time()\n        print(f\"\\nmAA: {result*100}%\")\n        print(\"Running time: %s\" % (end - start))        \n        results_per_dataset.append(result)\n    return float(np.array(results_per_dataset).mean())","metadata":{"execution":{"iopub.status.busy":"2024-05-08T16:56:22.497100Z","iopub.execute_input":"2024-05-08T16:56:22.497392Z","iopub.status.idle":"2024-05-08T16:56:22.558958Z","shell.execute_reply.started":"2024-05-08T16:56:22.497367Z","shell.execute_reply":"2024-05-08T16:56:22.557985Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os, time\ntrans = np.eye(4)\ntrans [:3, -1] = [0,0,0]\n\ngt_csv = '/kaggle/input/custom-dataset/test_gt.csv'\nuser_csv = '/kaggle/working/submission.csv'\n\ngt_df = pd.read_csv(gt_csv).sort_values(by='image_path')\nuser_df = pd.read_csv(user_csv)\n\nuser_df['image_path'] = user_df['image_path'].apply(lambda x: os.path.basename(x))\nuser_df = user_df.head(41)\nuser_df = user_df.sort_values(by='image_path')\n\n\n\nstart = time.time()\nres = score(gt_df, user_df)\nend = time.time()\n\nprint(f\"\\nGlobal mAA: {res*100}%\")\nprint(\"Total running time: %s\" % (end - start))","metadata":{"execution":{"iopub.status.busy":"2024-05-08T17:07:23.414558Z","iopub.execute_input":"2024-05-08T17:07:23.415005Z","iopub.status.idle":"2024-05-08T17:07:35.658735Z","shell.execute_reply.started":"2024-05-08T17:07:23.414967Z","shell.execute_reply":"2024-05-08T17:07:35.657668Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!ls","metadata":{"execution":{"iopub.status.busy":"2024-05-08T16:56:22.579027Z","iopub.execute_input":"2024-05-08T16:56:22.579621Z","iopub.status.idle":"2024-05-08T16:56:23.582634Z","shell.execute_reply.started":"2024-05-08T16:56:22.579587Z","shell.execute_reply":"2024-05-08T16:56:23.581352Z"},"trusted":true},"execution_count":null,"outputs":[]}]}