{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.10","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":49349,"databundleVersionId":5447706,"sourceType":"competition"},{"sourceId":11946120,"sourceType":"datasetVersion","datasetId":7510065},{"sourceId":11971685,"sourceType":"datasetVersion","datasetId":7528281},{"sourceId":11999320,"sourceType":"datasetVersion","datasetId":7548126}],"dockerImageVersionId":30498,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!python --version\n!pip show pycolmap","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install /kaggle/input/kaggle-wheels/kaggle_wheels/pycolmap-0.5.0-cp310-cp310-linux_x86_64.whl --no-index\n!pip install /kaggle/input/kaggle-wheels/kaggle_wheels/einops-0.8.1-py3-none-any.whl --no-index\n!pip install /kaggle/input/kaggle-wheels/kaggle_wheels/flash_attn-2.6.3+cu118torch2.0cxx11abiTRUE-cp310-cp310-linux_x86_64.whl --no-index","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pycolmap\noptions = pycolmap.IncrementalMapperOptions()\n# print(dir(options))","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport cv2\nimport sys\nimport h5py\nimport timm\nimport torch\nimport shutil\nimport sqlite3\nimport warnings\nimport pycolmap\nimport itertools\nimport numpy as np\nimport configparser\nimport pandas as pd\nfrom glob import glob\nfrom tqdm import tqdm\nfrom copy import deepcopy\nfrom PIL import Image, ExifTags\nimport torch.nn.functional as F\nfrom collections import defaultdict\nfrom timm.data import resolve_data_config\nfrom timm.data.transforms_factory import create_transform\nfrom pathlib import Path\n\nsys.path.append('/kaggle/input')\nfrom sppluslg.matching_lg import Matching\n\nINPUT_ROOT = '/kaggle/input/image-matching-challenge-2023'\nDATA_ROOT = '/kaggle/data'\nOUTPUT_ROOT = '/kaggle/working'\nmatching_name = 'SuperGlue'\nimage_size = 1460\nsimilarity_filter = True\nmanual_ransac = False","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"datasets_scenes = []\nsample_submission_df = pd.read_csv(f\"{INPUT_ROOT}/sample_submission.csv\")\nfor _, r in sample_submission_df[['dataset', 'scene']].iterrows():\n    ds = f\"{r.dataset}/{r.scene}\"\n    if ds not in datasets_scenes:\n        datasets_scenes.append(ds)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if similarity_filter:\n    similarity_model = timm.create_model('tf_efficientnet_b7', checkpoint_path='/kaggle/input/sppluslg/weights/tf_efficientnet_b7_ra-6c08e654.pth').cuda().half().eval()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"matching_config = {\n    'superpoint': {\n        'nms_radius': 3,\n        'keypoint_threshold': 0.001,\n        'max_keypoints': 2048\n    },\n    'superglue': {\n        'weights': 'outdoor',\n        'sinkhorn_iterations': 20,\n        'match_threshold': 0.2,\n    },\n    'lightglue': {\n        'depth_confidence': 0.95,\n        'width_confidence': 0.99,\n        'filter_threshold': 0.1\n    }\n}\n\nmatching_model = Matching(matching_config).cuda().eval()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_img_pairs_all(fnames):\n    \"\"\"\n    Generate pairs of indices for all possible combinations of image filenames.\n\n    Args:\n        filenames (list): List of image filenames.\n\n    Returns:\n        list: List of index pairs representing all possible combinations of image indices.\n    \"\"\"    \n    index_pairs = []\n    for i in range(len(fnames)):\n        for j in range(i+1, len(fnames)):\n            index_pairs.append((i,j))\n    return index_pairs\n\n\ndef get_global_desc(model, fnames):\n    \"\"\"\n    Get global descriptors for a list of image filenames using a similarity model.\n\n    Args:\n        model (torch.nn.Module): Similarity model.\n        filenames (list): List of image filenames.\n\n    Returns:\n        torch.Tensor: Global descriptors for all images.\n    \"\"\"    \n    config = resolve_data_config({}, model=model)\n    transform = create_transform(**config)\n    global_descs_convnext=[]\n    for fname in tqdm(fnames, desc='Get global features using similarity model'):\n        img = Image.open(fname).convert('RGB')\n        timg = transform(img).unsqueeze(0).cuda().half()\n        with torch.no_grad():\n            desc = model.forward_features(timg.cuda().half()).mean(dim=(-1,2))\n            desc = desc.view(1, -1)\n            desc_norm = F.normalize(desc, dim=1, p=2)\n        global_descs_convnext.append(desc_norm.detach().cpu())\n    global_descs_all = torch.cat(global_descs_convnext, dim=0)\n    return global_descs_all\n\n\ndef get_image_pairs_filtered(model, fnames, sim_th=0.5, min_pairs=20, all_if_less=20):\n    \"\"\"\n    Generate pairs of image indices based on similarity filtering using global descriptors.\n\n    Args:\n        model (torch.nn.Module): Similarity model.\n        filenames (list): List of image filenames.\n        similarity_threshold (float): Similarity threshold for filtering. Default is 0.5.\n        min_pairs (int): Minimum number of pairs to generate if the number of images is below all_if_less. Default is 20.\n        all_if_less (int): If the number of images is less than or equal to all_if_less, return all possible pairs. Default is 20.\n\n    Returns:\n        tuple: A tuple containing a list of matching pairs of image indices and a distance matrix.\n    \"\"\"   \n\n    num_imgs = len(fnames)\n\n    if num_imgs <= all_if_less:\n        return get_img_pairs_all(fnames), None\n\n    descs = get_global_desc(model, fnames).type(torch.FloatTensor)\n    dm = torch.cdist(descs, descs, p=2).detach().cpu().numpy()\n\n    mask = dm <= sim_th\n    total = 0\n    matching_list = []\n    ar = np.arange(num_imgs)\n    for st_idx in range(num_imgs-1):\n        mask_idx = mask[st_idx]\n        to_match = ar[mask_idx]\n        if len(to_match) < min_pairs:\n            to_match = np.argsort(dm[st_idx])[:min_pairs]  \n        for idx in to_match:\n            if st_idx == idx:\n                continue\n            if dm[st_idx, idx] < 1200:\n                matching_list.append(tuple(sorted((st_idx, idx.item()))))\n                total+=1\n    matching_list = sorted(list(set(matching_list)))\n\n    return matching_list, dm\n\n\ndef get_unique_idxs(A, dim=0):\n    \"\"\"\n    Get the indices of the first occurrence of unique elements along the specified dimension of the input tensor.\n\n    Args:\n        input_tensor (torch.Tensor): Input tensor.\n        dim (int): Dimension along which to find unique elements. Default is 0.\n\n    Returns:\n        torch.Tensor: Indices of the first occurrence of unique elements.\n    \"\"\"    \n    _, idx, counts = torch.unique(A, dim=dim, sorted=True, return_inverse=True, return_counts=True)\n    _, ind_sorted = torch.sort(idx, stable=True)\n    cum_sum = counts.cumsum(0)\n    cum_sum = torch.cat((torch.tensor([0],device=cum_sum.device), cum_sum[:-1]))\n    first_indices = ind_sorted[cum_sum]\n    return first_indices","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Above images show that KeyNetAffNetHardNet + Adalam can properly handle rotated features matching with low mean distance.","metadata":{}},{"cell_type":"code","source":"# Default settings\nMAX_IMAGE_ID = 2**31 - 1\n# from https://github.com/colmap/colmap/blob/dev/scripts/python/database.py\n\nCREATE_CAMERAS_TABLE = \"\"\"CREATE TABLE IF NOT EXISTS cameras (\n    camera_id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,\n    model INTEGER NOT NULL,\n    width INTEGER NOT NULL,\n    height INTEGER NOT NULL,\n    params BLOB,\n    prior_focal_length INTEGER NOT NULL)\"\"\"\n\n\nCREATE_DESCRIPTORS_TABLE = \"\"\"CREATE TABLE IF NOT EXISTS descriptors (\n    image_id INTEGER PRIMARY KEY NOT NULL,\n    rows INTEGER NOT NULL,\n    cols INTEGER NOT NULL,\n    data BLOB,\n    FOREIGN KEY(image_id) REFERENCES images(image_id) ON DELETE CASCADE)\"\"\"\n\n\nCREATE_IMAGES_TABLE = \"\"\"CREATE TABLE IF NOT EXISTS images (\n    image_id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,\n    name TEXT NOT NULL UNIQUE,\n    camera_id INTEGER NOT NULL,\n    prior_qw REAL,\n    prior_qx REAL,\n    prior_qy REAL,\n    prior_qz REAL,\n    prior_tx REAL,\n    prior_ty REAL,\n    prior_tz REAL,\n    CONSTRAINT image_id_check CHECK(image_id >= 0 and image_id < {}),\n    FOREIGN KEY(camera_id) REFERENCES cameras(camera_id))\n\"\"\".format(MAX_IMAGE_ID)\n\n\nCREATE_TWO_VIEW_GEOMETRIES_TABLE = \"\"\"\nCREATE TABLE IF NOT EXISTS two_view_geometries (\n    pair_id INTEGER PRIMARY KEY NOT NULL,\n    rows INTEGER NOT NULL,\n    cols INTEGER NOT NULL,\n    data BLOB,\n    config INTEGER NOT NULL,\n    F BLOB,\n    E BLOB,\n    H BLOB)\n\"\"\"\n\n\nCREATE_KEYPOINTS_TABLE = \"\"\"CREATE TABLE IF NOT EXISTS keypoints (\n    image_id INTEGER PRIMARY KEY NOT NULL,\n    rows INTEGER NOT NULL,\n    cols INTEGER NOT NULL,\n    data BLOB,\n    FOREIGN KEY(image_id) REFERENCES images(image_id) ON DELETE CASCADE)\n\"\"\"\n\n\nCREATE_MATCHES_TABLE = \"\"\"CREATE TABLE IF NOT EXISTS matches (\n    pair_id INTEGER PRIMARY KEY NOT NULL,\n    rows INTEGER NOT NULL,\n    cols INTEGER NOT NULL,\n    data BLOB)\"\"\"\n\n\nCREATE_NAME_INDEX = \\\n    \"CREATE UNIQUE INDEX IF NOT EXISTS index_name ON images(name)\"\n\n\nCREATE_ALL = \"; \".join([\n    CREATE_CAMERAS_TABLE,\n    CREATE_IMAGES_TABLE,\n    CREATE_KEYPOINTS_TABLE,\n    CREATE_DESCRIPTORS_TABLE,\n    CREATE_MATCHES_TABLE,\n    CREATE_TWO_VIEW_GEOMETRIES_TABLE,\n    CREATE_NAME_INDEX\n])\n\n\ndef image_ids_to_pair_id(image_id1, image_id2):\n    if image_id1 > image_id2:\n        image_id1, image_id2 = image_id2, image_id1\n    return image_id1 * MAX_IMAGE_ID + image_id2\n\n\ndef array_to_blob(array):\n    return array.tostring()\n\n\nclass COLMAPDatabase(sqlite3.Connection):\n\n    @staticmethod\n    def connect(database_path):\n        return sqlite3.connect(database_path, factory=COLMAPDatabase)\n\n    def __init__(self, *args, **kwargs):\n        super(COLMAPDatabase, self).__init__(*args, **kwargs)\n\n        self.create_tables = lambda: self.executescript(CREATE_ALL)\n        self.create_cameras_table = \\\n            lambda: self.executescript(CREATE_CAMERAS_TABLE)\n        self.create_descriptors_table = \\\n            lambda: self.executescript(CREATE_DESCRIPTORS_TABLE)\n        self.create_images_table = \\\n            lambda: self.executescript(CREATE_IMAGES_TABLE)\n        self.create_two_view_geometries_table = \\\n            lambda: self.executescript(CREATE_TWO_VIEW_GEOMETRIES_TABLE)\n        self.create_keypoints_table = \\\n            lambda: self.executescript(CREATE_KEYPOINTS_TABLE)\n        self.create_matches_table = \\\n            lambda: self.executescript(CREATE_MATCHES_TABLE)\n        self.create_name_index = lambda: self.executescript(CREATE_NAME_INDEX)\n\n    def add_camera(self, model, width, height, params,\n                   prior_focal_length=False, camera_id=None):\n        params = np.asarray(params, np.float64)\n        cursor = self.execute(\n            \"INSERT INTO cameras VALUES (?, ?, ?, ?, ?, ?)\",\n            (camera_id, model, width, height, array_to_blob(params),\n             prior_focal_length))\n        return cursor.lastrowid\n\n    def add_image(self, name, camera_id,\n                  prior_q=np.zeros(4), prior_t=np.zeros(3), image_id=None):\n        cursor = self.execute(\n            \"INSERT INTO images VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)\",\n            (image_id, name, camera_id, prior_q[0], prior_q[1], prior_q[2],\n             prior_q[3], prior_t[0], prior_t[1], prior_t[2]))\n        return cursor.lastrowid\n\n    def add_keypoints(self, image_id, keypoints):\n        assert(len(keypoints.shape) == 2)\n        assert(keypoints.shape[1] in [2, 4, 6])\n\n        keypoints = np.asarray(keypoints, np.float32)\n        self.execute(\n            \"INSERT INTO keypoints VALUES (?, ?, ?, ?)\",\n            (image_id,) + keypoints.shape + (array_to_blob(keypoints),))\n\n    def add_matches(self, image_id1, image_id2, matches):\n        assert(len(matches.shape) == 2)\n        assert(matches.shape[1] == 2)\n\n        if image_id1 > image_id2:\n            matches = matches[:,::-1]\n\n        pair_id = image_ids_to_pair_id(image_id1, image_id2)\n        matches = np.asarray(matches, np.uint32)\n        self.execute(\n            \"INSERT INTO matches VALUES (?, ?, ?, ?)\",\n            (pair_id,) + matches.shape + (array_to_blob(matches),))\n\n    def add_two_view_geometry(self, image_id1, image_id2, matches,\n                              F=np.eye(3), E=np.eye(3), H=np.eye(3), config=2):\n        assert(len(matches.shape) == 2)\n        assert(matches.shape[1] == 2)\n\n        if image_id1 > image_id2:\n            matches = matches[:,::-1]\n\n        pair_id = image_ids_to_pair_id(image_id1, image_id2)\n        matches = np.asarray(matches, np.uint32)\n        F = np.asarray(F, dtype=np.float64)\n        E = np.asarray(E, dtype=np.float64)\n        H = np.asarray(H, dtype=np.float64)\n        self.execute(\n            \"INSERT INTO two_view_geometries VALUES (?, ?, ?, ?, ?, ?, ?, ?)\",\n            (pair_id,) + matches.shape + (array_to_blob(matches), config,\n             array_to_blob(F), array_to_blob(E), array_to_blob(H)))\n\n\ndef get_focal(image_path, err_on_default=False):\n    image         = Image.open(image_path)\n    max_size      = max(image.size)\n\n    exif = image.getexif()\n    focal = None\n    if exif is not None:\n        focal_35mm = None\n        # https://github.com/colmap/colmap/blob/d3a29e203ab69e91eda938d6e56e1c7339d62a99/src/util/bitmap.cc#L299\n        for tag, value in exif.items():\n            focal_35mm = None\n            if ExifTags.TAGS.get(tag, None) == 'FocalLengthIn35mmFilm':\n                focal_35mm = float(value)\n                break\n\n        if focal_35mm is not None:\n            focal = focal_35mm / 35. * max_size\n    \n    if focal is None:\n        if err_on_default:\n            raise RuntimeError(\"Failed to find focal length\")\n\n        # failed to find it in exif, use prior\n        FOCAL_PRIOR = 1.2\n        focal = FOCAL_PRIOR * max_size\n\n    return focal\n\n\ndef create_camera(db, image_path, camera_model):\n\n    image = Image.open(image_path)\n    width, height = image.size\n\n    focal = get_focal(image_path)\n\n    if camera_model == 'simple-pinhole':\n        model = 0 # simple pinhole\n        param_arr = np.array([focal, width / 2, height / 2])\n    if camera_model == 'pinhole':\n        model = 1 # pinhole\n        param_arr = np.array([focal, focal, width / 2, height / 2])\n    elif camera_model == 'simple-radial':\n        model = 2 # simple radial\n        param_arr = np.array([focal, width / 2, height / 2, 0.1])\n    elif camera_model == 'opencv':\n        model = 4 # opencv\n        param_arr = np.array([focal, focal, width / 2, height / 2, 0., 0., 0., 0.])\n         \n    return db.add_camera(model, width, height, param_arr)\n\n\ndef add_keypoints(db, feature_dir, img_dir, camera_model, single_camera=True):\n    keypoint_f = h5py.File(os.path.join(feature_dir, 'keypoints.h5'), 'r')\n\n    camera_id = None\n    fname_to_id = {}\n    for filename in tqdm(list(keypoint_f.keys())):\n        keypoints = keypoint_f[filename][()]\n\n        path = os.path.join(img_dir, filename)\n        if not os.path.isfile(path):\n            raise IOError(f'Invalid image path {path}')\n\n        if camera_id is None or not single_camera:\n            camera_id = create_camera(db, path, camera_model)\n        image_id = db.add_image(filename, camera_id)\n        fname_to_id[filename] = image_id\n\n        db.add_keypoints(image_id, keypoints)\n\n    return fname_to_id\n\n\ndef add_matches(db, feature_dir, fname_to_id, FH=None):\n\n    match_file = h5py.File(os.path.join(feature_dir, 'matches.h5'), 'r')\n    added = set()\n    if FH:\n        all_pair_ids = list(itertools.combinations(range(1,len(fname_to_id)+1), 2))\n\n    for key_1 in match_file.keys():\n        group = match_file[key_1]\n        for key_2 in group.keys():\n            id_1 = fname_to_id[key_1]\n            id_2 = fname_to_id[key_2]\n\n            pair_id = (id_1, id_2)\n            if pair_id in added:\n                warnings.warn(f'Pair {pair_id} ({id_1}, {id_2}) already added!')\n                continue\n            added.add(pair_id)\n\n            matches = group[key_2][()]\n            db.add_matches(id_1, id_2, matches)\n            if FH:\n                db.add_two_view_geometry(id_1, id_2, matches, F=FH[0][(key_1, key_2)], E=np.eye(3), H=FH[1][(key_1, key_2)], config=3)\n\n    if FH:\n        for pair_id in all_pair_ids:\n            if pair_id not in added:\n                id_1, id_2 = pair_id\n                db.add_matches(id_1, id_2, np.empty((0,2)))\n                db.add_two_view_geometry(id_1, id_2, np.empty((0,2)), config=0)\n\n\ndef import_into_colmap(img_dir, feature_dir='.featureout', FH=None):\n    db = COLMAPDatabase.connect(f\"{feature_dir}/colmap.db\")\n    db.create_tables()\n    fname_to_id = add_keypoints(db, feature_dir, img_dir, 'simple-radial', single_camera=False)\n    add_matches(db, feature_dir, fname_to_id, FH=FH)\n    db.commit()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_homography_matrix(source, destination):\n    \"\"\" Calculates the entries of the Homography matrix between two sets of matching points.\n    Args\n    ----\n        - `source`: Source points where each point is int (x, y) format.\n        - `destination`: Destination points where each point is int (x, y) format.\n    Returns\n    ----\n        - A numpy array of shape (3, 3) representing the Homography matrix.\n    Raises\n    ----\n        - `source` and `destination` is lew than four points.\n        - `source` and `destination` is of different size.\n    \"\"\"\n    assert len(source) >= 4, \"must provide more than 4 source points\"\n    assert len(destination) >= 4, \"must provide more than 4 destination points\"\n    assert len(source) == len(destination), \"source and destination must be of equal length\"\n    A = []\n    b = []\n    for i in range(len(source)):\n        s_x, s_y = source[i]\n        d_x, d_y = destination[i]\n        A.append([s_x, s_y, 1, 0, 0, 0, (-d_x)*(s_x), (-d_x)*(s_y)])\n        A.append([0, 0, 0, s_x, s_y, 1, (-d_y)*(s_x), (-d_y)*(s_y)])\n        b += [d_x, d_y]\n    A = np.array(A)\n    h = np.linalg.lstsq(A, b, rcond=None)[0]\n    h = np.concatenate((h, [1]), axis=-1)\n    return np.reshape(h, (3, 3))\n\n\ndef resize(image, image_size):\n    \"\"\"\n    Resize the image while maintaining the aspect ratio.\n\n    Args:\n        image (np.ndarray): Input image.\n        image_size (int): Target size of the image.\n\n    Returns:\n        np.ndarray: Resized image.\n        tuple: New size of the image.\n    \"\"\"    \n    h, w = image.shape[:2]\n    aspect_ratio = h/w\n    smaller_side_size = int(image_size/max(aspect_ratio, 1/aspect_ratio))\n    if aspect_ratio > 1: # H > W\n        new_size = (image_size, smaller_side_size)\n    else: # H <= W\n        new_size = (smaller_side_size, image_size)\n    image = cv2.resize(image, new_size[::-1])\n    return image, new_size\n\n\ndef superglue_inference(model, img1, img2):\n    \"\"\"\n    Perform inference using the SuperGlue model.\n\n    Args:\n        model (torch.nn.Module): SuperGlue model.\n        img1 (torch.Tensor): Image 1.\n        img2 (torch.Tensor): Image 2.\n\n    Returns:\n        np.ndarray: Matched keypoints from image 1.\n        np.ndarray: Matched keypoints from image 2.\n    \"\"\"\n    with torch.no_grad():\n        pred = model({'image0': img1, 'image1': img2})\n\n    kpts1, kpts2 = pred['keypoints0'][0].cpu().numpy(), pred['keypoints1'][0].cpu().numpy()\n    matches = pred['matches0'][0].cpu().numpy()\n    valid_matches = matches > -1\n    mkpts1 = kpts1[valid_matches]\n    mkpts2 = kpts2[matches[valid_matches]]\n\n    return mkpts1, mkpts2\n\ndef rbd(data: dict) -> dict:\n    \"\"\"Remove batch dimension from elements in data\"\"\"\n    return {\n        k: v[0] if isinstance(v, (torch.Tensor, np.ndarray, list)) else v\n        for k, v in data.items()\n    }\n\ndef lightglue_inference(model, img1, img2):\n    with torch.no_grad():\n        pred = model({'image0': img1, 'image1': img2, 'image_size': image_size})\n        feats0, feats1, matches01 = [\n            rbd(x) for x in [pred['feats0'], pred['feats1'], pred['matches01']]\n        ]  # remove batch dimension\n        \n        kpts0, kpts1, matches = feats0[\"keypoints\"], feats1[\"keypoints\"], matches01[\"matches\"]\n        m_kpts0, m_kpts1 = kpts0[matches[..., 0]].cpu().numpy(), kpts1[matches[..., 1]].cpu().numpy()\n\n    return m_kpts0, m_kpts1\n\n\n\ndef matching_inference(model, fname1, fname2, cache=None):\n    \"\"\"\n    Perform matching inference using the matching model.\n\n    Args:\n        model (torch.nn.Module): Matching model.\n        filename1 (str): Path to the first image file.\n        filename2 (str): Path to the second image file.\n        cache (dict): Cache dictionary for storing preprocessed images. Default is None.\n\n    Returns:\n        np.ndarray: Matched keypoints from image 1.\n        np.ndarray: Matched keypoints from image 2.\n    \"\"\"\n    for fname in [fname1, fname2]:\n        if fname not in cache:\n            img = cv2.imread(fname, 0)\n            h, w = h_r, w_r = img.shape[:2]\n            if max(h, w) != image_size:\n                img, (h_r, w_r) = resize(img, image_size)\n            cache[fname] = {'img': fname, 'h': h, 'w': w, 'h_r': h_r, 'w_r': w_r}\n        \n    mkpts1, mkpts2 = lightglue_inference(model, cache[fname1]['img'], cache[fname2]['img'])\n\n    if max(cache[fname1]['h'], cache[fname1]['w']) != image_size:\n        mkpts1[:,0] *= cache[fname1]['w']/cache[fname1]['w_r']\n        mkpts1[:,1] *= cache[fname1]['h']/cache[fname1]['h_r']\n    if max(cache[fname2]['h'], cache[fname2]['w']) != image_size:\n        mkpts2[:,0] *= cache[fname2]['w']/cache[fname2]['w_r']\n        mkpts2[:,1] *= cache[fname2]['h']/cache[fname2]['h_r']\n\n    return mkpts1, mkpts2\n\n\ndef matching_pipeline(matching_model, fnames, index_pairs, feature_dir, manual_ransac=False):\n\n    cache = {}\n    with h5py.File(f\"{feature_dir}/matches_{matching_name}.h5\", mode='w') as f_match:\n\n        for pair_idx in tqdm(index_pairs, desc='Get matched keypoints using matching model'):\n            idx1, idx2 = pair_idx\n            fname1, fname2 = fnames[idx1], fnames[idx2]\n            key1, key2 = fname1.split('/')[-1], fname2.split('/')[-1]\n\n            mkpts1, mkpts2 = matching_inference(matching_model, fname1, fname2, cache)\n\n            n_matches = len(mkpts2)\n            group  = f_match.require_group(key1)\n            # print(f'{pair_idx}: {n_matches}')\n            if n_matches >= 150:\n                group.create_dataset(key2, data=np.concatenate([mkpts1, mkpts2], axis=1))\n\n    kpts = defaultdict(list)\n    total_kpts = defaultdict(int)\n    match_indexes = defaultdict(dict)\n\n    with h5py.File(f\"{feature_dir}/matches_{matching_name}.h5\", mode='r') as f_match:\n        for k1 in f_match.keys():\n            group  = f_match[k1]\n            for k2 in group.keys():\n                matches = group[k2][...]\n                total_kpts[k1]\n                kpts[k1].append(matches[:, :2])\n                kpts[k2].append(matches[:, 2:])\n                current_match = torch.arange(len(matches)).reshape(-1, 1).repeat(1, 2)\n                current_match[:, 0] += total_kpts[k1]\n                current_match[:, 1] += total_kpts[k2]\n                total_kpts[k1] += len(matches)\n                total_kpts[k2] += len(matches)\n                match_indexes[k1][k2] = current_match\n\n    for k in kpts.keys():\n        kpts[k] = np.round(np.concatenate(kpts[k], axis=0))\n\n    unique_kpts = {}\n    unique_match_idxs = {}\n    for k in kpts.keys():\n        uniq_kps, uniq_reverse_idxs = torch.unique(torch.from_numpy(kpts[k].astype(np.float32)), dim=0, return_inverse=True)\n        unique_match_idxs[k] = uniq_reverse_idxs\n        unique_kpts[k] = uniq_kps.numpy()\n\n    with h5py.File(f\"{feature_dir}/keypoints.h5\", mode='w') as f_kp:\n        for k, kpts1 in unique_kpts.items():\n            f_kp[k] = kpts1\n\n    out_match = defaultdict(dict)\n\n    for k1, group in match_indexes.items():\n        for k2, m in group.items():\n            m2 = deepcopy(m)\n            m2[:,0] = unique_match_idxs[k1][m2[:,0]]\n            m2[:,1] = unique_match_idxs[k2][m2[:,1]]\n            mkpts = np.concatenate([unique_kpts[k1][m2[:,0]], unique_kpts[k2][m2[:,1]]], axis=1)\n            unique_idxs_current = get_unique_idxs(torch.from_numpy(mkpts), dim=0)\n            m2_semiclean = m2[unique_idxs_current]\n            unique_idxs_current1 = get_unique_idxs(m2_semiclean[:, 0], dim=0)\n            m2_semiclean = m2_semiclean[unique_idxs_current1]\n            unique_idxs_current2 = get_unique_idxs(m2_semiclean[:, 1], dim=0)\n            m2_semiclean2 = m2_semiclean[unique_idxs_current2]\n            out_match[k1][k2] = m2_semiclean2.numpy()\n\n    with h5py.File(f\"{feature_dir}/matches.h5\", mode='w') as f_match:\n        for k1, gr in out_match.items():\n            group  = f_match.require_group(k1)\n            for k2, match in gr.items():\n                group[k2] = match\n\n    return None\n\n\ndef colmap_pipeline(img_dir, feature_dir, FH=None):\n\n    import_into_colmap(img_dir, feature_dir=feature_dir, FH=FH)\n\n    database_path=f\"{feature_dir}/colmap.db\"\n    if FH is None:\n        pycolmap.match_exhaustive(database_path)\n\n    output_path = f\"{feature_dir}/colmap_rec_{matching_name}\"\n    os.makedirs(output_path)\n\n    mapper_options = pycolmap.IncrementalMapperOptions()\n    mapper_options.min_model_size = 3\n    mapper_options.num_threads = 1\n    maps = pycolmap.incremental_mapping(database_path=database_path, image_path=img_dir, output_path=output_path, options=mapper_options)\n\n    return maps\n\n\ndef postprocessing(maps, dataset, scene):\n\n    results = {}\n    imgs_registered  = 0\n    best_idx = None\n    print (\"Looking for the best reconstruction\")\n    if isinstance(maps, dict):\n        for idx1, rec in maps.items():\n            print(idx1, rec.summary())\n            if len(rec.images) > imgs_registered:\n                imgs_registered = len(rec.images)\n                best_idx = idx1\n    if best_idx is not None:\n        print(maps[best_idx].summary())\n        for im in maps[best_idx].images.values():\n            key1 = f'{dataset}/{scene}/images/{im.name}'\n            results[key1] = {}\n            results[key1][\"R\"] = im.rotmat()\n            results[key1][\"t\"] = im.tvec\n\n    print(f'Registered: {dataset} / {scene} -> {len(results)} images')\n\n    return results\n\n\ndef arr_to_str(a):\n    return ';'.join([str(x) for x in a.reshape(-1)])","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# !mkdir -p ./2cfa01ab573141e4\n# !mkdir -p ./2cfa01ab573141e4/2fa124afd1f74f38\n# !mkdir -p ./2cfa01ab573141e4/2fa124afd1f74f38/images\n# !cp /kaggle/input/image-matching-challenge-2023/train/urban/kyiv-puppet-theater/images/IMG_20220127_165418.jpg ./2cfa01ab573141e4/2fa124afd1f74f38/images/IMG_20220127_165418.jpg\n# !cp /kaggle/input/image-matching-challenge-2023/train/urban/kyiv-puppet-theater/images/IMG_20220127_165612.jpg ./2cfa01ab573141e4/2fa124afd1f74f38/images/IMG_20220127_165612.jpg\n# !cp /kaggle/input/image-matching-challenge-2023/train/urban/kyiv-puppet-theater/images/IMG_20220127_165703.jpg ./2cfa01ab573141e4/2fa124afd1f74f38/images/IMG_20220127_165703.jpg\n# !cp /kaggle/input/image-matching-challenge-2023/train/urban/kyiv-puppet-theater/images/IMG_20220127_170633.jpg ./2cfa01ab573141e4/2fa124afd1f74f38/images/IMG_20220127_170633.jpg\n# !cp /kaggle/input/image-matching-challenge-2023/train/urban/kyiv-puppet-theater/images/IMG_20220127_170709.jpg ./2cfa01ab573141e4/2fa124afd1f74f38/images/IMG_20220127_170709.jpg\n# !cp /kaggle/input/image-matching-challenge-2023/train/urban/kyiv-puppet-theater/images/IMG_20220127_170022.jpg ./2cfa01ab573141e4/2fa124afd1f74f38/images/IMG_20220127_170022.jpg\n# !cp /kaggle/input/image-matching-challenge-2023/train/urban/kyiv-puppet-theater/images/IMG_20220127_170250.jpg ./2cfa01ab573141e4/2fa124afd1f74f38/images/IMG_20220127_170250.jpg\n# !cp /kaggle/input/image-matching-challenge-2023/train/urban/kyiv-puppet-theater/images/IMG_20220127_165343.jpg ./2cfa01ab573141e4/2fa124afd1f74f38/images/IMG_20220127_165343.jpg\n# !cp /kaggle/input/image-matching-challenge-2023/train/urban/kyiv-puppet-theater/images/IMG_20220127_165608.jpg ./2cfa01ab573141e4/2fa124afd1f74f38/images/IMG_20220127_165608.jpg\n# !cp /kaggle/input/image-matching-challenge-2023/train/urban/kyiv-puppet-theater/images/IMG_20220127_165700.jpg ./2cfa01ab573141e4/2fa124afd1f74f38/images/IMG_20220127_165700.jpg\n# !cp /kaggle/input/image-matching-challenge-2023/train/urban/kyiv-puppet-theater/images/IMG_20220127_165449.jpg ./2cfa01ab573141e4/2fa124afd1f74f38/images/IMG_20220127_165449.jpg\n# !cp /kaggle/input/image-matching-challenge-2023/train/urban/kyiv-puppet-theater/images/IMG_20220127_165811.jpg ./2cfa01ab573141e4/2fa124afd1f74f38/images/IMG_20220127_165811.jpg\n# !cp /kaggle/input/image-matching-challenge-2023/train/urban/kyiv-puppet-theater/images/IMG_20220127_165845.jpg ./2cfa01ab573141e4/2fa124afd1f74f38/images/IMG_20220127_165845.jpg\n# !cp /kaggle/input/image-matching-challenge-2023/train/urban/kyiv-puppet-theater/images/IMG_20220127_165722.jpg ./2cfa01ab573141e4/2fa124afd1f74f38/images/IMG_20220127_165722.jpg\n# !cp /kaggle/input/image-matching-challenge-2023/train/urban/kyiv-puppet-theater/images/IMG_20220127_165737.jpg ./2cfa01ab573141e4/2fa124afd1f74f38/images/IMG_20220127_165737.jpg\n# !cp /kaggle/input/image-matching-challenge-2023/train/urban/kyiv-puppet-theater/images/IMG_20220127_165748.jpg ./2cfa01ab573141e4/2fa124afd1f74f38/images/IMG_20220127_165748.jpg\n# !cp /kaggle/input/image-matching-challenge-2023/train/urban/kyiv-puppet-theater/images/IMG_20220127_165815.jpg ./2cfa01ab573141e4/2fa124afd1f74f38/images/IMG_20220127_165815.jpg\n# !cp /kaggle/input/image-matching-challenge-2023/train/urban/kyiv-puppet-theater/images/IMG_20220127_170339.jpg ./2cfa01ab573141e4/2fa124afd1f74f38/images/IMG_20220127_170339.jpg\n# !cp /kaggle/input/image-matching-challenge-2023/train/urban/kyiv-puppet-theater/images/IMG_20220127_170350.jpg ./2cfa01ab573141e4/2fa124afd1f74f38/images/IMG_20220127_170350.jpg\n# !cp /kaggle/input/image-matching-challenge-2023/train/urban/kyiv-puppet-theater/images/IMG_20220127_165920.jpg ./2cfa01ab573141e4/2fa124afd1f74f38/images/IMG_20220127_165920.jpg\n# !cp /kaggle/input/image-matching-challenge-2023/train/urban/kyiv-puppet-theater/images/IMG_20220127_165822.jpg ./2cfa01ab573141e4/2fa124afd1f74f38/images/IMG_20220127_165822.jpg\n# !cp /kaggle/input/image-matching-challenge-2023/train/urban/kyiv-puppet-theater/images/IMG_20220127_165709.jpg ./2cfa01ab573141e4/2fa124afd1f74f38/images/IMG_20220127_165709.jpg\n# !cp /kaggle/input/image-matching-challenge-2023/train/urban/kyiv-puppet-theater/images/IMG_20220127_170001.jpg ./2cfa01ab573141e4/2fa124afd1f74f38/images/IMG_20220127_170001.jpg\n# !cp /kaggle/input/image-matching-challenge-2023/train/urban/kyiv-puppet-theater/images/IMG_20220127_165957.jpg ./2cfa01ab573141e4/2fa124afd1f74f38/images/IMG_20220127_165957.jpg\n# !cp /kaggle/input/image-matching-challenge-2023/train/urban/kyiv-puppet-theater/images/IMG_20220127_165641.jpg ./2cfa01ab573141e4/2fa124afd1f74f38/images/IMG_20220127_165641.jpg\n# !cp /kaggle/input/image-matching-challenge-2023/train/urban/kyiv-puppet-theater/images/IMG_20220127_165549.jpg ./2cfa01ab573141e4/2fa124afd1f74f38/images/IMG_20220127_165549.jpg\n# !cp /kaggle/input/image-matching-challenge-2023/train/urban/kyiv-puppet-theater/images/IMG_20220127_170459.jpg ./2cfa01ab573141e4/2fa124afd1f74f38/images/IMG_20220127_170459.jpg","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"results_df = pd.DataFrame(columns=['image_path', 'dataset', 'scene', 'rotation_matrix', 'translation_vector'])\nfor dataset_scene in datasets_scenes:\n    \n    dataset, scene = dataset_scene.split('/')\n    print(f\"{dataset=}, {scene=}\")\n\n    img_dir = f\"{INPUT_ROOT}/test/{dataset}/{scene}/images\"\n    #img_dir = f\"./2cfa01ab573141e4/2fa124afd1f74f38/images\"\n    \n    if not os.path.exists(img_dir):\n        continue\n    \n    feature_dir = f\"{DATA_ROOT}/featureout/{dataset}/{scene}\"\n    os.makedirs(feature_dir, exist_ok=True)\n\n    fnames = sorted(glob(f\"{img_dir}/*\"))\n\n    # Similarity pipeline\n    if similarity_filter:\n        index_pairs, distance_matrix = get_image_pairs_filtered(similarity_model, fnames=fnames, sim_th=2.2, min_pairs=20, all_if_less=20)\n        if distance_matrix is not None:\n            distances = {fname: np.argsort(distance_matrix[idx])[1:] for idx, fname in enumerate(fnames)}\n    else:\n        index_pairs = get_img_pairs_all(fnames=fnames)\n\n    # Matching pipeline\n    FH = matching_pipeline(matching_model=matching_model,\n                           fnames=fnames,\n                           index_pairs=index_pairs,\n                           feature_dir=feature_dir,\n                           manual_ransac=manual_ransac)\n\n    # Colmap pipeline\n    maps = colmap_pipeline(img_dir, feature_dir, FH=FH)\n\n    # Postprocessing\n    results = postprocessing(maps, dataset, scene)\n\n    # Create submission\n    for fname in fnames:\n        image_id = '/'.join(fname.split('/')[-4:])\n        if image_id in results:\n            R = results[image_id]['R'].reshape(-1)\n            T = results[image_id]['t'].reshape(-1)\n        else:\n            R = np.eye(3).reshape(-1)\n            T = np.zeros((3))\n\n        new_row = pd.DataFrame({'image_path': image_id,\n                                'dataset': dataset,\n                                'scene': scene,\n                                'rotation_matrix': arr_to_str(R),\n                                'translation_vector': arr_to_str(T)}, index=[0])\n\n        results_df = pd.concat([results_df, new_row]).reset_index(drop=True)\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"results_df.to_csv(f\"{OUTPUT_ROOT}/submission.csv\", index=False)\nprint(\"file save1\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}