{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":71885,"databundleVersionId":8143495,"sourceType":"competition"},{"sourceId":7884485,"sourceType":"datasetVersion","datasetId":4628051},{"sourceId":7884725,"sourceType":"datasetVersion","datasetId":4628331},{"sourceId":4534,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":3326},{"sourceId":17191,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":14317},{"sourceId":17555,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":14611}],"dockerImageVersionId":30684,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-04-20T02:24:10.294602Z","iopub.execute_input":"2024-04-20T02:24:10.295246Z","iopub.status.idle":"2024-04-20T02:24:11.757518Z","shell.execute_reply.started":"2024-04-20T02:24:10.295213Z","shell.execute_reply":"2024-04-20T02:24:11.75666Z"},"collapsed":true,"jupyter":{"outputs_hidden":true},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\n\npath = '/kaggle/input'\nfile_list = os.listdir(path)\n\nprint (\"file_list: {}\".format(file_list))","metadata":{"execution":{"iopub.status.busy":"2024-04-20T02:24:11.759316Z","iopub.execute_input":"2024-04-20T02:24:11.759725Z","iopub.status.idle":"2024-04-20T02:24:11.764973Z","shell.execute_reply.started":"2024-04-20T02:24:11.759698Z","shell.execute_reply":"2024-04-20T02:24:11.764138Z"},"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/* /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-04-20T02:24:11.766006Z","iopub.execute_input":"2024-04-20T02:24:11.766263Z","iopub.status.idle":"2024-04-20T02:24:20.450696Z","shell.execute_reply.started":"2024-04-20T02:24:11.766241Z","shell.execute_reply":"2024-04-20T02:24:20.449587Z"},"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\n# Data importing into colmap\nimport sys\nsys.path.append(\"/kaggle/input/colmap-db-import\")\n\n\n# Provided by organizers\nfrom database import *\nfrom h5_to_db import *\n\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: \",device)\n\n# DEBUG = len([p for p in Path(\"/kaggle/input/image-matching-challenge-2024/test/\").iterdir() if p.is_dir()]) == 2\nDEBUG = True\nprint(\"DEBUG:\", DEBUG)\n","metadata":{"execution":{"iopub.status.busy":"2024-04-20T02:24:20.452251Z","iopub.execute_input":"2024-04-20T02:24:20.452604Z","iopub.status.idle":"2024-04-20T02:24:37.936741Z","shell.execute_reply.started":"2024-04-20T02:24:20.452573Z","shell.execute_reply":"2024-04-20T02:24:37.935751Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Finding Image Pairs\n## DINOv2","metadata":{}},{"cell_type":"code","source":"\n\ndef embed_images(\n    paths: list[Path],\n    model_name: str,\n    device: torch.device = torch.device('cpu'),\n) -> T:\n    \"\"\"\n    Compute image embeddings\n    Return 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_model():\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\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            ","metadata":{"execution":{"iopub.status.busy":"2024-04-20T02:24:37.939198Z","iopub.execute_input":"2024-04-20T02:24:37.93976Z","iopub.status.idle":"2024-04-20T02:24:37.947904Z","shell.execute_reply.started":"2024-04-20T02:24:37.939732Z","shell.execute_reply":"2024-04-20T02:24:37.94685Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_pairs_exhaustive(lst: list[Any]) -> list[tuple[int, int]]:\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    if len(paths) <= exhaustive_if_less:\n        return get_pairs_exhaustive(paths)\n    matches = []\n    \n    # Embed images and compute distances for filtering\n    embeddings = embed_images(paths, model_name, device) #JAYEONG device import\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            if other_image_index == current_image_index:\n                continue\n            \n            # We need to check if we are below a certain distance torlerance\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, othre_image_index] < tolerance:\n                matches.append(tuple(sorted((current_image_index, other_image_index.item()))))\n    return sorted(list(set(matches)))","metadata":{"execution":{"iopub.status.busy":"2024-04-20T02:24:37.94953Z","iopub.execute_input":"2024-04-20T02:24:37.949884Z","iopub.status.idle":"2024-04-20T02:24:37.978816Z","shell.execute_reply.started":"2024-04-20T02:24:37.949852Z","shell.execute_reply":"2024-04-20T02:24:37.97768Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"path = \"/kaggle/input/image-matching-challenge-2024/test/church/images/\"\nif DEBUG:\n    images_list = list(Path(path).glob(\"*.png\"))[:10]\n    index_pairs = get_image_pairs(images_list, \"/kaggle/input/dinov2/pytorch/base/1\", device)\n    print(index_pairs)","metadata":{"execution":{"iopub.status.busy":"2024-04-20T02:24:37.980753Z","iopub.execute_input":"2024-04-20T02:24:37.981178Z","iopub.status.idle":"2024-04-20T02:24:37.995212Z","shell.execute_reply.started":"2024-04-20T02:24:37.981143Z","shell.execute_reply":"2024-04-20T02:24:37.994188Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Coputing Key Points","metadata":{}},{"cell_type":"code","source":"\nprint(DEBUG)\nif 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\n    ","metadata":{"execution":{"iopub.status.busy":"2024-04-20T02:24:37.996236Z","iopub.execute_input":"2024-04-20T02:24:37.996513Z","iopub.status.idle":"2024-04-20T02:24:40.449325Z","shell.execute_reply.started":"2024-04-20T02:24:37.996474Z","shell.execute_reply":"2024-04-20T02:24:40.448377Z"},"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\") # JAYEONG cpu -> cuda\n) -> None: \n    \"\"\"\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    print(\"device\", device)\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            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()\n                    ","metadata":{"execution":{"iopub.status.busy":"2024-04-20T02:24:40.450489Z","iopub.execute_input":"2024-04-20T02:24:40.450784Z","iopub.status.idle":"2024-04-20T02:24:40.460384Z","shell.execute_reply.started":"2024-04-20T02:24:40.45076Z","shell.execute_reply":"2024-04-20T02:24:40.459503Z"},"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-04-20T02:24:40.461512Z","iopub.execute_input":"2024-04-20T02:24:40.461787Z","iopub.status.idle":"2024-04-20T02:25:08.831673Z","shell.execute_reply.started":"2024-04-20T02:24:40.461763Z","shell.execute_reply":"2024-04-20T02:25:08.830746Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Match and Compute keypoint distances","metadata":{}},{"cell_type":"code","source":"\n\nif 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        \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        \n        descriptors1= torch.from_numpy(f_descriptors[key1][...]).to(device)\n        descriptors2= torch.from_numpy(f_descriptors[key2][...]).to(device)\n        print(\"Descriptor: \", 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-04-20T02:25:08.833142Z","iopub.execute_input":"2024-04-20T02:25:08.833943Z","iopub.status.idle":"2024-04-20T02:25:09.46864Z","shell.execute_reply.started":"2024-04-20T02:25:08.833907Z","shell.execute_reply":"2024-04-20T02:25:09.467707Z"},"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(\"cuda\"),  # JAYEONG cup -> cuda\n) -> None:\n    \"\"\"\n    Computes distances between keypoints of images\n    Stores output at feature_dir/matches.h5\n    \"\"\"\n    match_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))\n\n                                ","metadata":{"execution":{"iopub.status.busy":"2024-04-20T02:25:09.469915Z","iopub.execute_input":"2024-04-20T02:25:09.470202Z","iopub.status.idle":"2024-04-20T02:25:09.48228Z","shell.execute_reply.started":"2024-04-20T02:25:09.470175Z","shell.execute_reply":"2024-04-20T02:25:09.48135Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(images_list, index_pairs,feature_dir)\nif DEBUG:\n    keypoint_distances(images_list, index_pairs, feature_dir, verbose=False)","metadata":{"execution":{"iopub.status.busy":"2024-04-20T02:25:09.483364Z","iopub.execute_input":"2024-04-20T02:25:09.483658Z","iopub.status.idle":"2024-04-20T02:25:19.60423Z","shell.execute_reply.started":"2024-04-20T02:25:09.483633Z","shell.execute_reply":"2024-04-20T02:25:19.603301Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 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=pycolmap.Database.connect\n    db = COLMAPDatabase.connect(database_path)\n#     db.close()\n#     db.cursor.close()\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-04-20T02:25:19.607776Z","iopub.execute_input":"2024-04-20T02:25:19.608135Z","iopub.status.idle":"2024-04-20T02:25:19.614153Z","shell.execute_reply.started":"2024-04-20T02:25:19.608108Z","shell.execute_reply":"2024-04-20T02:25:19.613145Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install pycolmap\nimport pycolmap","metadata":{"execution":{"iopub.status.busy":"2024-04-20T02:25:19.615402Z","iopub.execute_input":"2024-04-20T02:25:19.615745Z","iopub.status.idle":"2024-04-20T02:25:32.559193Z","shell.execute_reply.started":"2024-04-20T02:25:19.615715Z","shell.execute_reply":"2024-04-20T02:25:32.558167Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"database_path = 'colmap003.db' # 이전에 사용한 db 다시 사용할때 transaction - db is locked 문제 발생함. initializing 해줘야함\n\nif 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-04-20T02:25:32.560692Z","iopub.execute_input":"2024-04-20T02:25:32.561Z","iopub.status.idle":"2024-04-20T02:25:32.952692Z","shell.execute_reply.started":"2024-04-20T02:25:32.560971Z","shell.execute_reply":"2024-04-20T02:25:32.951778Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if DEBUG:\n    pycolmap.match_exhaustive(database_path)","metadata":{"execution":{"iopub.status.busy":"2024-04-20T02:25:32.953818Z","iopub.execute_input":"2024-04-20T02:25:32.954097Z","iopub.status.idle":"2024-04-20T02:25:33.632969Z","shell.execute_reply.started":"2024-04-20T02:25:32.954073Z","shell.execute_reply":"2024-04-20T02:25:33.63196Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Sparse Reconstruction","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":{"execution":{"iopub.status.busy":"2024-04-20T02:25:33.634129Z","iopub.execute_input":"2024-04-20T02:25:33.634409Z","iopub.status.idle":"2024-04-20T02:25:39.09905Z","shell.execute_reply.started":"2024-04-20T02:25:33.634386Z","shell.execute_reply":"2024-04-20T02:25:39.098033Z"},"collapsed":true,"jupyter":{"outputs_hidden":true},"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-04-20T02:25:39.100152Z","iopub.execute_input":"2024-04-20T02:25:39.100444Z","iopub.status.idle":"2024-04-20T02:25:39.108938Z","shell.execute_reply.started":"2024-04-20T02:25:39.100414Z","shell.execute_reply":"2024-04-20T02:25:39.1079Z"},"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    \"\"\"\n    Construct a dict describing the test data as\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: \", 1)\n            if l and i > 0:\n                image_path, dataset, scene, _, _ = l.strip().split(',')\n                \n                if dataset not in data_dict:\n                    data_dict[dataset] = {}\n                    \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-04-20T02:25:39.110131Z","iopub.execute_input":"2024-04-20T02:25:39.11039Z","iopub.status.idle":"2024-04-20T02:25:39.120541Z","shell.execute_reply.started":"2024-04-20T02:25:39.110367Z","shell.execute_reply":"2024-04-20T02:25:39.119696Z"},"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,\n        \"max_num_models\" : 2,\n    }","metadata":{"execution":{"iopub.status.busy":"2024-04-20T02:25:39.121687Z","iopub.execute_input":"2024-04-20T02:25:39.122068Z","iopub.status.idle":"2024-04-20T02:25:39.135971Z","shell.execute_reply.started":"2024-04-20T02:25:39.122017Z","shell.execute_reply":"2024-04-20T02:25:39.135098Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"database_path = 'colmap002.db'\n\ndef 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            result[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 / database_path # JAYEONG\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 matched 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 trianulating 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                # pycolamp 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                \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                # Parese 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, cofnig.base_path)\n                gc.collect()\n                \n            except Exception as e:    \n                print(e)\n               \n                \n               ","metadata":{"execution":{"iopub.status.busy":"2024-04-20T02:25:39.137125Z","iopub.execute_input":"2024-04-20T02:25:39.13768Z","iopub.status.idle":"2024-04-20T02:25:39.158211Z","shell.execute_reply.started":"2024-04-20T02:25:39.137648Z","shell.execute_reply":"2024-04-20T02:25:39.157036Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"global_database_path = 'colmap002.db'\n\ndef 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 / global_database_path # JAYEONG\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 matched into colmap for RANSAC\n                import_into_colmap(\n                    images_dir,\n                    feature_dir,\n                    database_path,\n                )\n                output_path = feature_dir / \"colmap_rec_aliked\"\n                output_path.mkdir(parents=True, exist_ok=True)\n                \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                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 trianulating 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                # pycolamp 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                 # Parese 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, cofnig.base_path)\n                gc.collect()   \n                    \n            except Exception as e:    \n                print(f\"===Error===\", e, \"======\")\n               \n                ","metadata":{"execution":{"iopub.status.busy":"2024-04-20T02:37:26.49993Z","iopub.execute_input":"2024-04-20T02:37:26.500306Z","iopub.status.idle":"2024-04-20T02:37:26.516128Z","shell.execute_reply.started":"2024-04-20T02:37:26.500261Z","shell.execute_reply":"2024-04-20T02:37:26.515056Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"run_from_config(Config)","metadata":{"execution":{"iopub.status.busy":"2024-04-20T02:37:29.344249Z","iopub.execute_input":"2024-04-20T02:37:29.345144Z"},"trusted":true},"execution_count":null,"outputs":[]}]}