{"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":"none","dataSources":[{"sourceId":71885,"databundleVersionId":8143495,"sourceType":"competition"},{"sourceId":2187279,"sourceType":"datasetVersion","datasetId":767955}],"dockerImageVersionId":30698,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"%cd /kaggle/working/\n!rm -rf /kaggle/working/Hierarchical-Localization\n!git clone --quiet --recursive https://github.com/cvg/Hierarchical-Localization/\n%cd /kaggle/working/Hierarchical-Localization\n!pip install -e .\n\nfrom hloc import extract_features, match_features, reconstruction, visualization, pairs_from_exhaustive\nfrom hloc.visualization import plot_images, read_image\nfrom hloc.utils import viz_3d\nimport pycolmap\n%cd /kaggle/working/\nimport os\n# import pycolmap\nimport numpy as np\nimport cv2\nfrom glob import glob\nfrom pathlib import Path\nfrom time import time\nfrom pprint import pprint\nfrom PIL import Image\n\n!pip install check_orientation  #> /dev/null\nfrom check_orientation.pre_trained_models import create_model\nimport torch","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-05-08T08:25:21.827618Z","iopub.execute_input":"2024-05-08T08:25:21.82816Z","iopub.status.idle":"2024-05-08T08:26:30.984478Z","shell.execute_reply.started":"2024-05-08T08:25:21.828123Z","shell.execute_reply":"2024-05-08T08:26:30.982726Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"imagesDirList = glob('/kaggle/input/image-matching-challenge-2024/train/*/images')\npprint(imagesDirList)\n\n# %cd /kaggle/input/image-matching-challenge-2024\n\ndirIdx = 4\nsrc = imagesDirList[dirIdx]\nimages = [cv2.cvtColor(cv2.imread(im), cv2.COLOR_BGR2RGB) for im in glob(f'{src}/*')]\nprint(f'Showing {len(images)} image from {src}')","metadata":{"execution":{"iopub.status.busy":"2024-05-08T08:26:30.987674Z","iopub.execute_input":"2024-05-08T08:26:30.989611Z","iopub.status.idle":"2024-05-08T08:26:33.406666Z","shell.execute_reply.started":"2024-05-08T08:26:30.989424Z","shell.execute_reply":"2024-05-08T08:26:33.40543Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\nfrom torch.utils.data import Dataset\nfrom torchvision.io import read_image\n\nclass IMCDataset(Dataset):\n    def __init__(self, imgDir, transform=None):\n        self.imgDir = imgDir\n        self.transform = transform\n        self.imgPathList = [imgPath for imgPath in glob(f'{imgDir}/*')]\n\n    def __len__(self):\n        return len(self.imgPathList)\n\n    def __getitem__(self, idx):\n        imgPath = self.imgPathList[idx]\n        image = read_image(imgPath)\n        if self.transform:\n            image = self.transform(image)\n        return image\n    ","metadata":{"execution":{"iopub.status.busy":"2024-05-08T08:26:33.408109Z","iopub.execute_input":"2024-05-08T08:26:33.409497Z","iopub.status.idle":"2024-05-08T08:26:33.418384Z","shell.execute_reply.started":"2024-05-08T08:26:33.409435Z","shell.execute_reply":"2024-05-08T08:26:33.417439Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torchvision import transforms as T\ntransform = 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])","metadata":{"execution":{"iopub.status.busy":"2024-05-08T08:26:33.420951Z","iopub.execute_input":"2024-05-08T08:26:33.422198Z","iopub.status.idle":"2024-05-08T08:26:33.447846Z","shell.execute_reply.started":"2024-05-08T08:26:33.422159Z","shell.execute_reply":"2024-05-08T08:26:33.446232Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"heritageDioscuriDataset = IMCDataset(imagesDirList[dirIdx], transform=transform)","metadata":{"execution":{"iopub.status.busy":"2024-05-08T08:26:33.449322Z","iopub.execute_input":"2024-05-08T08:26:33.449874Z","iopub.status.idle":"2024-05-08T08:26:33.466055Z","shell.execute_reply.started":"2024-05-08T08:26:33.4497Z","shell.execute_reply":"2024-05-08T08:26:33.465035Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = create_model(\"swsl_resnext50_32x4d\")\nmodel.eval();","metadata":{"execution":{"iopub.status.busy":"2024-05-08T08:26:33.467457Z","iopub.execute_input":"2024-05-08T08:26:33.467834Z","iopub.status.idle":"2024-05-08T08:26:34.719785Z","shell.execute_reply.started":"2024-05-08T08:26:33.467803Z","shell.execute_reply":"2024-05-08T08:26:34.718616Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"Image.fromarray(images[25])","metadata":{"execution":{"iopub.status.busy":"2024-05-08T08:26:34.72146Z","iopub.execute_input":"2024-05-08T08:26:34.72251Z","iopub.status.idle":"2024-05-08T08:26:35.077305Z","shell.execute_reply.started":"2024-05-08T08:26:34.722455Z","shell.execute_reply":"2024-05-08T08:26:35.076178Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"ROTATING IMAGES )","metadata":{}},{"cell_type":"code","source":"imgIdx=[i for i in range(len(heritageDioscuriDataset))]\nprediction=[]\npredictedRotationDeg=[]\nfor i in imgIdx:\n    with torch.no_grad():\n        temp=model(heritageDioscuriDataset[i][None, ...]).numpy()\n        prediction.append(temp)\n\n#     print(f'Prediction: {prediction}')\n#     print(f\"Prediction argmax: {prediction[0].argmax()}\")\n\n    predictedRotationDeg.append(temp.argmax() * 90)\n#     print(f\"Predicted rotation is {predictedRotationDeg} degrees.\")\n\npredictionToRotationCodeMap = {\n    90: cv2.ROTATE_90_CLOCKWISE,\n    180: cv2.ROTATE_180,\n    270: cv2.ROTATE_90_COUNTERCLOCKWISE\n}\nrotatedImages=[]\nfor i in imgIdx:\n    if predictedRotationDeg[i]== 0:\n        rotatedImage = images[i]\n        rotatedImages.append(rotatedImage)\n    else:\n        rotatedImage = cv2.rotate(images[i], predictionToRotationCodeMap[predictedRotationDeg[i]])\n        rotatedImages.append(rotatedImage)\n        \n        \n#SAVING IMAGES\n%mkdir \"/kaggle/working/imgs\"\nfrom PIL import Image\nimport os\n\n# Define the directory where you want to save the images\nsave_dir = \"/kaggle/working/imgs\"\n\n# Iterate through the rotatedImages list\nfor idx, rotated_image in enumerate(rotatedImages):\n    # Convert the image array to a PIL image\n    pil_image = Image.fromarray(rotated_image)\n    \n    # Define the filename (you can adjust this according to your needs)\n    filename = f\"rotated_image_{idx}.jpg\"\n    \n    # Save the PIL image to the directory\n    pil_image.save(os.path.join(save_dir, filename))\n","metadata":{"execution":{"iopub.status.busy":"2024-05-08T08:26:35.078504Z","iopub.execute_input":"2024-05-08T08:26:35.079219Z","iopub.status.idle":"2024-05-08T08:26:48.924458Z","shell.execute_reply.started":"2024-05-08T08:26:35.079181Z","shell.execute_reply":"2024-05-08T08:26:48.92322Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nfrom tqdm import tqdm\nfrom time import time\nfrom fastprogress import progress_bar\nimport gc\nimport numpy as np\nimport h5py\nfrom IPython.display import clear_output\nfrom collections import defaultdict\nfrom copy import deepcopy\n\n# CV/ML\nimport cv2\nimport torch\nimport torch.nn.functional as F\nimport kornia as K\nimport kornia.feature as KF\nfrom PIL import Image\nimport timm\nfrom timm.data import resolve_data_config\nfrom timm.data.transforms_factory import create_transform\n\n# 3D reconstruction\nimport pycolmap","metadata":{"execution":{"iopub.status.busy":"2024-05-08T08:26:48.926394Z","iopub.execute_input":"2024-05-08T08:26:48.927757Z","iopub.status.idle":"2024-05-08T08:26:48.939331Z","shell.execute_reply.started":"2024-05-08T08:26:48.92766Z","shell.execute_reply":"2024-05-08T08:26:48.937927Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print('Kornia version', K.__version__)\nprint('Pycolmap version', pycolmap.__version__)\n\nLOCAL_FEATURE = 'LoFTR'\ndevice=torch.device('cuda')\n# Can be LoFTR, KeyNetAffNetHardNet, or DISK","metadata":{"execution":{"iopub.status.busy":"2024-05-08T08:26:48.945044Z","iopub.execute_input":"2024-05-08T08:26:48.945492Z","iopub.status.idle":"2024-05-08T08:26:48.960573Z","shell.execute_reply.started":"2024-05-08T08:26:48.945462Z","shell.execute_reply":"2024-05-08T08:26:48.959546Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def arr_to_str(a):\n    return ';'.join([str(x) for x in a.reshape(-1)])\n\n\ndef load_torch_image(fname, device=torch.device('cpu')):\n    img = K.image_to_tensor(cv2.imread(fname), False).float() / 255.\n    img = K.color.bgr_to_rgb(img.to(device))\n    return img","metadata":{"execution":{"iopub.status.busy":"2024-05-08T08:26:48.961981Z","iopub.execute_input":"2024-05-08T08:26:48.962611Z","iopub.status.idle":"2024-05-08T08:26:48.979776Z","shell.execute_reply.started":"2024-05-08T08:26:48.962576Z","shell.execute_reply":"2024-05-08T08:26:48.978481Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_global_desc(fnames, model,\n                    device =  torch.device('cpu')):\n    model = model.eval()\n    model= model.to(device)\n    config = resolve_data_config({}, model=model)\n    transform = create_transform(**config)\n    global_descs_convnext=[]\n    print(\"lolololololololol\")\n    for i, img_fname_full in tqdm(enumerate(fnames),total= len(fnames)):\n        key = os.path.splitext(os.path.basename(img_fname_full))[0]\n        img = Image.open(img_fname_full).convert('RGB')\n        print(img_fname_full)\n        timg = transform(img).unsqueeze(0).to(device)\n        with torch.no_grad():\n            desc = model.forward_features(timg.to(device)).mean(dim=(-1,2))#\n            #print (desc.shape)\n            desc = desc.view(1, -1)\n            desc_norm = F.normalize(desc, dim=1, p=2)\n        #print (desc_norm)\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_img_pairs_exhaustive(img_fnames):\n    index_pairs = []\n    for i in range(len(img_fnames)):\n        for j in range(i+1, len(img_fnames)):\n            index_pairs.append((i,j))\n    return index_pairs\n\n\ndef get_image_pairs_shortlist(fnames,\n                              sim_th = 0.6, # should be strict\n                              min_pairs = 20,\n                              exhaustive_if_less = 20,\n                              device=torch.device('cpu')):\n    num_imgs = len(fnames)\n#     print(num_imgs,fnames)\n\n    if num_imgs <= exhaustive_if_less:\n        return get_img_pairs_exhaustive(fnames)\n    print(\"creating eff model\")\n    model = timm.create_model('tf_efficientnet_b7',checkpoint=\"/kaggle/input/timm-pretrained-efficientnet/efficientnet/tf_efficientnet_b7_ra-6c08e654.pth\")\n    print(\"created eff model\")\n    model.eval()\n    print(\"model set to erval\")\n    descs = get_global_desc(fnames, model, device=device)\n    dm = torch.cdist(descs, descs, p=2).detach().cpu().numpy()\n    # removing half\n    mask = dm <= sim_th\n    total = 0\n    matching_list = []\n    ar = np.arange(num_imgs)\n    already_there_set = []\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] < 1000:\n                matching_list.append(tuple(sorted((st_idx, idx.item()))))\n                total+=1\n    matching_list = sorted(list(set(matching_list)))\n    return matching_list\n","metadata":{"execution":{"iopub.status.busy":"2024-05-08T08:26:48.981954Z","iopub.execute_input":"2024-05-08T08:26:48.982419Z","iopub.status.idle":"2024-05-08T08:26:49.005428Z","shell.execute_reply.started":"2024-05-08T08:26:48.982383Z","shell.execute_reply":"2024-05-08T08:26:49.003799Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import sys\nimport sqlite3\nimport numpy as np\n\n\nIS_PYTHON3 = sys.version_info[0] >= 3\n\nMAX_IMAGE_ID = 2**31 - 1\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\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\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\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\"\"\"add Codeadd Markdown\n[ ]:\n\ncreate_submission(out_results, image_dict)\nadd Codeadd Markdown\n[ ]:\n\n​\nadd Codeadd Markdown\n[ ]:\n\n​\nadd Codeadd Markdown\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\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\nCREATE_NAME_INDEX = \\\n    \"CREATE UNIQUE INDEX IF NOT EXISTS index_name ON images(name)\"\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\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 pair_id_to_image_ids(pair_id):\n    image_id2 = pair_id % MAX_IMAGE_IDadd Codeadd Markdown\n[ ]:\n\ncreate_submission(out_results, image_dict)\nadd Codeadd Markdown\n[ ]:\n\n​\nadd Codeadd Markdown\n[ ]:\n\n​\nadd Codeadd Markdown\n    image_id1 = (pair_id - image_id2) / MAX_IMAGE_ID\n    return image_id1, image_id2\n\n\ndef array_to_blob(array):\n    if IS_PYTHON3:\n        return array.tostring()\n    else:\n        return np.getbuffer(array)\n\n\ndef blob_to_array(blob, dtype, shape=(-1,)):\n    if IS_PYTHON3:\n        return np.fromstring(blob, dtype=dtype).reshape(*shape)\n    else:\n        return np.frombuffer(blob, dtype=dtype).reshape(*shape)\n","metadata":{"execution":{"iopub.status.busy":"2024-05-08T08:26:49.00724Z","iopub.execute_input":"2024-05-08T08:26:49.007664Z","iopub.status.idle":"2024-05-08T08:26:49.026454Z","shell.execute_reply.started":"2024-05-08T08:26:49.00763Z","shell.execute_reply":"2024-05-08T08:26:49.025309Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os, argparse, h5py, warnings\nimport numpy as np\nfrom tqdm import tqdm\nfrom PIL import Image, ExifTags\n\n\ndef get_focal(image_path, err_on_default=False):\n    image         = Image.open(image_path)\n    \n    max_size      = max(image.size)\n\n    exif = image.getexif()\n    focal = None\n    if exif is not None:\n        focal_35mm = Noneadd Codeadd Markdown\n[ ]:\n\ncreate_submission(out_results, image_dict)\nadd Codeadd Markdown\n[ ]:\n\n​\nadd Codeadd Markdown\n[ ]:\n\n​\nadd Codeadd Markdown\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\ndef create_camera(db, image_path, camera_model):\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, h5_path, image_path, img_ext, camera_model, single_camera = True):\n    keypoint_f = h5py.File(os.path.join(h5_path, '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        fname_with_ext = filename# + img_ext\n        path = os.path.join(image_path, fname_with_ext)\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(fname_with_ext, camera_id)\n        fname_to_id[filename] = image_id\n\n        db.add_keypoints(image_id, keypoints)\n\n    return fname_to_id\n\ndef add_matches(db, h5_path, fname_to_id):\n    match_file = h5py.File(os.path.join(h5_path, 'matches.h5'), 'r')\n    \n    added = set()\n    n_keys = len(match_file.keys())\n    n_total = (n_keys * (n_keys - 1)) // 2\n\n    with tqdm(total=n_total) as pbar:\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 = image_ids_to_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            \n                matches = group[key_2][()]\n                db.add_matches(id_1, id_2, matches)\n\n                added.add(pair_id)\n\n                pbar.update(1)\n","metadata":{"execution":{"iopub.status.busy":"2024-05-08T08:26:49.028407Z","iopub.execute_input":"2024-05-08T08:26:49.028974Z","iopub.status.idle":"2024-05-08T08:26:49.055263Z","shell.execute_reply.started":"2024-05-08T08:26:49.02894Z","shell.execute_reply":"2024-05-08T08:26:49.053876Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Making kornia local features loading w/o internet\nclass KeyNetAffNetHardNet(KF.LocalFeature):\n    \"\"\"Convenience module, which implements KeyNet detector + AffNet + HardNet descriptor.\n\n    .. image:: _static/img/keynet_affnet.jpg\n    \"\"\"\n\n    def __init__(\n        self,\n        num_features: int = 5000,\n        upright: bool = False,\n        device = torch.device('cpu'),\n        scale_laf: float = 1.0,\n    ):\n        ori_module = KF.PassLAF() if upright else KF.LAFOrienter(angle_detector=KF.OriNet(False)).eval()\n        if not upright:\n            weights = torch.load('/kaggle/input/kornia-local-feature-weights/OriNet.pth')['state_dict']\n            ori_module.angle_detector.load_state_dict(weights)\n        detector = KF.KeyNetDetector(\n            False, num_features=num_features, ori_module=ori_module, aff_module=KF.LAFAffNetShapeEstimator(False).eval()\n        ).to(device)\n        kn_weights = torch.load('/kaggle/input/kornia-local-feature-weights/keynet_pytorch.pth')['state_dict']\n        detector.model.load_state_dict(kn_weights)\n        affnet_weights = torch.load('/kaggle/input/kornia-local-feature-weights/AffNet.pth')['state_dict']\n        detector.aff.load_state_dict(affnet_weights)\n        \n        hardnet = KF.HardNet(False).eval()\n        hn_weights = torch.load('/kaggle/input/kornia-local-feature-weights/HardNetLib.pth')['state_dict']\n        hardnet.load_state_dict(hn_weights)\n        descriptor = KF.LAFDescriptor(hardnet, patch_size=32, grayscale_descriptor=True).to(device)\n        super().__init__(detector, descriptor, scale_laf)","metadata":{"execution":{"iopub.status.busy":"2024-05-08T08:26:49.057046Z","iopub.execute_input":"2024-05-08T08:26:49.058369Z","iopub.status.idle":"2024-05-08T08:26:49.076593Z","shell.execute_reply.started":"2024-05-08T08:26:49.058329Z","shell.execute_reply":"2024-05-08T08:26:49.075413Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def detect_features(img_fnames,\n                    num_feats = 2048,\n                    upright = False,\n                    device=torch.device('cpu'),\n                    feature_dir = '.featureout',\n                    resize_small_edge_to = 600):\n    if LOCAL_FEATURE == 'DISK':\n        # Load DISK from Kaggle models so it can run when the notebook is offline.\n        disk = KF.DISK().to(device)\n        pretrained_dict = torch.load('/kaggle/input/disk/pytorch/depth-supervision/1/loftr_outdoor.ckpt', map_location=device)\n        disk.load_state_dict(pretrained_dict['extractor'])\n        disk.eval()\n    if LOCAL_FEATURE == 'KeyNetAffNetHardNet':\n        print(\"mai jarau\")\n        feature = KeyNetAffNetHardNet(num_feats, upright, device).to(device).eval()\n    if not os.path.isdir(feature_dir):\n        os.makedirs(feature_dir)\n    with h5py.File(f'{feature_dir}/lafs.h5', mode='w') as f_laf, \\\n         h5py.File(f'{feature_dir}/keypoints.h5', mode='w') as f_kp, \\\n         h5py.File(f'{feature_dir}/descriptors.h5', mode='w') as f_desc:\n        for img_path in progress_bar(img_fnames):\n            img_fname = img_path.split('/')[-1]\n            key = img_fname\n            with torch.inference_mode():\n                timg = load_torch_image(img_path, device=device)\n                H, W = timg.shape[2:]\n                if resize_small_edge_to is None:\n                    timg_resized = timg\n                else:\n                    timg_resized = K.geometry.resize(timg, resize_small_edge_to, antialias=True)\n                    print(f'Resized {timg.shape} to {timg_resized.shape} (resize_small_edge_to={resize_small_edge_to})')\n                h, w = timg_resized.shape[2:]\n                if LOCAL_FEATURE == 'DISK':\n                    features = disk(timg_resized, num_feats, pad_if_not_divisible=True)[0]\n                    kps1, descs = features.keypoints, features.descriptors\n                    \n                    lafs = KF.laf_from_center_scale_ori(kps1[None], torch.ones(1, len(kps1), 1, 1, device=device))\n                if LOCAL_FEATURE == 'KeyNetAffNetHardNet':\n                    print(\"mai sachhi jarau\")\n                    lafs, resps, descs = feature(K.color.rgb_to_grayscale(timg_resized))\n                lafs[:,:,0,:] *= float(W) / float(w)\n                lafs[:,:,1,:] *= float(H) / float(h)\n                desc_dim = descs.shape[-1]\n                kpts = KF.get_laf_center(lafs).reshape(-1, 2).detach().cpu().numpy()\n                descs = descs.reshape(-1, desc_dim).detach().cpu().numpy()\n                f_laf[key] = lafs.detach().cpu().numpy()\n                f_kp[key] = kpts\n                f_desc[key] = descs\n    return\n\ndef get_unique_idxs(A, dim=0):\n    # https://stackoverflow.com/questions/72001505/how-to-get-unique-elements-and-their-firstly-appeared-indices-of-a-pytorch-tenso\n    unique, 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\n\ndef match_features(img_fnames,\n                   index_pairs,\n                   feature_dir = '.featureout',\n                   device=torch.device('cpu'),\n                   min_matches=15, \n                   force_mutual = True,\n                   matching_alg='smnn'\n                  ):\n    assert matching_alg in ['smnn', 'adalam']\n    with h5py.File(f'{feature_dir}/lafs.h5', mode='r') as f_laf, \\\n         h5py.File(f'{feature_dir}/descriptors.h5', mode='r') as f_desc, \\\n        h5py.File(f'{feature_dir}/matches.h5', mode='w') as f_match:\n\n        for pair_idx in progress_bar(index_pairs):\n                    idx1, idx2 = pair_idx\n                    fname1, fname2 = img_fnames[idx1], img_fnames[idx2]\n                    key1, key2 = fname1.split('/')[-1], fname2.split('/')[-1]\n                    lafs1 = torch.from_numpy(f_laf[key1][...]).to(device)\n                    lafs2 = torch.from_numpy(f_laf[key2][...]).to(device)\n                    desc1 = torch.from_numpy(f_desc[key1][...]).to(device)\n                    desc2 = torch.from_numpy(f_desc[key2][...]).to(device)\n                    if matching_alg == 'adalam':\n                        img1, img2 = cv2.imread(fname1), cv2.imread(fname2)\n                        hw1, hw2 = img1.shape[:2], img2.shape[:2]\n                        adalam_config = KF.adalam.get_adalam_default_config()\n                        #adalam_config['orientation_difference_threshold'] = None\n                        #adalam_config['scale_rate_threshold'] = None\n                        adalam_config['force_seed_mnn']= False\n                        adalam_config['search_expansion'] = 16\n                        adalam_config['ransac_iters'] = 128\n                        adalam_config['device'] = device\n                        dists, idxs = KF.match_adalam(desc1, desc2,\n                                                      lafs1, lafs2, # Adalam takes into account also geometric information\n                                                      hw1=hw1, hw2=hw2,\n                                                      config=adalam_config) # Adalam also benefits from knowing image size\n                    else:\n                        dists, idxs = KF.match_smnn(desc1, desc2, 0.98)\n                    if len(idxs)  == 0:\n                        continue\n                    # Force mutual nearest neighbors\n                    if force_mutual:\n                        first_indices = get_unique_idxs(idxs[:,1])\n                        idxs = idxs[first_indices]\n                        dists = dists[first_indices]\n                    n_matches = len(idxs)\n                    if False:\n                        print (f'{key1}-{key2}: {n_matches} matches')\n                    group  = f_match.require_group(key1)\n                    if n_matches >= min_matches:\n                         group.create_dataset(key2, data=idxs.detach().cpu().numpy().reshape(-1, 2))\n    return\n\ndef match_loftr(img_fnames,\n                   index_pairs,\n                   feature_dir = '.featureout_loftr',\n                   device=torch.device('cpu'),\n                   min_matches=15, resize_to_ = (640, 480)):\n    matcher = KF.LoFTR(pretrained=None)\n    matcher.load_state_dict(torch.load('/kaggle/input/loftr/pytorch/outdoor/1/loftr_outdoor.ckpt')['state_dict'])\n    matcher = matcher.to(device).eval()\n\n    # First we do pairwise matching, and then extract \"keypoints\" from loftr matches.\n    with h5py.File(f'{feature_dir}/matches_loftr.h5', mode='w') as f_match:\n        for pair_idx in progress_bar(index_pairs):\n            idx1, idx2 = pair_idx\n            fname1, fname2 = img_fnames[idx1], img_fnames[idx2]\n            key1, key2 = fname1.split('/')[-1], fname2.split('/')[-1]\n            # Load img1\n            timg1 = K.color.rgb_to_grayscale(load_torch_image(fname1, device=device))\n            H1, W1 = timg1.shape[2:]\n            if H1 < W1:\n                resize_to = resize_to_[1], resize_to_[0]\n            else:\n                resize_to = resize_to_\n            timg_resized1 = K.geometry.resize(timg1, resize_to, antialias=True)\n            h1, w1 = timg_resized1.shape[2:]\n\n            # Load img2\n            timg2 = K.color.rgb_to_grayscale(load_torch_image(fname2, device=device))\n            H2, W2 = timg2.shape[2:]\n            if H2 < W2:\n                resize_to2 = resize_to[1], resize_to[0]\n            else:\n                resize_to2 = resize_to_\n            timg_resized2 = K.geometry.resize(timg2, resize_to2, antialias=True)\n            h2, w2 = timg_resized2.shape[2:]\n            with torch.inference_mode():\n                input_dict = {\"image0\": timg_resized1,\"image1\": timg_resized2}\n                correspondences = matcher(input_dict)\n            mkpts0 = correspondences['keypoints0'].cpu().numpy()\n            mkpts1 = correspondences['keypoints1'].cpu().numpy()\n\n            mkpts0[:,0] *= float(W1) / float(w1)\n            mkpts0[:,1] *= float(H1) / float(h1)\n\n            mkpts1[:,0] *= float(W2) / float(w2)\n            mkpts1[:,1] *= float(H2) / float(h2)\n\n            n_matches = len(mkpts1)\n            group  = f_match.require_group(key1)\n            if n_matches >= min_matches:\n                 group.create_dataset(key2, data=np.concatenate([mkpts0, mkpts1], axis=1))\n\n    # Let's find unique loftr pixels and group them together.\n    kpts = defaultdict(list)\n    match_indexes = defaultdict(dict)\n    total_kpts=defaultdict(int)\n    with h5py.File(f'{feature_dir}/matches_loftr.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    unique_kpts = {}\n    unique_match_idxs = {}\n    out_match = defaultdict(dict)\n    for k in kpts.keys():\n        uniq_kps, uniq_reverse_idxs = torch.unique(torch.from_numpy(kpts[k]),dim=0, return_inverse=True)\n        unique_match_idxs[k] = uniq_reverse_idxs\n        unique_kpts[k] = uniq_kps.numpy()\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]],\n                                    unique_kpts[k2][  m2[:,1]],\n                                   ],\n                                   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    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    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    return\n\ndef import_into_colmap(img_dir,\n                       feature_dir ='.featureout',\n                       database_path = 'colmap.db',\n                       img_ext='.jpg'):\n    db = COLMAPDatabase.connect(database_path)\n    db.create_tables()\n    single_camera = False\n    fname_to_id = add_keypoints(db, feature_dir, img_dir, img_ext, 'simple-radial', single_camera)\n    add_matches(\n        db,\n        feature_dir,\n        fname_to_id,\n    )\n\n    db.commit()\n    return\n","metadata":{"execution":{"iopub.status.busy":"2024-05-08T08:31:14.553241Z","iopub.execute_input":"2024-05-08T08:31:14.553667Z","iopub.status.idle":"2024-05-08T08:31:14.618792Z","shell.execute_reply.started":"2024-05-08T08:31:14.553636Z","shell.execute_reply":"2024-05-08T08:31:14.617387Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# src = '/kaggle/input/image-matching-challenge-2024'\n# Get data from csv.\n\n# data_dict = {}\n# with open(f'{src}/sample_submission.csv', 'r') as f:\n#     for i, l in enumerate(f):\n#         # Skip header.\n#         if l and i > 0:\n#             image, dataset, scene, _, _ = l.strip().split(',')\n#             if dataset not in data_dict:\n#                 data_dict[dataset] = {}\n#             if scene not in data_dict[dataset]:\n#                 data_dict[dataset][scene] = []\n#             data_dict[dataset][scene].append(image)\n# for dataset in data_dict:\n#     for scene in data_dict[dataset]:\n#         print(f'{dataset} / {scene} -> {len(data_dict[dataset][scene])} images')","metadata":{"execution":{"iopub.status.busy":"2024-05-08T08:26:49.144607Z","iopub.execute_input":"2024-05-08T08:26:49.145158Z","iopub.status.idle":"2024-05-08T08:26:49.164141Z","shell.execute_reply.started":"2024-05-08T08:26:49.145112Z","shell.execute_reply":"2024-05-08T08:26:49.162702Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\n\n# Directory containing your images\nimgs_dir = \"/kaggle/working/imgs\"\n\n# List all image files in the directory\nimage_files = os.listdir(imgs_dir)\n\n# Create a dictionary structure\nimage_dict = {'discouri': {'discouri': [os.path.join('/kaggle/working/imgs', img_file) for img_file in image_files]}}\n\n# Display the dictionary\nprint(image_dict)\n","metadata":{"execution":{"iopub.status.busy":"2024-05-08T08:26:49.165735Z","iopub.execute_input":"2024-05-08T08:26:49.166244Z","iopub.status.idle":"2024-05-08T08:26:49.17927Z","shell.execute_reply.started":"2024-05-08T08:26:49.166201Z","shell.execute_reply":"2024-05-08T08:26:49.177852Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"out_results = {}\ntimings = {\"shortlisting\":[],\n           \"feature_detection\": [],\n           \"feature_matching\":[],\n           \"RANSAC\": [],\n           \"Reconstruction\": []}\n# Function to create a submission file.\ndef create_submission(out_results, data_dict):\n    with open(f'submission.csv', 'w') as f:\n        f.write('image_path,dataset,scene,rotation_matrix,translation_vector\\n')\n        for dataset in data_dict:\n            if dataset in out_results:\n                res = out_results[dataset]\n            else:\n                res = {}\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                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                    f.write(f'{image},{dataset},{scene},{arr_to_str(R)},{arr_to_str(T)}\\n')\ngc.collect()\ndatasets = []\nfor dataset in image_dict:\n    datasets.append(dataset)\n\nfor dataset in datasets:\n    print(dataset,\"nooo\")\n    if dataset not in out_results:\n        out_results[dataset] = {}\n    for scene in image_dict[dataset]:\n        print(scene,\"hooooo\")\n        # Fail gently if the notebook has not been submitted and the test data is not populated.\n        # You may want to run this on the training data in that case?\n        img_dir = '/kaggle/working/imgs'\n        if not os.path.exists(img_dir):\n            print(img_dir,\"oh shit\")\n            continue\n        # Wrap the meaty part in a try-except block.\n        try:\n            out_results[dataset][scene] = {}\n            img_fnames = [f'{x}' for x in image_dict[dataset][scene]]\n#             print(img_fnames)\n            print (f\"Got {len(img_fnames)} images\")\n            feature_dir = f'/kaggle/working/featureout/{dataset}_{scene}'\n            print(feature_dir)\n            if not os.path.isdir(feature_dir):\n                print(\"making\")\n                os.makedirs(feature_dir, exist_ok=True)\n            t=time()\n            print(\"hiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiii\")\n            index_pairs = get_image_pairs_shortlist(img_fnames,\n                                  sim_th = 0.5, # should be strict\n                                  min_pairs = 20, # we select at least min_pairs PER IMAGE with biggest similarity\n                                  exhaustive_if_less = 20,\n                                  device=device)\n            print(\"hiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiii\")\n            \n            t=time() -t \n            timings['shortlisting'].append(t)\n            print (f'{len(index_pairs)}, pairs to match, {t:.4f} sec')\n            gc.collect()\n            t=time()\n            if LOCAL_FEATURE != 'LoFTR':\n                detect_features(img_fnames, \n                                2048,\n                                feature_dir=feature_dir,\n                                upright=True,\n                                device=device,\n                                resize_small_edge_to=600\n                               )\n                gc.collect()\n                t=time() -t \n                timings['feature_detection'].append(t)\n                print(f'Features detected in  {t:.4f} sec')\n                t=time()\n                match_features(img_fnames, index_pairs, feature_dir=feature_dir,device=device)\n            else:\n                match_loftr(img_fnames, index_pairs, feature_dir=feature_dir, device=device, resize_to_=(632, 832))\n            t=time() -t \n            timings['feature_matching'].append(t)\n            print(f'Features matched in  {t:.4f} sec')\n            database_path = f'{feature_dir}/colmap.db'\n            if os.path.isfile(database_path):\n                os.remove(database_path)\n            gc.collect()\n            import_into_colmap(img_dir, feature_dir=feature_dir,database_path=database_path)\n            output_path = f'{feature_dir}/colmap_rec_{LOCAL_FEATURE}'\n\n            t=time()\n            pycolmap.match_exhaustive(database_path)\n            t=time() - t \n            timings['RANSAC'].append(t)\n            print(f'RANSAC in  {t:.4f} sec')\n\n            t=time()\n            # By default colmap does not generate a reconstruction if less than 10 images are registered. Lower it to 3.\n            mapper_options = pycolmap.IncrementalMapperOptions()\n            mapper_options.min_model_size = 3\n            os.makedirs(output_path, exist_ok=True)\n            maps = pycolmap.incremental_mapping(database_path=database_path, image_path=img_dir, output_path=output_path, options=mapper_options)\n            print(maps)\n            #clear_output(wait=False)\n            t=time() - t\n            timings['Reconstruction'].append(t)\n            print(f'Reconstruction done in  {t:.4f} sec')\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 k, im in maps[best_idx].images.items():\n                    key1 = f'{dataset}/{scene}/images/{im.name}'\n                    out_results[dataset][scene][key1] = {}\n                    out_results[dataset][scene][key1][\"R\"] = deepcopy(im.rotmat())\n                    out_results[dataset][scene][key1][\"t\"] = deepcopy(np.array(im.tvec))\n            print(f'Registered: {dataset} / {scene} -> {len(out_results[dataset][scene])} images')\n            print(f'Total: {dataset} / {scene} -> {len(image_dict[dataset][scene])} images')\n            create_submission(out_results, image_dict)\n            gc.collect()\n        except:\n            pass\n        ","metadata":{"execution":{"iopub.status.busy":"2024-05-08T08:31:29.288181Z","iopub.execute_input":"2024-05-08T08:31:29.288611Z","iopub.status.idle":"2024-05-08T08:31:29.572732Z","shell.execute_reply.started":"2024-05-08T08:31:29.288579Z","shell.execute_reply":"2024-05-08T08:31:29.57157Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"create_submission(out_results, image_dict)","metadata":{"execution":{"iopub.status.busy":"2024-05-08T08:26:49.5065Z","iopub.execute_input":"2024-05-08T08:26:49.506908Z","iopub.status.idle":"2024-05-08T08:26:49.527032Z","shell.execute_reply.started":"2024-05-08T08:26:49.506875Z","shell.execute_reply":"2024-05-08T08:26:49.525429Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}