{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":71885,"databundleVersionId":8143495,"sourceType":"competition"},{"sourceId":7884485,"sourceType":"datasetVersion","datasetId":4628051},{"sourceId":7884725,"sourceType":"datasetVersion","datasetId":4628331},{"sourceId":172469456,"sourceType":"kernelVersion"},{"sourceId":173217852,"sourceType":"kernelVersion"},{"sourceId":4534,"sourceType":"modelInstanceVersion","modelInstanceId":3326},{"sourceId":17191,"sourceType":"modelInstanceVersion","modelInstanceId":14317},{"sourceId":17199,"sourceType":"modelInstanceVersion","modelInstanceId":14324},{"sourceId":17555,"sourceType":"modelInstanceVersion","modelInstanceId":14611}],"dockerImageVersionId":30683,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## IMC24 Starter for Image Matching Challenge 2024 Hexathlon.\n\nIMC24 Starter came to existance thanks to [imc-understanding-the-baseline][1]. The metric score for mean average accuracy came from [IMC2024-3D-metric-evaluation-example][2].<br>\nTo further modify the code for submission and scoring, the utility script can be accessed [here][3].\n\n[1]: https://www.kaggle.com/code/asarvazyan/imc-understanding-the-baseline\n[2]: https://www.kaggle.com/code/fabiobellavia/imc2024-3d-metric-evaluation-example\n[3]: https://www.kaggle.com/code/nartaa/imc24","metadata":{}},{"cell_type":"markdown","source":"# SETUP","metadata":{}},{"cell_type":"code","source":"from imc24 import *","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2024-05-12T06:25:33.775430Z","iopub.execute_input":"2024-05-12T06:25:33.775895Z","iopub.status.idle":"2024-05-12T06:26:39.565197Z","shell.execute_reply.started":"2024-05-12T06:25:33.775861Z","shell.execute_reply":"2024-05-12T06:26:39.564190Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# SIMILLIAR PAIRS","metadata":{}},{"cell_type":"code","source":"def get_pairs(images_list,device=DEVICE):\n    if EXHAUSTIVE:\n        return list(combinations(range(len(images_list)), 2)) \n    \n    processor = AutoImageProcessor.from_pretrained('/kaggle/input/dinov2/pytorch/base/1/')\n    model = AutoModel.from_pretrained('/kaggle/input/dinov2/pytorch/base/1/').eval().to(DEVICE)\n    embeddings = []\n    \n    for img_path in images_list:\n        image = K.io.load_image(img_path, K.io.ImageLoadType.RGB32, device=DEVICE)[None, ...]\n        with torch.inference_mode():\n            inputs = processor(images=image, return_tensors=\"pt\", do_rescale=False ,do_resize=True, \n                               do_center_crop=True, size=224).to(DEVICE)\n            outputs = model(**inputs)\n            embedding = F.normalize(outputs.last_hidden_state.max(dim=1)[0])\n        embeddings.append(embedding)\n        \n    embeddings = torch.cat(embeddings, dim=0)\n    distances = torch.cdist(embeddings,embeddings).cpu()\n    distances_ = (distances <= DISTANCES_THRESHOLD).numpy()\n    np.fill_diagonal(distances_,False)\n    z = distances_.sum(axis=1)\n    idxs0 = np.where(z == 0)[0]\n    for idx0 in idxs0:\n        t = np.argsort(distances[idx0])[1:MIN_PAIRS]\n        distances_[idx0,t] = True\n        \n    s = np.where(distances >= TOLERANCE)\n    distances_[s] = False\n    \n    idxs = []\n    for i in range(len(images_list)):\n        for j in range(len(images_list)):\n            if distances_[i][j]:\n                idxs += [(i,j)] if i<j else [(j,i)]\n    \n    idxs = list(set(idxs))\n    return idxs","metadata":{"execution":{"iopub.status.busy":"2024-05-12T06:26:39.567340Z","iopub.execute_input":"2024-05-12T06:26:39.567822Z","iopub.status.idle":"2024-05-12T06:26:39.579723Z","shell.execute_reply.started":"2024-05-12T06:26:39.567786Z","shell.execute_reply":"2024-05-12T06:26:39.578700Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# KEYPOINTS EXTRACTOR AND MATCHER","metadata":{}},{"cell_type":"code","source":"import torch\nfrom torchvision import io\nimport time\n\ndef rotate_image_and_keypoints(image_path, keypoints, angle):\n    with torch.no_grad():\n        image = K.io.load_image(image_path, K.io.ImageLoadType.RGB32, device=DEVICE)\n\n        # Get image center\n        height, width = image.shape[-2], image.shape[-1]\n\n        # Determine the shape after rotation\n        if angle == 90 or angle == 270:\n            rotated_height, rotated_width = width, height\n        else:\n            rotated_height, rotated_width = height, width\n        # Rotate keypoints\n        if angle == 90:\n            rotated_keypoints = torch.stack([height - keypoints[:, 1], keypoints[:, 0]], dim=1)\n        elif angle == 180:\n            rotated_keypoints = torch.stack([width - keypoints[:, 0], height - keypoints[:, 1]], dim=1)\n        elif angle == 270:\n            rotated_keypoints = torch.stack([keypoints[:, 1], width - keypoints[:, 0]], dim=1)\n        elif angle == 0:\n            rotated_keypoints = keypoints.clone()\n        #print(time.time())\n        return (rotated_height, rotated_width), rotated_keypoints.to(DEVICE)\n\n\ndef rotate_image_and_keypoints_back(rotated_image_shape, rotated_keypoints, angle):\n    with torch.no_grad():\n        rotated_height, rotated_width = rotated_image_shape\n\n        # Rotate keypoints back\n        if angle == 90:\n            keypoints_back = torch.stack([rotated_keypoints[:, 1], rotated_height - rotated_keypoints[:, 0] - 1], dim=1)\n        elif angle == 180:\n            keypoints_back = torch.stack([rotated_width - rotated_keypoints[:, 0] - 1, rotated_height - rotated_keypoints[:, 1] - 1], dim=1)\n        elif angle == 270:\n            keypoints_back = torch.stack([rotated_width - rotated_keypoints[:, 1] - 1, rotated_keypoints[:, 0]], dim=1)\n        elif angle == 0:\n            keypoints_back = rotated_keypoints.clone()\n    #     print(\"HI\", HI)\n    #     print(time.time())\n        return keypoints_back.to(DEVICE)\n","metadata":{"execution":{"iopub.status.busy":"2024-05-12T06:29:25.585626Z","iopub.execute_input":"2024-05-12T06:29:25.586023Z","iopub.status.idle":"2024-05-12T06:29:25.598626Z","shell.execute_reply.started":"2024-05-12T06:29:25.585992Z","shell.execute_reply":"2024-05-12T06:29:25.597643Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def keypoints_matches(images_list,pairs):\n    extractor = ALIKED(max_num_keypoints=MAX_NUM_KEYPOINTS,detection_threshold=DETECTION_THRESHOLD,resize=RESIZE_TO).eval().to(DEVICE)\n    matcher = KF.LightGlueMatcher(\"aliked\", {'width_confidence':-1, 'depth_confidence':-1, 'mp':True if 'cuda' in str(DEVICE) else False}).eval().to(DEVICE)\n    rotation = create_model(\"swsl_resnext50_32x4d\").eval().to(DEVICE)\n    \n    with h5py.File(\"keypoints.h5\", mode=\"w\") as f_kp, h5py.File(\"descriptors.h5\", mode=\"w\") as f_desc:  \n        for image_path in images_list:\n            with torch.inference_mode():\n                image = load_image(image_path).to(DEVICE)\n                feats = extractor.extract(image)\n                f_kp[image_path.name] = feats[\"keypoints\"].squeeze().cpu().numpy()\n                f_desc[image_path.name] = feats[\"descriptors\"].squeeze().detach().cpu().numpy()\n                \n    with h5py.File(\"keypoints.h5\", mode=\"r\") as f_kp, h5py.File(\"descriptors.h5\", mode=\"r\") as f_desc, \\\n         h5py.File(\"matches.h5\", mode=\"w\") as f_matches:  \n        print(\"length of pairs : \" , len(pairs))\n        for number,pair in enumerate(pairs):\n            key1, key2 = images_list[pair[0]].name, images_list[pair[1]].name\n            kp1 = torch.from_numpy(f_kp[key1][...]).to(DEVICE)\n            kp2 = torch.from_numpy(f_kp[key2][...]).to(DEVICE)\n            desc1 = torch.from_numpy(f_desc[key1][...]).to(DEVICE)\n            desc2 = torch.from_numpy(f_desc[key2][...]).to(DEVICE)\n            \n            best_k = 0\n            best_n_matches = 0\n            min_n_match = 9999\n            best_keypoints = kp1\n            for k_rotation in range(4):\n\n                rotated_image_shape, rotated_keypoints = rotate_image_and_keypoints(images_list[pair[0]], kp1, k_rotation*90)\n                if k_rotation == 0:\n                    best_rotated_image_shape = rotated_image_shape\n\n                with torch.inference_mode():\n                    _, idxs = matcher(desc1, desc2, \n                                      KF.laf_from_center_scale_ori(rotated_keypoints[None]), \n                                      KF.laf_from_center_scale_ori(kp2[None]))\n                n_matches = len(idxs)\n                if n_matches <= SOFT_MIN_MATCH:\n                    print(n_matches,\"too low\")\n                    break\n                #print(\"n_matches:\", n_matches)\n                min_n_match = min(min_n_match,n_matches)\n                if(n_matches > best_n_matches):\n                    best_k = k_rotation\n                    best_n_matches = n_matches\n                    best_keypoints = rotated_keypoints\n                    best_rotated_image_shape = rotated_image_shape\n                \n            print(number,\".best k rotation:\", best_k*90, \". with match count = \",best_n_matches , \"MIN:\",min_n_match)\n            keypoints_back = rotate_image_and_keypoints_back(best_rotated_image_shape, best_keypoints, best_k*90)\n            if best_n_matches: group = f_matches.require_group(key1)\n            if best_n_matches >= MIN_MATCHES: group.create_dataset(key2, data=idxs.detach().cpu().numpy())\n\n            ##################################################################\n\n                ","metadata":{"execution":{"iopub.status.busy":"2024-05-12T06:29:26.581203Z","iopub.execute_input":"2024-05-12T06:29:26.581549Z","iopub.status.idle":"2024-05-12T06:29:26.597778Z","shell.execute_reply.started":"2024-05-12T06:29:26.581525Z","shell.execute_reply":"2024-05-12T06:29:26.596729Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# RANSAC AND SPARSE RECONSTRUCTION","metadata":{}},{"cell_type":"code","source":"def ransac_and_sparse_reconstruction(images_path):\n    now = datetime.datetime.now()\n    time_str = now.strftime(\"%Y-%m-%d_%H-%M-%S\")\n    db_name = f'colmap_{time_str}.db'\n    db = COLMAPDatabase.connect(db_name)\n    db.create_tables()\n    fname_to_id = add_keypoints(db, '/kaggle/working/', images_path, '', 'simple-pinhole', False)\n    add_matches(db, '/kaggle/working/',fname_to_id)\n    db.commit()\n    \n    pycolmap.match_exhaustive(db_name, sift_options={'num_threads':1})\n    maps = pycolmap.incremental_mapping(\n        database_path=db_name, \n        image_path=images_path,\n        output_path='/kaggle/working/', \n        options=pycolmap.IncrementalPipelineOptions({'min_model_size':MIN_MODEL_SIZE, 'max_num_models':MAX_NUM_MODELS, 'num_threads':1})\n    )\n    return maps","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2024-05-12T06:29:27.174246Z","iopub.execute_input":"2024-05-12T06:29:27.174613Z","iopub.status.idle":"2024-05-12T06:29:27.182191Z","shell.execute_reply.started":"2024-05-12T06:29:27.174582Z","shell.execute_reply":"2024-05-12T06:29:27.181165Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# HYPERPARAMETER TUNING","metadata":{}},{"cell_type":"code","source":"# SIMILLIAR PAIRS\nEXHAUSTIVE = True\nMIN_PAIRS = 50\nDISTANCES_THRESHOLD = 0.3\nTOLERANCE = 500\n\n# KEYPOINTS EXTRACTOR AND MATCHER\nMAX_NUM_KEYPOINTS = 4096\nRESIZE_TO = 1280\nDETECTION_THRESHOLD = 0.002\nMIN_MATCHES = 100\nSOFT_MIN_MATCH = 5\n\n# RANSAC AND SPARSE RECONSTRUCTION\nMIN_MODEL_SIZE = 3\nMAX_NUM_MODELS = 3\n\n# CROSS VALIDATION\nN_SAMPLES = 50\n\nSUBMISSION = True","metadata":{"execution":{"iopub.status.busy":"2024-05-12T06:29:27.232595Z","iopub.execute_input":"2024-05-12T06:29:27.232913Z","iopub.status.idle":"2024-05-12T06:29:27.238321Z","shell.execute_reply.started":"2024-05-12T06:29:27.232888Z","shell.execute_reply":"2024-05-12T06:29:27.237458Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## CROSS VALIDATION","metadata":{}},{"cell_type":"raw","source":"if not SUBMISSION:\n    def image_path(row):\n        row['image_path'] = 'train/' + row['dataset'] + '/images/' + row['image_name']\n        return row\n\n    train_df = pd.read_csv(f'{IMC_PATH}/train/train_labels.csv')\n    train_df = train_df.apply(image_path,axis=1).drop_duplicates(subset=['image_path'])\n    G = train_df.groupby(['dataset','scene'])['image_path']\n    image_paths = []\n    \n    for g in G:\n        n = N_SAMPLES\n        n = n if n < len(g[1]) else len(g[1])\n        g = g[0],g[1].sample(n,random_state=42).reset_index(drop=True)\n        for image_path in g[1]:\n            image_paths.append(image_path)\n        \n    gt_df = train_df[train_df.image_path.isin(image_paths)].reset_index(drop=True)\n    pred_df = gt_df[['image_path','dataset','scene','rotation_matrix','translation_vector']]\n    pred_df.to_csv('pred_df.csv',index=False)\n    run('pred_df.csv', get_pairs, keypoints_matches, ransac_and_sparse_reconstruction, submit=False)\n    pred_df = pd.read_csv('submission.csv')\n    mAA = round(score(gt_df, pred_df),4)\n    print('*** Total mean Average Accuracy ***')\n    print(f\"mAA: {mAA}\")","metadata":{"execution":{"iopub.status.busy":"2024-05-11T13:49:06.311484Z","iopub.execute_input":"2024-05-11T13:49:06.311867Z","iopub.status.idle":"2024-05-11T13:49:06.322278Z","shell.execute_reply.started":"2024-05-11T13:49:06.311824Z","shell.execute_reply":"2024-05-11T13:49:06.321172Z"}}},{"cell_type":"markdown","source":"# SUBMISSION","metadata":{}},{"cell_type":"code","source":"if SUBMISSION:\n    data_path = IMC_PATH + \"/sample_submission.csv\"\n    run(data_path, get_pairs, keypoints_matches, ransac_and_sparse_reconstruction)","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2024-05-12T06:29:28.118395Z","iopub.execute_input":"2024-05-12T06:29:28.118802Z","iopub.status.idle":"2024-05-12T06:29:53.323269Z","shell.execute_reply.started":"2024-05-12T06:29:28.118769Z","shell.execute_reply":"2024-05-12T06:29:53.321694Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}