{"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":8581151,"sourceType":"datasetVersion","datasetId":5016841},{"sourceId":170460983,"sourceType":"kernelVersion"},{"sourceId":174014325,"sourceType":"kernelVersion"},{"sourceId":179091099,"sourceType":"kernelVersion"},{"sourceId":179566701,"sourceType":"kernelVersion"},{"sourceId":179656770,"sourceType":"kernelVersion"},{"sourceId":3840,"sourceType":"modelInstanceVersion","modelInstanceId":2742},{"sourceId":17191,"sourceType":"modelInstanceVersion","modelInstanceId":14317},{"sourceId":17555,"sourceType":"modelInstanceVersion","modelInstanceId":14611}],"dockerImageVersionId":30674,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"baseline, exp521, min_matches=[100,100], val=0.41, lb=0.18\n\nversion5, exp556, use_all_images=100, val=0.xx, lb=0.\n\nversion7, exp557, min_matches=[100, 125], val=0.45, lb=0.19\n\nversion9, exp557+alpha, min_matches=[100,125,150]","metadata":{}},{"cell_type":"code","source":"DEBUG = False\n\nimport os\nif os.getenv('KAGGLE_IS_COMPETITION_RERUN'):\n    DEBUG = False\n\nSAVE_FLAG = (not os.getenv('KAGGLE_IS_COMPETITION_RERUN') and not DEBUG)","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:56:07.324537Z","iopub.execute_input":"2024-06-03T09:56:07.324799Z","iopub.status.idle":"2024-06-03T09:56:07.335288Z","shell.execute_reply.started":"2024-06-03T09:56:07.324773Z","shell.execute_reply":"2024-06-03T09:56:07.334306Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if DEBUG:\n    DATA_TYPE = \"train\"\n    #SCENE_TYPE = \"transp_obj_glass_cylinder\"\n    SCENE_TYPE = \"dioscuri\"\nelse:\n    DATA_TYPE = \"test\"\n    SCENE_TYPE = \"submission\"","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:56:07.336739Z","iopub.execute_input":"2024-06-03T09:56:07.337277Z","iopub.status.idle":"2024-06-03T09:56:07.346409Z","shell.execute_reply.started":"2024-06-03T09:56:07.337248Z","shell.execute_reply":"2024-06-03T09:56:07.345447Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from pathlib import Path\noutput_dir = Path(\"/kaggle/working/\")","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:56:07.347461Z","iopub.execute_input":"2024-06-03T09:56:07.347714Z","iopub.status.idle":"2024-06-03T09:56:07.357902Z","shell.execute_reply.started":"2024-06-03T09:56:07.347692Z","shell.execute_reply":"2024-06-03T09:56:07.356229Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# COPY models","metadata":{}},{"cell_type":"code","source":"!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\n!cp /kaggle/input/check-orientation/2020-11-16_resnext50_32x4d.zip /root/.cache/torch/hub/checkpoints/\n!cp /kaggle/input/dinov2-repo/dinov2_vits14_pretrain.pth /root/.cache/torch/hub/checkpoints/\n!cp /kaggle/input/dinov2-repo/dinov2_vits14_voc2012_linear_head.pth /root/.cache/torch/hub/checkpoints/","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:56:07.360615Z","iopub.execute_input":"2024-06-03T09:56:07.360860Z","iopub.status.idle":"2024-06-03T09:56:17.173306Z","shell.execute_reply.started":"2024-06-03T09:56:07.360838Z","shell.execute_reply":"2024-06-03T09:56:17.171946Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import argparse\nimport cv2\nimport gc\nimport h5py\nimport importlib\nimport itertools\nimport kornia as K\nimport kornia.feature as KF\nimport math\nimport mmcv\nimport numpy as np\nimport pandas as pd\nimport pycolmap\nimport random\nimport shutil\nimport sys\nimport torch\nimport torch.nn.functional as F\n\nfrom check_orientation.pre_trained_models import create_model\nfrom concurrent.futures import ThreadPoolExecutor\nfrom copy import deepcopy\nfrom functools import partial\nfrom kneed import KneeLocator\nfrom lightglue import ALIKED\nfrom mmcv.runner import load_checkpoint\nfrom mmseg.apis import init_segmentor, inference_segmentor\nfrom sklearn.cluster import DBSCAN\nfrom sklearn.neighbors import NearestNeighbors\nfrom torch.backends import cudnn\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision.io import read_image as T_read_image\nfrom torchvision.io import ImageReadMode\nfrom torchvision import transforms as T\nfrom tqdm.auto import tqdm","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:56:17.177995Z","iopub.execute_input":"2024-06-03T09:56:17.178304Z","iopub.status.idle":"2024-06-03T09:57:05.924760Z","shell.execute_reply.started":"2024-06-03T09:56:17.178268Z","shell.execute_reply":"2024-06-03T09:57:05.923776Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# DINOv2 Segmenter\n\nhttps://github.com/facebookresearch/dinov2","metadata":{}},{"cell_type":"code","source":"sys.path.append('/kaggle/input/dinov2-repo/dinov2')\nimport dinov2.eval.segmentation.models","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:05.925924Z","iopub.execute_input":"2024-06-03T09:57:05.926486Z","iopub.status.idle":"2024-06-03T09:57:06.092283Z","shell.execute_reply.started":"2024-06-03T09:57:05.926459Z","shell.execute_reply":"2024-06-03T09:57:06.091360Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CenterPadding(torch.nn.Module):\n    def __init__(self, multiple):\n        super().__init__()\n        self.multiple = multiple\n\n    def _get_pad(self, size):\n        new_size = math.ceil(size / self.multiple) * self.multiple\n        pad_size = new_size - size\n        pad_size_left = pad_size // 2\n        pad_size_right = pad_size - pad_size_left\n        return pad_size_left, pad_size_right\n\n    @torch.inference_mode()\n    def forward(self, x):\n        pads = list(itertools.chain.from_iterable(self._get_pad(m) for m in x.shape[:1:-1]))\n        output = F.pad(x, pads)\n        return output\n\ndef create_segmenter(cfg, backbone_model):\n    model = init_segmentor(cfg)\n    model.backbone.forward = partial(\n        backbone_model.get_intermediate_layers,\n        n=cfg.model.backbone.out_indices,\n        reshape=True,\n    )\n    if hasattr(backbone_model, \"patch_size\"):\n        model.backbone.register_forward_pre_hook(lambda _, x: CenterPadding(backbone_model.patch_size)(x[0]))\n    model.init_weights()\n    return model\n\ndef dinov2_segmentation(\n    paths,\n    feature_dir,\n    device_id,\n):\n\n    device = torch.device(f\"cuda:{device_id}\")\n\n    backbone_model = torch.hub.load('/kaggle/input/dinov2-repo/dinov2', model=\"dinov2_vits14\", source='local')\n    backbone_model.eval()\n    backbone_model.to(device)\n\n    #head_config_url = \"https://dl.fbaipublicfiles.com/dinov2/dinov2_vits14/dinov2_vits14_voc2012_linear_config.py\"\n    cfg = mmcv.Config.fromfile(\"/kaggle/input/dinov2-repo/dinov2_vits14_voc2012_linear_config.py\")\n\n    model = create_segmenter(cfg, backbone_model=backbone_model)\n    #head_checkpoint_url = \"https://dl.fbaipublicfiles.com/dinov2/dinov2_vits14/dinov2_vits14_voc2012_linear_head.pth\"\n    load_checkpoint(model, \"/kaggle/input/dinov2-repo/dinov2_vits14_voc2012_linear_head.pth\", map_location=\"cpu\")\n    model.to(device)\n    model.eval()\n\n    with h5py.File(feature_dir / f\"fg_mask{device_id}.h5\", mode=\"w\") as f_fg_mask:\n        for img_path in tqdm(paths, dynamic_ncols=True, desc=f\"Segmenting[GPU{device_id}]\"):\n            image = cv2.imread(str(img_path))\n\n            segmentation_logits = inference_segmentor(model, image)[0]\n            mask = np.zeros_like(segmentation_logits)\n            mask[segmentation_logits == 5] = 255\n\n            key = img_path.name\n            f_fg_mask[key] = mask\n\n    return","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:06.093999Z","iopub.execute_input":"2024-06-03T09:57:06.094398Z","iopub.status.idle":"2024-06-03T09:57:06.110403Z","shell.execute_reply.started":"2024-06-03T09:57:06.094361Z","shell.execute_reply":"2024-06-03T09:57:06.109593Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def merge_single_h5(\n    feature_dir,\n    input_name_list,\n    output_name,\n):\n    with h5py.File(feature_dir / output_name, mode=\"w\") as f_output:\n\n        for input_name in input_name_list:\n            with h5py.File(feature_dir / input_name, mode=\"r\") as f_input:\n                for key in f_input.keys():\n                    f_output[key] = f_input[key][...]\n\n    return\n\ndef merge_double_h5(\n    feature_dir,\n    input_name_list,\n    output_name,\n):\n    with h5py.File(feature_dir / output_name, mode=\"w\") as f_output:\n\n        for input_matches_name in input_name_list:\n            with h5py.File(feature_dir / input_matches_name, mode=\"r\") as f_input:\n                for key1 in f_input.keys():\n                    group  = f_output.require_group(key1)\n                    for key2 in f_input[key1].keys():\n                        group.create_dataset(key2, data=f_input[key1][key2][...])\n\n    return","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:06.117216Z","iopub.execute_input":"2024-06-03T09:57:06.117565Z","iopub.status.idle":"2024-06-03T09:57:06.126139Z","shell.execute_reply.started":"2024-06-03T09:57:06.117539Z","shell.execute_reply":"2024-06-03T09:57:06.125172Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Detect keypoints by ALIKED","metadata":{}},{"cell_type":"code","source":"def pad_to_square(image):\n    if len(image.shape) == 3:\n        height, width, _ = image.shape\n        max_dim = max(height, width)\n        padded_image = np.zeros((max_dim, max_dim, 3), dtype=np.uint8)\n        padded_image[:height, :width, :] = image\n    elif len(image.shape) == 2:\n        height, width = image.shape\n        max_dim = max(height, width)\n        padded_image = np.zeros((max_dim, max_dim), dtype=np.uint8)\n        padded_image[:height, :width] = image\n    else:\n        raise ValueError(\"Invalid image shape\")\n    return padded_image","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:06.127186Z","iopub.execute_input":"2024-06-03T09:57:06.127462Z","iopub.status.idle":"2024-06-03T09:57:06.140242Z","shell.execute_reply.started":"2024-06-03T09:57:06.127439Z","shell.execute_reply":"2024-06-03T09:57:06.139370Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## rotate image 90 degrees clockwise","metadata":{}},{"cell_type":"code","source":"def detect_keypoints(\n    paths,\n    feature_dir,\n    device_id,\n    num_features,\n    detection_threshold,\n    resize_to,\n):\n    device = torch.device(f\"cuda:{device_id}\")\n\n    dtype = torch.float32 # ALIKED has issues with float16\n\n    extractor = ALIKED(\n        max_num_keypoints=num_features,\n        detection_threshold=detection_threshold,\n        resize=resize_to,\n    ).eval().to(device, dtype)\n\n    with h5py.File(feature_dir / f\"keypoints_deg{device_id}.h5\", mode=\"w\") as f_keypoints_deg, \\\n         h5py.File(feature_dir / f\"descriptors_deg{device_id}.h5\", mode=\"w\") as f_descriptors_deg, \\\n         h5py.File(feature_dir / f\"offsets{device_id}.h5\", mode=\"w\") as f_offsets_deg, \\\n         h5py.File(feature_dir / f\"keypoints{device_id}.h5\", mode=\"w\") as f_keypoints:\n\n        for path in tqdm(paths, desc=f\"Computing keypoints[GPU{device_id}]\", dynamic_ncols=True):\n\n            _image = cv2.imread(str(path))\n            _image = pad_to_square(_image)\n\n            key = path.name\n\n            f_keypoints_deg.create_group(key)\n            f_descriptors_deg.create_group(key)\n            f_offsets_deg.create_group(key)\n\n            keypoints_list = []\n            offset = 0\n            for deg in [\"0deg\", \"90deg\", \"180deg\", \"270deg\"]:\n                if deg != \"0deg\":\n                    # rotate 90 degrees\n                    _image = cv2.rotate(_image, cv2.ROTATE_90_CLOCKWISE)\n                image = K.image_to_tensor(_image, False).float() / 255.\n                image = K.color.bgr_to_rgb(image).to(device).to(dtype)\n\n                with torch.inference_mode():\n                    features = extractor.extract(image)\n\n                keypoints = features[\"keypoints\"].squeeze().detach().cpu().numpy()\n                f_keypoints_deg[key][deg] = keypoints\n                f_descriptors_deg[key][deg] = features[\"descriptors\"].squeeze().detach().cpu().numpy()\n                f_offsets_deg[key][deg] = offset\n                offset += keypoints.shape[0]\n\n                # rotate back the keypoints\n                temp = keypoints.copy()\n                if deg == \"90deg\":\n                    keypoints[:, 0] = temp[:, 1]\n                    keypoints[:, 1] = _image.shape[1] - temp[:, 0]\n                elif deg == \"180deg\":\n                    keypoints[:, 0] = _image.shape[1] - temp[:, 0]\n                    keypoints[:, 1] = _image.shape[0] - temp[:, 1]\n                elif deg == \"270deg\":\n                    keypoints[:, 0] = _image.shape[0] - temp[:, 1]\n                    keypoints[:, 1] = temp[:, 0]\n\n                keypoints_list.append(keypoints)\n\n            f_keypoints[key] = np.concatenate(keypoints_list)\n\n    return","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:06.141560Z","iopub.execute_input":"2024-06-03T09:57:06.141883Z","iopub.status.idle":"2024-06-03T09:57:06.159303Z","shell.execute_reply.started":"2024-06-03T09:57:06.141848Z","shell.execute_reply":"2024-06-03T09:57:06.158548Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## special function for transparent object","metadata":{}},{"cell_type":"code","source":"def detect_keypoints_transparent(\n    paths,\n    feature_dir,\n    device_id,\n    num_features,\n    detection_threshold,\n    resize_to,\n):\n    device = torch.device(f\"cuda:{device_id}\")\n\n    dtype = torch.float32 # ALIKED has issues with float16\n\n    extractor = ALIKED(\n        max_num_keypoints=num_features,\n        detection_threshold=detection_threshold,\n        resize=resize_to,\n    ).eval().to(device, dtype)\n\n    num_grids = 0\n\n    def mask_keypoints(mask_area, keypoints):\n        fg_idx = []\n        for kpt in keypoints.astype(int):\n            flag = (mask_area[kpt[1], kpt[0]] > 0)\n            fg_idx.append(flag)\n        return np.array(fg_idx)\n\n    with h5py.File(feature_dir / \"fg_mask.h5\", mode=\"r\") as f_fg_mask, \\\n         h5py.File(feature_dir / f\"descriptors_grid{device_id}.h5\", mode=\"w\") as f_descriptors_grid, \\\n         h5py.File(feature_dir / f\"keypoints_grid{device_id}.h5\", mode=\"w\") as f_keypoints_grid, \\\n         h5py.File(feature_dir / f\"offsets_grid{device_id}.h5\", mode=\"w\") as f_offsets_grid, \\\n         h5py.File(feature_dir / f\"keypoints{device_id}.h5\", mode=\"w\") as f_keypoints:\n\n        for path in tqdm(paths, desc=f\"Computing keypoints[GPU{device_id}]\", dynamic_ncols=True):\n\n            _image = cv2.imread(str(path))\n            _mask = f_fg_mask[path.name][...]\n            height, width = _image.shape[:2]\n            step = 1024\n\n            key = path.name\n\n            f_descriptors_grid.create_group(key)\n            f_keypoints_grid.create_group(key)\n            f_offsets_grid.create_group(key)\n\n            grid_id = 0\n            offset = 0\n            keypoints_list = []\n            for y in range(0, height, step):\n                for x in range(0, width, step):\n\n                    mask = _mask[y:y+step, x:x+step]\n                    if mask.sum() == 0:\n                        grid_id += 1\n                        continue\n\n                    image = K.image_to_tensor(_image[y:y+step, x:x+step], False).float() / 255.\n                    image = K.color.bgr_to_rgb(image).to(device).to(dtype)\n\n                    with torch.inference_mode():\n                        features = extractor.extract(image)\n\n                    keypoints = features[\"keypoints\"].detach().cpu().numpy().reshape(-1, 2)\n                    features = features[\"descriptors\"].detach().cpu().numpy().reshape(-1, 128)\n\n                    # mask out keypoints that are not in the foreground\n                    fg_idx = mask_keypoints(mask, keypoints)\n                    keypoints = keypoints[fg_idx]\n                    features = features[fg_idx]\n\n                    if len(keypoints) == 0:\n                        grid_id += 1\n                        continue\n\n                    f_keypoints_grid[key][f\"{grid_id}\"] = keypoints\n                    f_descriptors_grid[key][f\"{grid_id}\"] = features\n                    f_offsets_grid[key][f\"{grid_id}\"] = offset\n                    grid_id += 1\n\n                    offset += keypoints.shape[0]\n\n                    keypoints[:, 0] += x\n                    keypoints[:, 1] += y\n                    keypoints_list.append(keypoints)\n\n            f_keypoints[key] = np.concatenate(keypoints_list)\n\n            num_grids = max(num_grids, grid_id)\n\n    return num_grids","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:06.160664Z","iopub.execute_input":"2024-06-03T09:57:06.161013Z","iopub.status.idle":"2024-06-03T09:57:06.178693Z","shell.execute_reply.started":"2024-06-03T09:57:06.160982Z","shell.execute_reply":"2024-06-03T09:57:06.177959Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Match keypoints by LightGlue","metadata":{}},{"cell_type":"markdown","source":"## rotate image 90 degrees clockwise","metadata":{}},{"cell_type":"code","source":"def keypoint_distances(\n    paths,\n    index_pairs,\n    feature_dir,\n    device_id,\n    early_stopping_thr,\n    min_matches,\n    verbose,\n):\n    device = torch.device(f\"cuda:{device_id}\")\n\n    matcher_params = {\n        \"width_confidence\": -1,\n        \"depth_confidence\": -1,\n        \"mp\": True,\n    }\n    matcher = KF.LightGlueMatcher(\"aliked\", matcher_params).eval().to(device)\n\n    with h5py.File(feature_dir / \"keypoints_deg.h5\", mode=\"r\") as f_keypoints_deg, \\\n         h5py.File(feature_dir / \"descriptors_deg.h5\", mode=\"r\") as f_descriptors_deg, \\\n         h5py.File(feature_dir / \"offsets.h5\", mode=\"r\") as f_offsets_deg, \\\n         h5py.File(feature_dir / f\"matches{device_id}.h5\", mode=\"w\") as f_matches:\n\n            for idx1, idx2 in tqdm(index_pairs, desc=f\"Computing keypoint distances[GPU{device_id}]\", dynamic_ncols=True):\n                key1, key2 = paths[idx1].name, paths[idx2].name\n\n                keypoints1 = torch.from_numpy(f_keypoints_deg[key1][\"0deg\"][...]).to(device)\n                descriptors1 = torch.from_numpy(f_descriptors_deg[key1][\"0deg\"][...]).to(device)\n\n                best_deg = \"0deg\"\n                max_matches = 0\n                for deg in [\"0deg\", \"90deg\", \"180deg\", \"270deg\"]:\n\n                    keypoints2 = torch.from_numpy(f_keypoints_deg[key2][deg][...]).to(device)\n                    descriptors2 = torch.from_numpy(f_descriptors_deg[key2][deg][...]).to(device)\n                    offset = f_offsets_deg[key2][deg][...]\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                    n_matches = len(indices)\n                    if verbose:\n                        print(f\"{key1}-{key2}: {n_matches} matches, rotation: {deg}\")\n                    if n_matches > max_matches:\n                        best_deg = deg\n                        max_matches = n_matches\n                        best_indices = indices.detach().cpu().numpy().reshape(-1, 2)\n                        best_indices[:, 1] += offset\n                    if n_matches >= early_stopping_thr:\n                        break\n\n                # We have matches to consider\n                if max_matches:\n                    # Store the matches in the group of one image\n                    if max_matches >= min_matches:\n                        if verbose:\n                            print(f\"{key1}-{key2}: {max_matches} matches, best rotation: {best_deg}\")\n                        group  = f_matches.require_group(key1)\n                        group.create_dataset(key2, data=best_indices)\n\n    return","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:06.179775Z","iopub.execute_input":"2024-06-03T09:57:06.180085Z","iopub.status.idle":"2024-06-03T09:57:06.194785Z","shell.execute_reply.started":"2024-06-03T09:57:06.180061Z","shell.execute_reply":"2024-06-03T09:57:06.193946Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## special function for transparent object","metadata":{}},{"cell_type":"code","source":"def keypoint_distances_transparent(\n    paths,\n    index_pairs,\n    feature_dir,\n    num_grids,\n    device_id,\n    min_matches,\n    verbose,\n):\n    device = torch.device(f\"cuda:{device_id}\")\n\n    matcher_params = {\n        \"width_confidence\": -1,\n        \"depth_confidence\": -1,\n        \"mp\": True,\n    }\n    matcher = KF.LightGlueMatcher(\"aliked\", matcher_params).eval().to(device)\n\n    with h5py.File(feature_dir / \"keypoints_grid.h5\", mode=\"r\") as f_keypoints_grid, \\\n         h5py.File(feature_dir / \"descriptors_grid.h5\", mode=\"r\") as f_descriptors_grid, \\\n         h5py.File(feature_dir / \"offsets_grid.h5\", mode=\"r\") as f_offsets_grid, \\\n         h5py.File(feature_dir / f\"matches{device_id}.h5\", mode=\"w\") as f_matches:\n\n            for idx1, idx2 in tqdm(index_pairs, desc=f\"Computing keypoint distances[GPU{device_id}]\", dynamic_ncols=True):\n                key1, key2 = paths[idx1].name, paths[idx2].name\n\n                index_list = []\n                for grid_id in range(num_grids):\n\n                    if f_keypoints_grid[key1].get(f\"{grid_id}\") is None or f_keypoints_grid[key2].get(f\"{grid_id}\") is None:\n                        continue\n\n                    keypoints1 = torch.from_numpy(f_keypoints_grid[key1][f\"{grid_id}\"][...]).to(device)\n                    descriptors1 = torch.from_numpy(f_descriptors_grid[key1][f\"{grid_id}\"][...]).to(device)\n                    offset1 = f_offsets_grid[key1][f\"{grid_id}\"][...]\n\n                    keypoints2 = torch.from_numpy(f_keypoints_grid[key2][f\"{grid_id}\"][...]).to(device)\n                    descriptors2 = torch.from_numpy(f_descriptors_grid[key2][f\"{grid_id}\"][...]).to(device)\n                    offset2 = f_offsets_grid[key2][f\"{grid_id}\"][...]\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                    n_matches = len(indices)\n\n                    # Store the matches in the group of one image\n                    if n_matches >= min_matches:\n                        if verbose:\n                            print(f\"grid_id1:{grid_id}, {key1}-{key2}: {n_matches} matches\")\n\n                        indices = indices.detach().cpu().numpy().reshape(-1, 2)\n                        indices[:, 0] += offset1\n                        indices[:, 1] += offset2\n                        index_list.append(indices)\n\n                if len(index_list) == 0:\n                    continue\n\n                indices = np.concatenate(index_list)\n                group  = f_matches.require_group(key1)\n                group.create_dataset(key2, data=indices)\n\n    return","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:06.195936Z","iopub.execute_input":"2024-06-03T09:57:06.196204Z","iopub.status.idle":"2024-06-03T09:57:06.212128Z","shell.execute_reply.started":"2024-06-03T09:57:06.196171Z","shell.execute_reply":"2024-06-03T09:57:06.211358Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def merge_matches(\n    feature_dir,\n    input_matches_name_list,\n    output_matches_name,\n):\n    with h5py.File(feature_dir / output_matches_name, mode=\"w\") as f_matches:\n\n        for input_matches_name in input_matches_name_list:\n            with h5py.File(feature_dir / input_matches_name, mode=\"r\") as f_input_matches:\n                for key1 in f_input_matches.keys():\n                    group  = f_matches.require_group(key1)\n                    for key2 in f_input_matches[key1].keys():\n                        group.create_dataset(key2, data=f_input_matches[key1][key2][...])\n\n    return","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:06.213039Z","iopub.execute_input":"2024-06-03T09:57:06.213369Z","iopub.status.idle":"2024-06-03T09:57:06.226440Z","shell.execute_reply.started":"2024-06-03T09:57:06.213336Z","shell.execute_reply":"2024-06-03T09:57:06.225643Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Check orientation\n\nhttps://github.com/ternaus/check_orientation","metadata":{}},{"cell_type":"code","source":"def convert_rot_k(index):\n    if index == 0:\n        return 0\n    elif index == 1:\n        return 3\n    elif index == 2:\n        return 2\n    else:\n        return 1\n\nclass CheckRotationDataset(Dataset):\n    def __init__(self, files, transform=None):\n        self.transform = transform\n        self.files = files\n\n    def __len__(self):\n        return len(self.files)\n\n    def __getitem__(self, idx):\n        imgPath = self.files[idx]\n        image = T_read_image(str(imgPath), mode=ImageReadMode.RGB)\n        if self.transform:\n            image = self.transform(image)\n        return image\n\ndef get_CheckRotation_dataloader(images, batch_size=1):\n    transform = T.Compose([\n        T.Resize((224, 224)),\n        T.ConvertImageDtype(torch.float),\n        T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))\n    ])\n\n    dataset = CheckRotationDataset(images, transform=transform)\n    dataloader = DataLoader(\n        dataset=dataset,\n        shuffle=False,\n        batch_size=batch_size,\n        pin_memory=True,\n        num_workers=2,\n        drop_last=False\n    )\n    return dataloader\n\ndef exec_rotation_detection(img_files):\n\n    model = create_model(\"swsl_resnext50_32x4d\")\n    model.eval().cuda()\n\n    dataloader = get_CheckRotation_dataloader(img_files)\n\n    rots = []\n    for idx, image in enumerate(dataloader):\n        image = image.to(torch.float32).cuda()\n        with torch.no_grad():\n            prediction = model(image).detach().cpu().numpy()\n            detected_rot = prediction[0].argmax()\n            rot_k = convert_rot_k(detected_rot)\n            rots.append(rot_k)\n    return rots","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:06.227509Z","iopub.execute_input":"2024-06-03T09:57:06.228551Z","iopub.status.idle":"2024-06-03T09:57:06.241075Z","shell.execute_reply.started":"2024-06-03T09:57:06.228526Z","shell.execute_reply":"2024-06-03T09:57:06.240354Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def output_rot_images(\n    paths,\n    output_dir,\n    rots,\n):\n    corrected_image_paths = []\n    for rot, path in tqdm(zip(rots, paths), total=len(paths), desc=f\"Rotating images\", dynamic_ncols=True):\n\n        img = cv2.imread(str(path))\n\n        if rot == 1:\n            img = cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)\n        elif rot == 2:\n            img = cv2.rotate(img, cv2.ROTATE_180)\n        elif rot == 3:\n            img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)\n\n        cv2.imwrite(str(output_dir / path.name), img)\n        corrected_image_paths.append(output_dir / path.name)\n\n    return corrected_image_paths\n\ndef exec_rotation_correction(paths, output_dir):\n\n    rots = exec_rotation_detection(paths)\n\n    corrected_image_paths = []\n    with ThreadPoolExecutor() as executor:\n        results = executor.map(\n            output_rot_images,\n            np.array_split(paths, 2),\n            itertools.repeat(output_dir),\n            np.array_split(rots, 2),\n        )\n        for data in results:\n            corrected_image_paths.append(data)\n\n    corrected_image_paths = list(itertools.chain.from_iterable(corrected_image_paths))\n    gc.collect()\n\n    return corrected_image_paths","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:06.242094Z","iopub.execute_input":"2024-06-03T09:57:06.242366Z","iopub.status.idle":"2024-06-03T09:57:06.255397Z","shell.execute_reply.started":"2024-06-03T09:57:06.242337Z","shell.execute_reply":"2024-06-03T09:57:06.254578Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Doppelgangers: Learning to Disambiguate Images of Similar Structures\n\nhttps://github.com/RuojinCai/Doppelgangers","metadata":{}},{"cell_type":"code","source":"sys.path.append(\"/kaggle/input/doppelgangers-repo/doppelgangers\")\nfrom doppelgangers.third_party.loftr import LoFTR, default_cfg\nfrom doppelgangers.utils.loftr_matches import read_image","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:06.256576Z","iopub.execute_input":"2024-06-03T09:57:06.256918Z","iopub.status.idle":"2024-06-03T09:57:06.637180Z","shell.execute_reply.started":"2024-06-03T09:57:06.256885Z","shell.execute_reply":"2024-06-03T09:57:06.636248Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def save_loftr_matches(data_path, pairs_info, pairs_info_index, output_path, device_id, model_weight_path=\"weights/outdoor_ds.ckpt\"):\n    # The default config uses dual-softmax.\n    # The outdoor and indoor models share the same config.\n    # You can change the default values like thr and coarse_match_type.\n\n    device = torch.device(f\"cuda:{device_id}\")\n\n    matcher = LoFTR(config=default_cfg)\n    matcher.load_state_dict(torch.load(model_weight_path)['state_dict'])\n    matcher = matcher.eval().to(device)\n\n    img_size = 1024\n    df = 8\n    padding = True\n\n    for _pairs_info, idx in tqdm(zip(pairs_info, pairs_info_index), total=len(pairs_info_index), desc=f\"Running LOFTR [GPU{device_id}]\"):\n\n        output_file_path = output_path / \"loftr_match\" / f\"{idx}.npy\"\n        if output_file_path.exists():\n            continue\n\n        name0, name1, _, _, _ = _pairs_info\n\n        img0_pth = data_path / name0\n        img1_pth = data_path / name1\n        img0_raw, mask0 = read_image(str(img0_pth), img_size, df, padding)\n        img1_raw, mask1 = read_image(str(img1_pth), img_size, df, padding)\n        img0 = torch.from_numpy(img0_raw).to(device)\n        img1 = torch.from_numpy(img1_raw).to(device)\n        mask0 = torch.from_numpy(mask0).to(device)\n        mask1 = torch.from_numpy(mask1).to(device)\n        batch = {'image0': img0, 'image1': img1, 'mask0': mask0, 'mask1':mask1}\n\n        # Inference with LoFTR and get prediction\n        with torch.no_grad():\n            matcher(batch)\n            mkpts0 = batch['mkpts0_f'].cpu().numpy()\n            mkpts1 = batch['mkpts1_f'].cpu().numpy()\n            mconf = batch['mconf'].cpu().numpy()\n\n            np.save(output_file_path, {\"kpt0\": mkpts0, \"kpt1\": mkpts1, \"conf\": mconf})\n\ndef doppelgangers_classifier(cfg, pretrained_model_path, pair_path):\n    # basic setup\n    cudnn.benchmark = True\n\n    # initial dataset\n    data_lib = importlib.import_module(cfg.data.type)\n    loaders = data_lib.get_data_loaders(cfg.data)\n    test_loader = loaders[\"test_loader\"]\n\n    # initial model\n    decoder_lib = importlib.import_module(cfg.models.decoder.type)\n    decoder = decoder_lib.decoder(cfg.models.decoder)\n    decoder = decoder.cuda()\n\n    # load pretrained model\n    ckpt = torch.load(pretrained_model_path)\n    new_ckpt = deepcopy(ckpt[\"dec\"])\n    for key, _ in ckpt[\"dec\"].items():\n        if \"module.\" in key:\n            new_ckpt[key[len(\"module.\"):]] = new_ckpt.pop(key)\n    decoder.load_state_dict(new_ckpt, strict=True)\n\n    # evaluate on test set\n    decoder.eval()\n    prob_list = list()\n    with torch.no_grad():\n        for _, data in tqdm(enumerate(test_loader), total=len(test_loader), desc=\"Running doppelgangers\"):\n\n            # TTA\n            img = data[\"image\"].cuda()\n\n            img_list = []\n            img_list.append(img)\n            # lrflip\n            img = torch.flip(img, [3])\n            img_list.append(img)\n            # stack\n            img = torch.cat(img_list)\n\n            logits = decoder(img)\n            prob = torch.nn.functional.softmax(logits, dim=1)\n            prob = torch.mean(prob, dim=0)\n            prob_list.append(prob[1].detach().cpu().numpy())\n\n    y_scores = np.array(prob_list)\n\n    pairs_info = np.load(pair_path, allow_pickle=True)\n    pair_probability_file_path = cfg.data.output_path / \"pair_probability_list.h5\"\n    with h5py.File(pair_probability_file_path, mode=\"w\") as f_matches:\n\n        for idx in range(pairs_info.shape[0]):\n\n            key1, key2, _, _, _ = pairs_info[idx]\n            score = y_scores[idx]\n\n            group  = f_matches.require_group(key1)\n            group.create_dataset(key2, data=score)\n\n    return pair_probability_file_path\n\ndef create_image_pair_list(matches_path, output_path):\n\n    dummy = 0\n    pairs_list = []\n    with h5py.File(matches_path, mode=\"r\") as f_matches:\n\n        for key1 in f_matches.keys():\n            group = f_matches[key1]\n            for key2 in group.keys():\n                pairs_list.append([key1, key2, dummy, dummy, dummy])\n\n    pairs_list = np.concatenate(pairs_list, axis=0).reshape(-1, 5)\n\n    pairs_list_path = output_path / \"pairs_list.npy\"\n    np.save(pairs_list_path, pairs_list)\n\n    return pairs_list_path\n\ndef exec_doppelgangers_classifier(\n    images_dir,\n    feature_dir,\n    matches_path,\n    num_gpus,\n    loftr_weight_path,\n    doppelgangers_weight_path,\n):\n    loftr_matches_path = feature_dir / \"loftr_match\"\n    loftr_matches_path.mkdir(parents=True, exist_ok=True)\n\n    pair_path = create_image_pair_list(matches_path, feature_dir)\n    pairs_info = np.load(pair_path, allow_pickle=True)\n    pairs_info_index = np.arange(pairs_info.shape[0])\n\n    with ThreadPoolExecutor() as executor:\n        executor.map(\n            save_loftr_matches,\n            itertools.repeat(images_dir),\n            np.array_split(pairs_info, num_gpus),\n            np.array_split(pairs_info_index, num_gpus),\n            itertools.repeat(feature_dir),\n            range(num_gpus),\n            itertools.repeat(loftr_weight_path),\n        )\n    gc.collect()\n    torch.cuda.empty_cache()\n\n    config = {\n        \"data\": {\n            \"image_dir\": images_dir,\n            \"loftr_match_dir\": loftr_matches_path,\n            \"output_path\": feature_dir,\n            \"type\": \"doppelgangers.datasets.sfm_disambiguation_dataset\",\n            \"num_workers\": 1,\n            \"test\": {\n                \"batch_size\": 1,\n                \"img_size\": 1024,\n                \"pair_path\": pair_path,\n            },\n        },\n        \"models\": {\n            \"decoder\": {\n                \"type\": \"doppelgangers.models.cnn_classifier\",\n                \"input_dim\": 10,\n            },\n        },\n    }\n\n    def dict2namespace(config):\n        namespace = argparse.Namespace()\n        for key, value in config.items():\n            if isinstance(value, dict):\n                new_value = dict2namespace(value)\n            else:\n                new_value = value\n            setattr(namespace, key, new_value)\n        return namespace\n    \n    config = dict2namespace(config)\n\n    # Running Doppelgangers classifier model on image pairs\n    print(\"Running Doppelgangers classifier model on image pairs\")\n    pair_probability_file_path = doppelgangers_classifier(config, doppelgangers_weight_path, pair_path)\n\n    shutil.rmtree(loftr_matches_path)\n\n    return pair_probability_file_path","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:06.638462Z","iopub.execute_input":"2024-06-03T09:57:06.638729Z","iopub.status.idle":"2024-06-03T09:57:06.668609Z","shell.execute_reply.started":"2024-06-03T09:57:06.638706Z","shell.execute_reply":"2024-06-03T09:57:06.667749Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Identify duplicate structures by 3D point cloud","metadata":{}},{"cell_type":"code","source":"sys.path.append(\"/kaggle/input/colmap-db-import\")\nfrom prepare_colmap import read_images_binary, read_cameras_binary, read_points3D_binary, project_3d_to_2d, get_camera_param","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:06.669806Z","iopub.execute_input":"2024-06-03T09:57:06.670110Z","iopub.status.idle":"2024-06-03T09:57:06.689152Z","shell.execute_reply.started":"2024-06-03T09:57:06.670085Z","shell.execute_reply":"2024-06-03T09:57:06.688517Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_points_inside_cameraview(\n        T_pointcloud_camera, K, width, height, # target camera\n        point_cloud, # 3D point cloud\n):\n    point_cloud_uv = project_3d_to_2d(T_pointcloud_camera, K, point_cloud)\n    mask = (point_cloud_uv[:, 0] >= 0) & (point_cloud_uv[:, 0] < width) & (point_cloud_uv[:, 1] >= 0) & (point_cloud_uv[:, 1] < height)\n\n    return point_cloud_uv[mask]\n\ndef keep_multi_camera_points(images, point_cloud_id, point_cloud, pcd_used_num_images_ratio):\n\n    pcd_id_list = np.concatenate([np.unique(image['points_ids']) for image in images.values()])\n    unique, counts = np.unique(pcd_id_list, return_counts=True)\n    pcd_id_count = dict(zip(unique, counts))\n\n    pcd_used_num_images = int(len(images) * pcd_used_num_images_ratio)\n    filtered_pcd_ids = [k for k, v in pcd_id_count.items() if v > pcd_used_num_images]\n    target_pcd_idx = np.isin(point_cloud_id, filtered_pcd_ids)\n\n    return point_cloud_id[target_pcd_idx], point_cloud[:, target_pcd_idx]","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:06.690129Z","iopub.execute_input":"2024-06-03T09:57:06.690403Z","iopub.status.idle":"2024-06-03T09:57:06.697857Z","shell.execute_reply.started":"2024-06-03T09:57:06.690380Z","shell.execute_reply":"2024-06-03T09:57:06.696914Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def remove_ambiguous_area(paths, recon_data_dir, pcd_used_num_images_ratio, erode_mask_ratio):\n\n    images = read_images_binary(recon_data_dir / \"images.bin\")\n    cameras = read_cameras_binary(recon_data_dir / \"cameras.bin\")\n    points = read_points3D_binary(recon_data_dir / \"points3D.bin\")\n\n    point_cloud_id = points[\"id\"].values\n    point_cloud = points[['x', 'y', 'z']].values\n    point_cloud = point_cloud.T\n\n    # step1: Extract point cloud observed by more than half of the images\n    point_cloud_id, point_cloud = keep_multi_camera_points(images, point_cloud_id, point_cloud, pcd_used_num_images_ratio)\n\n    # step2: Cluster and remove outliers\n    if len(point_cloud_id) < 10:\n        num_labels = 0\n    else:\n        nearest_neighbors = NearestNeighbors(n_neighbors=10)\n        neighbors = nearest_neighbors.fit(point_cloud.T)\n        distances, _ = neighbors.kneighbors(point_cloud.T)\n        distances = np.sort(distances[:,-1], axis=0)\n\n        i = np.arange(len(distances))\n        knee = KneeLocator(i, distances, S=1, curve='convex', direction='increasing', interp_method='polynomial')\n        eps = distances[knee.knee]\n        print(f\"Knee: {knee.knee}, eps: {eps}\")\n\n        labels = DBSCAN(eps=eps, min_samples=5).fit(point_cloud.T).labels_\n        num_labels = len(np.unique(labels))-1 if -1 in labels else len(np.unique(labels))\n\n    print(f\"Number of clusters: {num_labels}\")\n\n    kernel = np.ones((3,3), np.uint8)\n    mask_dict = {}\n    point_cloud_uv_dict = {}\n    for path in tqdm(paths, total=len(paths), desc=\"remove ambiguous area\"):\n\n        key = path.name\n        if key not in images:\n            img = cv2.imread(str(path))\n            height, width, _ = img.shape\n            mask = np.zeros((height, width), dtype=np.uint8)\n            mask_dict[key] = mask\n            continue\n\n        image = images[key]\n\n        pcd_ids_cam = image['points_ids']\n        T_pointcloud_camera, K, width, height = get_camera_param(cameras, image)\n\n        all_mask = np.zeros((height, width), dtype=np.uint8)\n        all_point_cloud_uv = []\n        for label_id in range(num_labels):\n\n            lbl_id = labels==label_id\n            _point_cloud_id = point_cloud_id[lbl_id]\n            _point_cloud = point_cloud[:,lbl_id]\n\n            target_idx = np.isin(_point_cloud_id, pcd_ids_cam)\n            point_cloud_cam = _point_cloud[:, target_idx]\n\n            # step3: Get the point cloud in the camera view\n            point_cloud_uv = get_points_inside_cameraview(T_pointcloud_camera, K, width, height, point_cloud_cam)\n\n            # step4: Create a mask image\n            mask = np.zeros((height, width), dtype=np.uint8)\n            if len(point_cloud_uv) > 0:\n                point_cloud_uv_int = point_cloud_uv.astype(int)\n                mask[point_cloud_uv_int[:,1],point_cloud_uv_int[:,0]] = 255\n\n                dilate_count = 0\n                for i in range(1000):\n                    n_labels, _ = cv2.connectedComponents(mask)\n                    if n_labels <= 2:\n                        break\n                    mask = cv2.dilate(mask, kernel=kernel, iterations=1)\n                    dilate_count += 1\n\n                erode_count = int(dilate_count * erode_mask_ratio)\n                if erode_count > 0:\n                    mask = cv2.erode(mask, kernel=kernel, iterations=erode_count)\n\n            all_mask[mask>0] = 255\n            all_point_cloud_uv.append(point_cloud_uv)\n\n        mask_dict[key] = all_mask\n        point_cloud_uv_dict[key] = np.concatenate(all_point_cloud_uv, axis=0) if len(all_point_cloud_uv) > 0 else np.zeros((0, 2))\n\n    return mask_dict, point_cloud_uv_dict","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:06.704484Z","iopub.execute_input":"2024-06-03T09:57:06.704903Z","iopub.status.idle":"2024-06-03T09:57:06.722265Z","shell.execute_reply.started":"2024-06-03T09:57:06.704879Z","shell.execute_reply":"2024-06-03T09:57:06.721275Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Prune the matching graphs caused by duplicate structures","metadata":{}},{"cell_type":"code","source":"def filter_FundamentalMatrix(mkpts0, mkpts1, filter_iterations=10, filter_threshold=8):\n\n    store_inliers = { idx:0 for idx in range(mkpts0.shape[0]) }\n    idxs = np.array(range(mkpts0.shape[0]))\n    for _ in range(filter_iterations):\n        try:\n            Fm, inliers = cv2.findFundamentalMat(\n                mkpts0, mkpts1, cv2.USAC_MAGSAC, 0.15, 0.9999, 20000)\n            if Fm is not None:\n                inliers = (inliers > 0).reshape(-1)\n                inlier_idxs = idxs[inliers]\n                for idx in inlier_idxs:\n                    store_inliers[idx] += 1\n        except:\n            print(f\"Failed to cv2.findFundamentalMat. mkpts.shape={mkpts0.shape}\")\n    inliers = np.array([ count for (idx, count) in store_inliers.items() ]) >= filter_threshold\n\n    return inliers\n\ndef disambiguous_keypoints(mask_area, keypoints):\n    disambiguous_idx = []\n    for kpt in keypoints.astype(int):\n        if kpt[0] < 0 or kpt[1] < 0 or kpt[0] >= mask_area.shape[1] or kpt[1] >= mask_area.shape[0]:\n            disambiguous_idx.append(True)\n            continue\n        flag = (mask_area[kpt[1], kpt[0]] == 0)\n        disambiguous_idx.append(flag)\n    return np.array(disambiguous_idx)\n\ndef verify_matches(\n    feature_dir,\n    pair_probability_file_path,\n    mask_area,\n    doppelgangers_min_thr,\n    doppelgangers_max_thr,\n    filter_iterations,\n    filter_threshold,\n):\n\n    index_pairs = []\n    with h5py.File(feature_dir / \"keypoints.h5\", mode=\"r\") as f_keypoints, \\\n         h5py.File(feature_dir / \"offsets.h5\", mode=\"r\") as f_offsets_deg, \\\n         h5py.File(feature_dir / \"matches.h5\", mode=\"r\") as f_matches, \\\n         h5py.File(pair_probability_file_path, mode=\"r\") as f_scores:\n\n        for key1 in tqdm(f_matches.keys(), desc=\"verify matches\"):\n\n            group = f_matches[key1]\n            keypoints1 = f_keypoints[key1][...]\n            mask_area_1 = pad_to_square(mask_area[key1])\n\n            group_score = f_scores[key1]\n\n            for key2 in group.keys():\n\n                keypoints2 = f_keypoints[key2][...]\n                matches = group[key2][...]\n                mask_area_2 = pad_to_square(mask_area[key2])\n\n                doppelgangers_score = group_score[key2][...]\n\n                if doppelgangers_score < doppelgangers_min_thr:\n                    continue\n\n                if doppelgangers_score > doppelgangers_max_thr:\n                    index_pairs.append((key1, key2, matches))\n                    continue\n\n                # rotate mask_area_2\n                for deg in [\"90deg\", \"180deg\", \"270deg\"]:\n                    offset = f_offsets_deg[key2][deg][...]\n                    if offset > matches[0, 1]:\n                        break\n                    mask_area_2 = cv2.rotate(mask_area_2, cv2.ROTATE_90_CLOCKWISE)\n\n                # Verify using Two-View Geometry\n                mkpts1 = keypoints1[matches[:, 0]]\n                mkpts2 = keypoints2[matches[:, 1]]\n                inliers_fmat = filter_FundamentalMatrix(mkpts1, mkpts2, filter_iterations, filter_threshold)\n\n                # Verify using Ambiguity Mask\n                inliers_disambiguous_kpts1 = disambiguous_keypoints(mask_area_1, mkpts1)\n                inliers_disambiguous_kpts2 = disambiguous_keypoints(mask_area_2, mkpts2)\n\n                # Integrate verification results\n                inliers = inliers_fmat & inliers_disambiguous_kpts1 & inliers_disambiguous_kpts2\n\n                index_pairs.append((key1, key2, matches[inliers]))\n\n    with h5py.File(feature_dir / \"matches.h5\", mode=\"w\") as f_matches:\n\n        for key1, key2, matches in index_pairs:\n            group  = f_matches.require_group(key1)\n            group.create_dataset(key2, data=matches)\n\n    return","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:06.723661Z","iopub.execute_input":"2024-06-03T09:57:06.723959Z","iopub.status.idle":"2024-06-03T09:57:06.742411Z","shell.execute_reply.started":"2024-06-03T09:57:06.723936Z","shell.execute_reply":"2024-06-03T09:57:06.741501Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Get focal length from incremental mapping","metadata":{}},{"cell_type":"code","source":"def get_focal_length_prior(paths, recon_data_dir):\n\n    images = read_images_binary(recon_data_dir / \"images.bin\")\n    cameras = read_cameras_binary(recon_data_dir / \"cameras.bin\")\n\n    focal_length_dict = {}\n    for path in paths:\n\n        key = path.name\n\n        if key not in images:\n            img = cv2.imread(str(path))\n            height, width, _ = img.shape\n            FOCAL_PRIOR = 1.2\n            focal_length_dict[key] = FOCAL_PRIOR * max(height, width)\n            continue\n\n        _, K, _, _ = get_camera_param(cameras, images[key])\n\n        focal_length_dict[key] = K[0][0]\n\n    return focal_length_dict","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:06.743413Z","iopub.execute_input":"2024-06-03T09:57:06.743665Z","iopub.status.idle":"2024-06-03T09:57:06.754839Z","shell.execute_reply.started":"2024-06-03T09:57:06.743643Z","shell.execute_reply":"2024-06-03T09:57:06.754095Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# COLMAP Utils","metadata":{}},{"cell_type":"code","source":"sys.path.append(\"/kaggle/input/colmap-db-import\")\nfrom database import COLMAPDatabase\nfrom h5_to_db import add_keypoints, add_matches","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:06.755997Z","iopub.execute_input":"2024-06-03T09:57:06.756379Z","iopub.status.idle":"2024-06-03T09:57:06.776160Z","shell.execute_reply.started":"2024-06-03T09:57:06.756324Z","shell.execute_reply":"2024-06-03T09:57:06.775539Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def import_into_colmap(\n    path,\n    feature_dir,\n    database_path,\n    camera_model,\n    focal_length_dict=None,\n    single_camera=False,\n):\n    db = COLMAPDatabase.connect(database_path)\n    db.create_tables()\n    fname_to_id = add_keypoints(\n        db,\n        feature_dir,\n        path,\n        camera_model,\n        single_camera,\n        focal_length_dict,\n    )\n    add_matches(\n        db,\n        feature_dir,\n        fname_to_id,\n    )\n    db.commit()\n\n    return","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:06.777091Z","iopub.execute_input":"2024-06-03T09:57:06.777366Z","iopub.status.idle":"2024-06-03T09:57:06.782623Z","shell.execute_reply.started":"2024-06-03T09:57:06.777333Z","shell.execute_reply":"2024-06-03T09:57:06.781667Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Submission Utils","metadata":{}},{"cell_type":"code","source":"def parse_sample_submission(data_type, scene_type):\n\n    df_scene = pd.read_csv(f\"/kaggle/input/image-matching-challenge-2024/{data_type}/categories.csv\")\n    if data_type == \"train\":\n        df = pd.read_csv(\"/kaggle/input/imc2024-validation/validation.csv\")\n        if scene_type != \"submission\":\n            df = df[df[\"scene\"] == scene_type]\n            df_scene = df_scene[df_scene[\"scene\"] == scene_type]\n    else:\n        df = pd.read_csv(\"/kaggle/input/image-matching-challenge-2024/sample_submission.csv\")\n\n    data_dict = {}\n    category_dict = {}\n    for dataset in df[\"dataset\"].unique():\n        data_dict[dataset] = {}\n        category_dict[dataset] = {}\n        for scene in df[df[\"dataset\"] == dataset][\"scene\"].unique():\n            data_dict[dataset][scene] = []\n\n            image_path_list = df[(df[\"dataset\"] == dataset) & (df[\"scene\"] == scene)][\"image_path\"].tolist()\n            for image_path in image_path_list:\n                data_dict[dataset][scene].append(Path(f\"/kaggle/input/image-matching-challenge-2024/{image_path}\"))\n\n            category_dict[dataset][scene] = df_scene[df_scene[\"scene\"] == scene][\"categories\"].values[0]\n\n    return data_dict, category_dict\n\ndef arr_to_str(a):\n    return \";\".join([str(x) for x in a.reshape(-1)])\n\ndef create_submission(\n    output_dir,\n    results,\n    data_dict,\n    base_path,\n    scene_type,\n):\n    \"\"\"Prepares a submission file.\"\"\"\n\n    with open(output_dir / f\"{scene_type}.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\")\n    return","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:06.784159Z","iopub.execute_input":"2024-06-03T09:57:06.784499Z","iopub.status.idle":"2024-06-03T09:57:06.799570Z","shell.execute_reply.started":"2024-06-03T09:57:06.784469Z","shell.execute_reply":"2024-06-03T09:57:06.798683Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Feature matching","metadata":{}},{"cell_type":"code","source":"def feature_matching_default(\n    config,\n    feature_dir,\n    image_paths,\n    index_pairs,\n    trial_index,\n):\n    # Detect keypoints of all images\n    with ThreadPoolExecutor() as executor:\n        executor.map(\n            detect_keypoints,\n            np.array_split(image_paths, config.num_gpus),\n            itertools.repeat(feature_dir),\n            range(config.num_gpus),\n            itertools.repeat(config.keypoint_detection_args[\"num_features\"]),\n            itertools.repeat(config.keypoint_detection_args[\"detection_threshold\"]),\n            itertools.repeat(config.keypoint_detection_args[\"resize_to\"]),\n        )\n    torch.cuda.empty_cache()\n    gc.collect()\n\n    merge_single_h5(\n        feature_dir,\n        [ f\"keypoints{i}.h5\" for i in range(config.num_gpus) ],\n        \"keypoints.h5\",\n    )\n    for i in range(config.num_gpus):\n        (feature_dir / f\"keypoints{i}.h5\").unlink()\n\n    for file_prefix in [\"descriptors_deg\", \"keypoints_deg\", \"offsets\"]:\n        merge_double_h5(\n            feature_dir,\n            [ f\"{file_prefix}{i}.h5\" for i in range(config.num_gpus) ],\n            f\"{file_prefix}.h5\",\n        )\n        for i in range(config.num_gpus):\n            (feature_dir / f\"{file_prefix}{i}.h5\").unlink()\n\n    # Match keypoints of pairs of similar images\n    with ThreadPoolExecutor() as executor:\n        executor.map(\n            keypoint_distances,\n            itertools.repeat(image_paths),\n            np.array_split(index_pairs, config.num_gpus),\n            itertools.repeat(feature_dir),\n            range(config.num_gpus),\n            itertools.repeat(config.keypoint_distances_args[\"early_stopping_thr\"]),\n            itertools.repeat(config.keypoint_distances_args[\"min_matches\"][trial_index]),\n            itertools.repeat(config.keypoint_distances_args[\"verbose\"]),\n        )\n    (feature_dir / \"descriptors_deg.h5\").unlink()\n    (feature_dir / \"keypoints_deg.h5\").unlink()\n    torch.cuda.empty_cache()\n    gc.collect()\n\n    # Merge the matching results of each GPU\n    merge_matches(\n        feature_dir,\n        [ f\"matches{i}.h5\" for i in range(config.num_gpus) ],\n        \"matches.h5\",\n    )\n    for i in range(config.num_gpus):\n        (feature_dir / f\"matches{i}.h5\").unlink()\n    gc.collect()\n\n    return\n\n# Special processing for transparent objects\ndef feature_matching_transparent(\n    config,\n    feature_dir,\n    image_paths,\n    index_pairs,\n    trial_index,\n):\n    # Correct the orientation of the images\n    corrected_images_dir = feature_dir / \"corrected_images\"\n    corrected_images_dir.mkdir(parents=True, exist_ok=True)\n    corrected_image_paths = exec_rotation_correction(\n        image_paths,\n        corrected_images_dir,\n    )\n    torch.cuda.empty_cache()\n    gc.collect()\n\n    # Extract foreground masks of all images\n    with ThreadPoolExecutor() as executor:\n        executor.map(\n            dinov2_segmentation,\n            np.array_split(corrected_image_paths, config.num_gpus),\n            itertools.repeat(feature_dir),\n            range(config.num_gpus),\n        )\n    torch.cuda.empty_cache()\n    gc.collect()\n\n    merge_single_h5(\n        feature_dir,\n        [ f\"fg_mask{i}.h5\" for i in range(config.num_gpus) ],\n        \"fg_mask.h5\",\n    )\n    for i in range(config.num_gpus):\n        (feature_dir / f\"fg_mask{i}.h5\").unlink()\n    gc.collect()\n\n    # Detect keypoints of all images\n    num_grids = []\n    with ThreadPoolExecutor() as executor:\n        results = executor.map(\n            detect_keypoints_transparent,\n            np.array_split(corrected_image_paths, config.num_gpus),\n            itertools.repeat(feature_dir),\n            range(config.num_gpus),\n            itertools.repeat(config.keypoint_detection_transparent_args[\"num_features\"]),\n            itertools.repeat(config.keypoint_detection_transparent_args[\"detection_threshold\"]),\n            itertools.repeat(config.keypoint_detection_transparent_args[\"resize_to\"]),\n        )\n        for data in results:\n            num_grids.append(data)\n    num_grids = max(num_grids)\n    torch.cuda.empty_cache()\n    gc.collect()\n\n    (feature_dir / \"fg_mask.h5\").unlink()\n    shutil.rmtree(corrected_images_dir)\n\n    merge_single_h5(\n        feature_dir,\n        [ f\"keypoints{i}.h5\" for i in range(config.num_gpus) ],\n        \"keypoints.h5\",\n    )\n    for i in range(config.num_gpus):\n        (feature_dir / f\"keypoints{i}.h5\").unlink()\n\n    for file_prefix in [\"descriptors_grid\", \"keypoints_grid\", \"offsets_grid\"]:\n        merge_double_h5(\n            feature_dir,\n            [ f\"{file_prefix}{i}.h5\" for i in range(config.num_gpus) ],\n            f\"{file_prefix}.h5\",\n        )\n        for i in range(config.num_gpus):\n            (feature_dir / f\"{file_prefix}{i}.h5\").unlink()\n\n    gc.collect()\n\n    # Match keypoints of pairs of similar images\n    with ThreadPoolExecutor() as executor:\n        executor.map(\n            keypoint_distances_transparent,\n            itertools.repeat(corrected_image_paths),\n            np.array_split(index_pairs, config.num_gpus),\n            itertools.repeat(feature_dir),\n            itertools.repeat(num_grids),\n            range(config.num_gpus),\n            itertools.repeat(config.keypoint_distances_args[\"min_matches\"][trial_index]),\n            itertools.repeat(config.keypoint_distances_args[\"verbose\"]),\n        )\n    (feature_dir / \"descriptors_grid.h5\").unlink()\n    (feature_dir / \"keypoints_grid.h5\").unlink()\n    (feature_dir / \"offsets_grid.h5\").unlink()\n    torch.cuda.empty_cache()\n    gc.collect()\n\n    # Merge the matching results of each GPU\n    merge_matches(\n        feature_dir,\n        [ f\"matches{i}.h5\" for i in range(config.num_gpus) ],\n        \"matches.h5\",\n    )\n    for i in range(config.num_gpus):\n        (feature_dir / f\"matches{i}.h5\").unlink()\n    gc.collect()\n\n    return\n\ndef feature_matching(args, config):\n\n    trial_index = args[\"trial_index\"]\n    image_paths = args[\"image_paths\"]\n    categories = args[\"categories\"]\n    feature_dir = args[\"feature_dir\"]\n\n    if feature_dir.exists():\n        shutil.rmtree(feature_dir)\n    feature_dir.mkdir(parents=True, exist_ok=True)\n\n    database_path = feature_dir / \"database.db\"\n\n    # Get the pairs of images that are somewhat similar\n    index_pairs = get_pairs_exhaustive(image_paths)\n    gc.collect()\n\n    # KeyPoint Detection and Matching\n    if \"transparent\" in categories:\n        feature_matching_transparent(\n            config,\n            feature_dir,\n            image_paths,\n            index_pairs,\n            trial_index,\n        )\n    else:\n        feature_matching_default(\n            config,\n            feature_dir,\n            image_paths,\n            index_pairs,\n            trial_index,\n        )\n\n    # Import keypoint distances of matches into colmap\n    import_into_colmap(\n        image_paths[0].parent,\n        feature_dir,\n        database_path,\n        camera_model = \"simple-radial\" if (\"historical_preservation\" in categories) or (\"transparent\" in categories) else \"simple-pinhole\",\n        single_camera = True if \"transparent\" in categories else False,\n    )\n    if (\"symmetries-and-repeats\" not in categories):\n        for file_path in feature_dir.glob(\"*.h5\"):\n            file_path.unlink()\n    if (\"transparent\" in categories):\n        for file_path in feature_dir.glob(\"*.h5\"):\n            file_path.unlink()\n\n    results = {\n        \"image_paths\": image_paths,\n        \"categories\": categories,\n        \"feature_dir\": feature_dir,\n        \"database_path\": database_path,\n    }\n    gc.collect()\n\n    return results","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:06.801001Z","iopub.execute_input":"2024-06-03T09:57:06.801283Z","iopub.status.idle":"2024-06-03T09:57:06.832099Z","shell.execute_reply.started":"2024-06-03T09:57:06.801260Z","shell.execute_reply":"2024-06-03T09:57:06.831367Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 3D Reconstruction","metadata":{}},{"cell_type":"code","source":"def reconstruction(args, config):\n\n    feature_dir = args[\"feature_dir\"]\n    image_paths = args[\"image_paths\"]\n    database_path = args[\"database_path\"]\n\n    # 4.2. Feature Matching and Geometric Verification\n    pycolmap.match_exhaustive(database_path)\n    gc.collect()\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    output_path = feature_dir / \"colmap_rec_aliked\"\n    output_path.mkdir(parents=True, exist_ok=True)\n\n    maps = pycolmap.incremental_mapping(\n        database_path=database_path,\n        image_path=image_paths[0].parent,\n        output_path=output_path,\n        options=pycolmap.IncrementalPipelineOptions(**config.colmap_mapper_options),\n    )\n    gc.collect()\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    num_points3D = 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                    num_points3D = len(rec.points3D)\n                    best_idx = idx1\n                elif len(rec.images) == images_registered:\n                    if len(rec.points3D) > num_points3D:\n                        images_registered = len(rec.images)\n                        num_points3D = len(rec.points3D)\n                        best_idx = idx1\n\n            except Exception:\n                continue\n\n    if best_idx is not None:\n        results = {}\n        results[\"images_registered\"] = images_registered\n        results[\"num_points3D\"] = num_points3D\n        results[\"best_idx\"] = best_idx\n        results[\"maps\"] = deepcopy(maps)\n    else:\n        results = None\n    gc.collect()\n\n    return results\n\ndef refine_symmetries_and_repeats(\n    config,\n    feature_dir,\n    image_paths,\n    recon_data_dir,\n):\n    # 6.1 Correct the orientation of the images\n    corrected_images_dir = feature_dir / \"corrected_images\"\n    corrected_images_dir.mkdir(parents=True, exist_ok=True)\n    exec_rotation_correction(image_paths, corrected_images_dir)\n    torch.cuda.empty_cache()\n    gc.collect()\n\n    # 6.2 Classify doppelgangers\n    pair_probability_file_path = exec_doppelgangers_classifier(\n        corrected_images_dir,\n        feature_dir,\n        feature_dir / \"matches.h5\",\n        config.num_gpus,\n        **config.doppelgangers_classifier_args,\n    )\n    shutil.rmtree(corrected_images_dir)\n    torch.cuda.empty_cache()\n    gc.collect()\n\n    # 6.3 Remove duplicate structures\n    mask_area, _ = remove_ambiguous_area(\n        image_paths,\n        recon_data_dir,\n        **config.remove_ambiguous_area_args,\n    )\n    gc.collect()\n\n    # 6.4 Prune the matching graphs caused by duplicate structures\n    verify_matches(\n        feature_dir,\n        pair_probability_file_path,\n        mask_area,\n        **config.verify_matches_args,\n    )\n    for file_path in feature_dir.glob(\"*.npy\"):\n        file_path.unlink()\n    gc.collect()\n\n    # 7. Reconstruct the scene again\n    database_path = feature_dir / \"database_refine.db\"\n    if database_path.exists():\n        database_path.unlink()\n\n    import_into_colmap(\n        image_paths[0].parent,\n        feature_dir,\n        database_path,\n        camera_model = \"simple-pinhole\", # TODO:\"simple-radial\"\n        focal_length_dict = get_focal_length_prior(image_paths, recon_data_dir),\n        single_camera = False,\n    )\n    for file_path in feature_dir.glob(\"*.h5\"):\n        file_path.unlink()\n\n    pycolmap.match_exhaustive(database_path)\n    gc.collect()\n\n    output_path = feature_dir / \"colmap_rec_aliked_refine\"\n    output_path.mkdir(parents=True, exist_ok=True)\n\n    maps = pycolmap.incremental_mapping(\n        database_path=database_path,\n        image_path=image_paths[0].parent,\n        output_path=output_path,\n        options=pycolmap.IncrementalPipelineOptions(**config.colmap_mapper_options),\n    )\n    gc.collect()\n\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    if not DEBUG:\n        shutil.rmtree(feature_dir)\n\n    if best_idx is not None:\n        results = {}\n        results[\"best_idx\"] = best_idx\n        results[\"maps\"] = deepcopy(maps)\n    else:\n        results = None\n    gc.collect()\n\n    return results","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:06.833078Z","iopub.execute_input":"2024-06-03T09:57:06.833305Z","iopub.status.idle":"2024-06-03T09:57:06.854562Z","shell.execute_reply.started":"2024-06-03T09:57:06.833285Z","shell.execute_reply":"2024-06-03T09:57:06.853657Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Pipeline process","metadata":{}},{"cell_type":"code","source":"def prepare_input(data_dict, config):\n\n    datasets = list(data_dict.keys())\n\n    image_num_list = []\n    dataset_list = []\n    scene_list = []\n    trial_list = []\n    for dataset in datasets:\n        for scene in data_dict[dataset]:\n            image_num = len(data_dict[dataset][scene])\n            for trial_idx in range(config.num_trials):\n                dataset_list.append(dataset)\n                scene_list.append(scene)\n                trial_list.append(trial_idx)\n                image_num_list.append(image_num)\n\n    # sort by the number of images\n    dataset_list = np.array(dataset_list)\n    scene_list = np.array(scene_list)\n    trial_list = np.array(trial_list)\n    image_num_list = np.array(image_num_list)\n    sort_idx = np.argsort(image_num_list)[::-1]\n    dataset_list = dataset_list[sort_idx]\n    scene_list = scene_list[sort_idx]\n    trial_list = trial_list[sort_idx]\n\n    return dataset_list, scene_list, trial_list","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:06.855626Z","iopub.execute_input":"2024-06-03T09:57:06.855891Z","iopub.status.idle":"2024-06-03T09:57:06.868571Z","shell.execute_reply.started":"2024-06-03T09:57:06.855861Z","shell.execute_reply":"2024-06-03T09:57:06.867704Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# core function for multithreading\nimport threading\nimport queue\n\ndef wrap_func_for_mt(func, params):\n    def wrap_func(queue_input, queue_output):\n        while True:\n            input = queue_input.get()\n            if input is None:\n                queue_output.put(None)\n                continue\n\n            result = func(input, params)\n\n            queue_output.put(result)\n\n    return wrap_func\n\ndef loop_proc(queues_input, queues_output, inputs):\n    for queue_input, input in zip(queues_input, inputs):\n        queue_input.put(input)\n\n    outputs = []\n    for queue_output in queues_output:\n        output = queue_output.get()\n        outputs.append(output)\n\n    return outputs\n\ndef prepare_multithreading(params):\n\n    # funcs to proc in pipeline\n    func_params = [\n        (feature_matching, (params)),\n        (reconstruction, (params)),\n    ]\n    wrap_funcs = list(map(lambda func_param: wrap_func_for_mt(func_param[0], func_param[1]), func_params))\n\n    # prepare queues\n    queues_input = [queue.Queue() for _ in range(len(wrap_funcs))]\n    queues_output = [queue.Queue() for _ in range(len(wrap_funcs))]\n\n    # create Threads\n    threads = []\n    for wrap_func, queue_input, queue_output in zip(wrap_funcs, queues_input, queues_output):\n        t = threading.Thread(target=wrap_func, args=(queue_input, queue_output), daemon=True)\n        threads.append(t)\n\n    for t in threads:\n        t.start()\n\n    return queues_input, queues_output, len(wrap_funcs)\n","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:06.869716Z","iopub.execute_input":"2024-06-03T09:57:06.870000Z","iopub.status.idle":"2024-06-03T09:57:06.881275Z","shell.execute_reply.started":"2024-06-03T09:57:06.869977Z","shell.execute_reply":"2024-06-03T09:57:06.880369Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Main process","metadata":{}},{"cell_type":"code","source":"def get_pairs_exhaustive(lst):\n    return list(itertools.combinations(range(len(lst)), 2))","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:06.882481Z","iopub.execute_input":"2024-06-03T09:57:06.882749Z","iopub.status.idle":"2024-06-03T09:57:06.893147Z","shell.execute_reply.started":"2024-06-03T09:57:06.882727Z","shell.execute_reply":"2024-06-03T09:57:06.892462Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def run_from_config(config):\n\n    data_dict, category_dict = parse_sample_submission(DATA_TYPE, SCENE_TYPE)\n    datasets = list(data_dict.keys())\n\n    ############################\n    # Pipeline process\n    dataset_list, scene_list, trial_list = prepare_input(data_dict, config)\n\n    queues_input, queues_output, len_wrap_funcs = prepare_multithreading(config)\n\n    idx = 0\n    results_trial_list = []\n    image_paths_list = {}\n    while len(results_trial_list) < len(dataset_list):\n\n        if idx >= len(dataset_list):\n            args = None\n        else:\n            dataset = dataset_list[idx]\n            scene = scene_list[idx]\n            trial_idx = trial_list[idx]\n\n            # use all images in the directory\n            images_dir = data_dict[dataset][scene][0].parent\n            image_paths = list(images_dir.glob(\"*\"))\n\n            if len(image_paths) > 100:\n                # random sampling\n                test_image_paths = data_dict[dataset][scene]\n                additional_image_paths = [path for path in image_paths if path not in test_image_paths]\n\n                random.shuffle(additional_image_paths)\n                additional_image_size = 100 - len(test_image_paths)\n                additional_image_paths = additional_image_paths[:additional_image_size]\n\n                image_paths = test_image_paths + additional_image_paths\n\n            if dataset not in image_paths_list:\n                image_paths_list[dataset] = {}\n\n            if scene not in image_paths_list[dataset]:\n                image_paths_list[dataset][scene] = {}\n\n            image_paths_list[dataset][scene][trial_idx] = image_paths\n                \n            args = {\n                \"trial_index\": trial_idx,\n                \"image_paths\": image_paths,\n                \"categories\": category_dict[dataset][scene],\n                \"feature_dir\": config.feature_dir / f\"{dataset}_{scene}_trial{trial_idx}\",\n            }\n\n        if idx == 0:\n            init_inputs = [args] + [None]*(len_wrap_funcs - 1)  # [[], None, None, ...]\n            inputs = init_inputs\n        else:\n            inputs = [args] + outputs[:-1]\n\n        outputs = loop_proc(queues_input, queues_output, inputs)\n        result = outputs[-1]\n\n        if result is not None:\n            results_trial_list.append(result)\n\n        idx = idx + 1\n\n    ############################\n    # Pipeline post-processing\n    results_trial = {}\n    for dataset, scene, trial_idx, results in zip(dataset_list, scene_list, trial_list, results_trial_list):\n\n        if dataset not in results_trial:\n            results_trial[dataset] = {}\n\n        if scene not in results_trial[dataset]:\n            results_trial[dataset][scene] = {}\n\n        results_trial[dataset][scene][trial_idx] = results\n\n    ############################\n    # Gather the best results with the most images registered and the most 3D points\n    for dataset in datasets:\n        for scene in data_dict[dataset]:\n\n            best_trial_ID = 0\n            images_registered_trial = 0\n            num_points3D_trial = 0\n            maps_trial = None\n            best_idx_trial = None\n            for trial_idx in results_trial[dataset][scene].keys():\n\n                recon_result = results_trial[dataset][scene][trial_idx]\n\n                if recon_result is not None:\n\n                    flag = False\n                    if recon_result[\"images_registered\"] > images_registered_trial:\n                        flag = True\n                    elif recon_result[\"images_registered\"] == images_registered_trial:\n                        if recon_result[\"num_points3D\"] > num_points3D_trial:\n                            flag = True\n\n                    if flag:\n                        best_trial_ID = trial_idx\n                        images_registered_trial = recon_result[\"images_registered\"]\n                        num_points3D_trial = recon_result[\"num_points3D\"]\n                        best_idx_trial = recon_result[\"best_idx\"]\n                        maps_trial = deepcopy(recon_result[\"maps\"])\n\n            if best_idx_trial is not None:\n                print(f\"Best Trial is {best_trial_ID}\")\n                print(f\"Best idx:{best_idx_trial}, num_images:{images_registered_trial}, num_points3D:{num_points3D_trial}\")\n\n                results_trial[dataset][scene][\"best_trial_ID\"] = best_trial_ID\n                results_trial[dataset][scene][\"best_idx\"] = best_idx_trial\n                results_trial[dataset][scene][\"maps\"] = deepcopy(maps_trial)\n\n    ##########################\n    # Parse the reconstruction object to get the rotation matrix and translation vector\n    # symmetries-and-repeats: refine the reconstruction\n    results = {}\n    for dataset in datasets:\n\n        if dataset not in results:\n            results[dataset] = {}\n\n        for scene in data_dict[dataset]:\n\n            best_idx = None\n            categories = category_dict[dataset][scene]\n\n            if (\"symmetries-and-repeats\" in categories) and (\"transparent\" not in categories):\n\n                best_trial_ID = results_trial[dataset][scene][\"best_trial_ID\"]\n                best_idx = results_trial[dataset][scene][\"best_idx\"]\n\n                feature_dir = config.feature_dir / f\"{dataset}_{scene}_trial{best_trial_ID}\"\n                recon_data_dir = feature_dir / \"colmap_rec_aliked\" / f\"{best_idx}\"\n                image_paths = image_paths_list[dataset][scene][best_trial_ID]\n\n                _result = refine_symmetries_and_repeats(\n                    config,\n                    feature_dir,\n                    image_paths,\n                    recon_data_dir,\n                )\n\n                if _result is not None:\n                    maps = _result[\"maps\"]\n                    best_idx = _result[\"best_idx\"]\n\n            else:\n                maps = results_trial[dataset][scene][\"maps\"]\n                best_idx = results_trial[dataset][scene][\"best_idx\"]\n\n            # Parse the reconstruction object to get the rotation matrix and translation vector\n            # obtained for each image in the reconstruction\n            results[dataset][scene] = {}\n            if best_idx is not None:\n                for k, im in maps[best_idx].images.items():\n                    key = config.base_path / f\"{DATA_TYPE}\" / 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            gc.collect()\n\n\n    if not DEBUG:\n        for dataset in datasets:\n            for scene in data_dict[dataset]:\n                for trial_idx in range(config.num_trials):\n                    feature_dir = config.feature_dir / f\"{dataset}_{scene}_trial{trial_idx}\"\n                    if feature_dir.exists():\n                        shutil.rmtree(feature_dir)\n\n    create_submission(output_dir, results, data_dict, config.base_path, SCENE_TYPE)","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:06.894479Z","iopub.execute_input":"2024-06-03T09:57:06.894745Z","iopub.status.idle":"2024-06-03T09:57:06.921245Z","shell.execute_reply.started":"2024-06-03T09:57:06.894722Z","shell.execute_reply":"2024-06-03T09:57:06.920355Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Config:\n    base_path = Path(\"/kaggle/input/image-matching-challenge-2024\")\n    feature_dir = output_dir / \".feature_outputs\"\n\n    device = K.utils.get_cuda_device_if_available(0)\n    num_gpus = torch.cuda.device_count()\n\n    num_trials = 3\n\n    keypoint_detection_args = {\n        \"num_features\": 4096,\n        \"detection_threshold\": 0.01,\n        \"resize_to\": 1024,\n    }\n\n    keypoint_detection_transparent_args = {\n        \"num_features\": 4096,\n        \"detection_threshold\": 0.5,\n        \"resize_to\": 1024,\n    }\n\n    keypoint_distances_args = {\n        \"early_stopping_thr\": 1000,\n        \"min_matches\": [100, 125, 150],\n        \"verbose\": False,\n    }\n\n    doppelgangers_classifier_args = {\n        \"loftr_weight_path\": Path(\"/kaggle/input/loftr/pytorch/outdoor/1/loftr_outdoor.ckpt\"),\n        \"doppelgangers_weight_path\": Path(\"/kaggle/input/doppelgangers-repo/doppelgangers/weights/doppelgangers_classifier_loftr.pt\"),\n    }\n\n    remove_ambiguous_area_args = {\n        \"pcd_used_num_images_ratio\": 0.5,\n        \"erode_mask_ratio\": 0.0,\n    }\n\n    verify_matches_args = {\n        \"doppelgangers_min_thr\": 0.3,\n        \"doppelgangers_max_thr\": 0.5,\n        \"filter_iterations\": 1,\n        \"filter_threshold\": 1,\n    }\n\n    colmap_mapper_options = {\n        \"min_model_size\": 3,\n        \"max_num_models\": 2,\n    }","metadata":{"execution":{"iopub.status.busy":"2024-06-03T09:57:06.922342Z","iopub.execute_input":"2024-06-03T09:57:06.922611Z","iopub.status.idle":"2024-06-03T09:57:06.959992Z","shell.execute_reply.started":"2024-06-03T09:57:06.922589Z","shell.execute_reply":"2024-06-03T09:57:06.959361Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if SAVE_FLAG:\n    !cp /kaggle/input/image-matching-challenge-2024/sample_submission.csv /kaggle/working/submission.csv\nelse:\n    run_from_config(Config)","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2024-06-03T09:57:06.960837Z","iopub.execute_input":"2024-06-03T09:57:06.961087Z"},"trusted":true},"execution_count":null,"outputs":[]}]}