{"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":8539410,"sourceType":"datasetVersion","datasetId":5101247},{"sourceId":8560092,"sourceType":"datasetVersion","datasetId":5075608}],"dockerImageVersionId":30699,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Install dependency","metadata":{}},{"cell_type":"code","source":"!pip install --no-index /kaggle/input/imc2024-packages-lightglue-rerun-kornia/* --no-deps","metadata":{"execution":{"iopub.status.busy":"2024-05-30T14:42:25.813850Z","iopub.execute_input":"2024-05-30T14:42:25.814439Z","iopub.status.idle":"2024-05-30T14:42:30.210580Z","shell.execute_reply.started":"2024-05-30T14:42:25.814409Z","shell.execute_reply":"2024-05-30T14:42:30.209217Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!mkdir -p /root/.cache/torch\n!cp -r /kaggle/input/hierarchical-localization/hub /root/.cache/torch\n!cp -r /kaggle/input/sfd2-semantic-guidedfeaturedetectionanddescription/weights ./","metadata":{"execution":{"iopub.status.busy":"2024-05-30T14:42:30.212985Z","iopub.execute_input":"2024-05-30T14:42:30.213815Z","iopub.status.idle":"2024-05-30T14:42:51.492899Z","shell.execute_reply.started":"2024-05-30T14:42:30.213779Z","shell.execute_reply":"2024-05-30T14:42:51.491612Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import sys\nsys.path.append(\"/kaggle/input/hierarchical-localization\")\nsys.path.append(\"/kaggle/input/hierarchical-localization/third_party\")\nsys.path.append(\"/kaggle/input/sfd2-semantic-guidedfeaturedetectionanddescription\")","metadata":{"execution":{"iopub.status.busy":"2024-05-30T14:42:51.494579Z","iopub.execute_input":"2024-05-30T14:42:51.494915Z","iopub.status.idle":"2024-05-30T14:42:51.500141Z","shell.execute_reply.started":"2024-05-30T14:42:51.494882Z","shell.execute_reply":"2024-05-30T14:42:51.499279Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import tqdm, tqdm.notebook\ntqdm.tqdm = tqdm.notebook.tqdm  # notebook-friendly progress bars\nfrom pathlib import Path\n\nfrom hloc import (\n    pairs_from_exhaustive, pairs_from_retrieval, \n    extract_features, \n    match_features, match_dense, \n    reconstruction, \n    visualization, \n)\nfrom hloc.visualization import plot_images, read_image\nfrom hloc.utils import viz_3d\n\nimport extract_localization as extract_features_sfd2\n\nimport pycolmap\nimport gc\nimport numpy as np\nfrom copy import deepcopy","metadata":{"execution":{"iopub.status.busy":"2024-05-30T14:42:51.502879Z","iopub.execute_input":"2024-05-30T14:42:51.503500Z","iopub.status.idle":"2024-05-30T14:42:59.681039Z","shell.execute_reply.started":"2024-05-30T14:42:51.503469Z","shell.execute_reply":"2024-05-30T14:42:59.680020Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Preprocess","metadata":{}},{"cell_type":"code","source":"# Rotation Correction, Color Correction, etc.","metadata":{"execution":{"iopub.status.busy":"2024-05-30T14:42:59.682577Z","iopub.execute_input":"2024-05-30T14:42:59.682991Z","iopub.status.idle":"2024-05-30T14:42:59.687273Z","shell.execute_reply.started":"2024-05-30T14:42:59.682963Z","shell.execute_reply":"2024-05-30T14:42:59.685933Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# SfM Pipeline","metadata":{}},{"cell_type":"code","source":"config = {\n    \"retrieval_conf\": \"netvlad\", # {\"dir\", \"netvlad\", \"openibl\", \"eigenplaces\"}\n    \"feature_conf\": \"aliked-n16-rot\", # {\"superpoint_aachen\", \"superpoint_max\", \"superpoint_inloc\", \"r2d2\", \"d2net-ss\", \"sift\", \"sosnet\", \"disk\", \"aliked-n16-rot\", \"aliked-n32\"}\n    \"feature_conf_sfd2\": \"ressegnetv2-20220810-wapv2-sd2mfsf-uspg-0001-n4096-r1024\", \n    \"match_method\": \"dense\", # {\"feature\", \"dense\"}\n    \"matcher_conf_feature\": \"aliked+lightglue\", # {\"superpoint+lightglue\", \"disk+lightglue\", \"superglue\", \"superglue-fast\", \"NN-superpoint\", \"NN-ratio\", \"NN-mutual\", \"adalam\", \"aliked+lightglue\"}\n    \"matcher_conf_dense\": \"loftr\", # {\"loftr\", \"loftr_aachen\", \"loftr_superpoint\"}\n}","metadata":{"execution":{"iopub.status.busy":"2024-05-30T14:42:59.689422Z","iopub.execute_input":"2024-05-30T14:42:59.690122Z","iopub.status.idle":"2024-05-30T14:43:01.394113Z","shell.execute_reply.started":"2024-05-30T14:42:59.690089Z","shell.execute_reply":"2024-05-30T14:43:01.393072Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"extract_features.confs[config[\"retrieval_conf\"]]","metadata":{"execution":{"iopub.status.busy":"2024-05-30T14:43:01.395426Z","iopub.execute_input":"2024-05-30T14:43:01.396237Z","iopub.status.idle":"2024-05-30T14:43:01.454691Z","shell.execute_reply.started":"2024-05-30T14:43:01.396204Z","shell.execute_reply":"2024-05-30T14:43:01.453826Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"extract_features.confs[config[\"feature_conf\"]]","metadata":{"execution":{"iopub.status.busy":"2024-05-30T14:43:01.456260Z","iopub.execute_input":"2024-05-30T14:43:01.456626Z","iopub.status.idle":"2024-05-30T14:43:01.528153Z","shell.execute_reply.started":"2024-05-30T14:43:01.456595Z","shell.execute_reply":"2024-05-30T14:43:01.527271Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"match_features.confs[config[\"matcher_conf_feature\"]]","metadata":{"execution":{"iopub.status.busy":"2024-05-30T14:43:01.529309Z","iopub.execute_input":"2024-05-30T14:43:01.531562Z","iopub.status.idle":"2024-05-30T14:43:01.590105Z","shell.execute_reply.started":"2024-05-30T14:43:01.531530Z","shell.execute_reply":"2024-05-30T14:43:01.589215Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"match_dense.confs[config[\"matcher_conf_dense\"]]","metadata":{"execution":{"iopub.status.busy":"2024-05-30T14:43:01.593936Z","iopub.execute_input":"2024-05-30T14:43:01.594237Z","iopub.status.idle":"2024-05-30T14:43:01.657775Z","shell.execute_reply.started":"2024-05-30T14:43:01.594205Z","shell.execute_reply":"2024-05-30T14:43:01.656815Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def sfm_pipeline(images: Path, outputs: Path, config: dict):\n    \"\"\" Pipeline of Structure from Motion \"\"\"\n    \n    # Set configures\n    # Output paths\n    sfm_pairs = outputs / str(\"pairs-\" + config[\"retrieval_conf\"] + \".txt\")\n    sfm_dir = outputs / config[\"feature_conf\"]\n    # Global descriptors : {\"dir\", \"netvlad\", \"openibl\", \"eigenplaces\"}\n    retrieval_conf = extract_features.confs[config[\"retrieval_conf\"]]\n    # Local descriptors : {\"superpoint_aachen\", \"superpoint_max\", \"superpoint_inloc\", \"r2d2\", \"d2net-ss\", \"sift\", \"sosnet\", \"disk\", \"aliked-n16-rot\", \"aliked-n32\"}\n    feature_conf = extract_features.confs[config[\"feature_conf\"]]\n    # feature_conf['model']['detection_threshold'] = 0.1\n    # feature_conf['model']['nms_radius'] = 1\n    \n    # SFD2: Semantic-guided Feature Detection and Description\n    # feature_conf = extract_features_sfd2.confs[config[\"feature_conf_sfd2\"]]\n    \n    # Get image pairs\n    if len(image_paths) <= -1:\n        print(\"Get exhaustive pairs.\")\n        outputs.mkdir(parents=True, exist_ok=True)\n        sfm_pairs.touch(exist_ok=True)\n        pairs_from_exhaustive.main(sfm_pairs, image_list=[p.relative_to(images).as_posix() for p in (images).iterdir()])\n        \n        # retrieval_path = extract_features.main(retrieval_conf, images, outputs)\n        # pairs_from_exhaustive.main(sfm_pairs, features=retrieval_path)\n    else:\n        print(\"Get retrieval pairs.\")\n        retrieval_path = extract_features.main(retrieval_conf, images, outputs)\n        pairs_from_retrieval.main(retrieval_path, sfm_pairs, num_matched=10)\n    \n    # Match\n    if(config[\"match_method\"]==\"dense\"):\n        # Match : {\"loftr\", \"loftr_aachen\", \"loftr_superpoint\"}\n        matcher_conf = match_dense.confs[config[\"matcher_conf_dense\"]]\n        # Match dense\n        feature_path, match_path = match_dense.main(matcher_conf, sfm_pairs, images, export_dir=outputs)\n        \n    else:\n        # Match : {\"superpoint+lightglue\", \"disk+lightglue\", \"superglue\", \"superglue-fast\", \"NN-superpoint\", \"NN-ratio\", \"NN-mutual\", \"adalam\", \"aliked+lightglue\"}\n        matcher_conf = match_features.confs[config[\"matcher_conf_feature\"]]\n        matcher_conf['max_error'] = 2\n        matcher_conf['cell_size'] = 6\n        \n        # \"outdoor\" is better than \"indoor\"\n        # if((config[\"matcher_conf_feature\"]==\"superglue\") & (config[\"feature_conf\"]==\"superpoint_inloc\")):\n        #    matcher_conf['model']['weights'] = \"indoor\"\n        \n        # Extract features and match\n        feature_path = extract_features.main(feature_conf, images, outputs)\n        # feature_path = extract_features_sfd2.main(feature_conf, images, outputs)\n        match_path = match_features.main(matcher_conf, sfm_pairs, feature_conf[\"output\"], outputs)\n    \n    # Reconstruct\n    model = reconstruction.main(sfm_dir, images, sfm_pairs, feature_path, match_path,\n                                # camera_mode=pycolmap.CameraMode.SINGLE, 　# default : pycolmap.CameraMode.AUTO\n                                # skip_geometric_verification=True,\n                                # min_match_score=0.5,\n                                # image_options=dict(camera_model=\"SIMPLE_PINHOLE\",\n                                #                   ),\n                                # mapper_options=dict(min_num_matches=5,\n                                #                     max_num_models=3,\n                                #                     min_model_size=2,\n                                #                     mapper=dict(max_reg_trials=10,\n                                #                                init_min_tri_angle=4,\n                                #                                abs_pose_min_num_inliers=5),\n                                #                     triangulation=dict(create_max_angle_error=10,\n                                #                                       re_max_trials=10,\n                                #                                       min_angle=0.1,\n                                #                                       ),\n                                #                    )\n                               )\n    # Save\n    model.write_text(outputs)\n    \n    return model\n","metadata":{"execution":{"iopub.status.busy":"2024-05-30T14:43:01.659574Z","iopub.execute_input":"2024-05-30T14:43:01.659937Z","iopub.status.idle":"2024-05-30T14:43:01.725736Z","shell.execute_reply.started":"2024-05-30T14:43:01.659906Z","shell.execute_reply":"2024-05-30T14:43:01.724699Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Functions for submission","metadata":{}},{"cell_type":"code","source":"def arr_to_str(a):\n    \"\"\"Returns ;-separated string representing the input\"\"\"\n    return \";\".join([str(x) for x in a.reshape(-1)])","metadata":{"execution":{"iopub.status.busy":"2024-05-30T14:43:01.726632Z","iopub.execute_input":"2024-05-30T14:43:01.726886Z","iopub.status.idle":"2024-05-30T14:43:01.826426Z","shell.execute_reply.started":"2024-05-30T14:43:01.726864Z","shell.execute_reply":"2024-05-30T14:43:01.825367Z"},"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-30T14:43:01.827633Z","iopub.execute_input":"2024-05-30T14:43:01.828046Z","iopub.status.idle":"2024-05-30T14:43:01.917122Z","shell.execute_reply.started":"2024-05-30T14:43:01.828016Z","shell.execute_reply":"2024-05-30T14:43:01.915436Z"},"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    \"\"\"Prepare a submission file\"\"\"\n    \n    with open(\"/kaggle/working/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-30T14:43:01.918277Z","iopub.execute_input":"2024-05-30T14:43:01.918662Z","iopub.status.idle":"2024-05-30T14:43:02.013340Z","shell.execute_reply.started":"2024-05-30T14:43:01.918626Z","shell.execute_reply":"2024-05-30T14:43:02.010853Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Main","metadata":{}},{"cell_type":"code","source":"base_path = Path(\"/kaggle/input/image-matching-challenge-2024\")\noutput_path = Path(\"/kaggle/working/outputs\")\n\nresults = {}\n\ndata_dict = parse_sample_submission(base_path)\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\"Got {len(image_paths)} images\")\n        \n        try:\n            images = Path(images_dir)\n            outputs = Path(output_path) / dataset / scene\n            \n            #### SfM pipline ####\n            model = sfm_pipeline(images=images, outputs=outputs, config=config)\n            \n            for k, im in model.images.items():\n                key = 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, base_path)\n            gc.collect()\n            \n        except Exception as e:\n            print(e)","metadata":{"execution":{"iopub.status.busy":"2024-05-30T14:43:02.014823Z","iopub.execute_input":"2024-05-30T14:43:02.015513Z","iopub.status.idle":"2024-05-30T14:48:58.611743Z","shell.execute_reply.started":"2024-05-30T14:43:02.015471Z","shell.execute_reply":"2024-05-30T14:48:58.610727Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Visualize","metadata":{}},{"cell_type":"code","source":"fig = viz_3d.init_figure()\nviz_3d.plot_reconstruction(fig, model, color='rgba(255,0,0,0.5)', name=\"mapping\", points_rgb=True)\nfig.show()","metadata":{"execution":{"iopub.status.busy":"2024-05-30T14:48:58.613106Z","iopub.execute_input":"2024-05-30T14:48:58.613561Z","iopub.status.idle":"2024-05-30T14:49:02.028485Z","shell.execute_reply.started":"2024-05-30T14:48:58.613533Z","shell.execute_reply":"2024-05-30T14:49:02.027536Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"visualization.visualize_sfm_2d(model, images, color_by='visibility', n=5)","metadata":{"execution":{"iopub.status.busy":"2024-05-30T14:49:02.029685Z","iopub.execute_input":"2024-05-30T14:49:02.029971Z","iopub.status.idle":"2024-05-30T14:49:04.354367Z","shell.execute_reply.started":"2024-05-30T14:49:02.029946Z","shell.execute_reply":"2024-05-30T14:49:04.353384Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"visualization.visualize_sfm_2d(model, images, color_by=\"track_length\", n=5)","metadata":{"execution":{"iopub.status.busy":"2024-05-30T14:49:04.355497Z","iopub.execute_input":"2024-05-30T14:49:04.355760Z","iopub.status.idle":"2024-05-30T14:49:07.114585Z","shell.execute_reply.started":"2024-05-30T14:49:04.355737Z","shell.execute_reply":"2024-05-30T14:49:07.113631Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"visualization.visualize_sfm_2d(model, images, color_by=\"depth\", n=5)","metadata":{"execution":{"iopub.status.busy":"2024-05-30T14:49:07.115871Z","iopub.execute_input":"2024-05-30T14:49:07.116181Z","iopub.status.idle":"2024-05-30T14:49:09.704568Z","shell.execute_reply.started":"2024-05-30T14:49:07.116155Z","shell.execute_reply":"2024-05-30T14:49:09.703609Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# import h5py\n# from hloc.utils.read_write_model import read_model\n\n# cameras, images, points3D = read_model('/kaggle/working/outputs/church/church/aliked-n16-rot', ext='.bin')\n# features = h5py.File(\"/kaggle/working/outputs/church/church/feats-aliked-n16-rot.h5\", \"r\")\n# print(len(images))\n# for id_ in images:\n#     n_model = len(images[id_].point3D_ids)\n#     name = images[id_].name\n#     n_feats = features[name]['keypoints'].__array__().shape[0]\n#     print(name, id_, n_model, n_feats)","metadata":{"execution":{"iopub.status.busy":"2024-05-30T14:49:09.705827Z","iopub.execute_input":"2024-05-30T14:49:09.706193Z","iopub.status.idle":"2024-05-30T14:49:09.710994Z","shell.execute_reply.started":"2024-05-30T14:49:09.706164Z","shell.execute_reply":"2024-05-30T14:49:09.710080Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# matches = h5py.File(\"/kaggle/working/outputs/church/church/feats-aliked-n16-rot_matches-aliked-lightglue_pairs-netvlad.h5\", \"r\")\n\n# for id_ in images:\n#     name = images[id_].name\n#     try:\n#         for matched in matches[name]:\n#             n_matches = matches[name][matched]['matches0'].__array__().shape[0]\n#             score_matches = matches[name][matched]['matching_scores0'].__array__().mean()\n#             print(id_, name, matched, n_matches, score_matches)\n#     except:\n#         print('error', name)","metadata":{"execution":{"iopub.status.busy":"2024-05-30T14:49:09.712304Z","iopub.execute_input":"2024-05-30T14:49:09.713016Z","iopub.status.idle":"2024-05-30T14:49:09.723890Z","shell.execute_reply.started":"2024-05-30T14:49:09.712976Z","shell.execute_reply":"2024-05-30T14:49:09.722948Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}