{"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":173368045,"sourceType":"kernelVersion"},{"sourceId":4534,"sourceType":"modelInstanceVersion","modelInstanceId":3326},{"sourceId":17191,"sourceType":"modelInstanceVersion","modelInstanceId":14317},{"sourceId":17555,"sourceType":"modelInstanceVersion","modelInstanceId":14611}],"dockerImageVersionId":30683,"isInternetEnabled":true,"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-04-22T17:24:35.125878Z","iopub.execute_input":"2024-04-22T17:24:35.126263Z","iopub.status.idle":"2024-04-22T17:24:59.939167Z","shell.execute_reply.started":"2024-04-22T17:24:35.126235Z","shell.execute_reply":"2024-04-22T17:24:59.937861Z"},"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-04-22T17:25:54.483471Z","iopub.execute_input":"2024-04-22T17:25:54.483839Z","iopub.status.idle":"2024-04-22T17:25:54.496174Z","shell.execute_reply.started":"2024-04-22T17:25:54.483808Z","shell.execute_reply":"2024-04-22T17:25:54.495264Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# KEYPOINTS EXTRACTOR AND MATCHER","metadata":{}},{"cell_type":"code","source":"def rotate_image_v2(image,rotation):\n    with torch.no_grad():\n        pred = rotation(image[None,...]).argmax()\n    angle = (pred*90).cpu().numpy()\n    \n    if angle == 0:\n        image = image\n    elif angle == 90:\n        image = image.rot90(dims=[1,2])\n        image = image.rot90(dims=[1,2])\n        image = image.rot90(dims=[1,2])\n    elif angle == 180:\n        image = image.rot90(dims=[1,2])\n        image = image.rot90(dims=[1,2])\n    elif angle == 270:\n        image = image.rot90(dims=[1,2])\n    return image, angle","metadata":{"execution":{"iopub.status.busy":"2024-04-22T17:25:56.956028Z","iopub.execute_input":"2024-04-22T17:25:56.956635Z","iopub.status.idle":"2024-04-22T17:25:56.963421Z","shell.execute_reply.started":"2024-04-22T17:25:56.956604Z","shell.execute_reply":"2024-04-22T17:25:56.962491Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def inv_rot_matrix(angle):\n    rot_matrix =  np.vstack((\n        [[np.cos(angle*np.pi/180), -np.sin(angle*np.pi/180), 0]],\n        [[np.sin(angle*np.pi/180), np.cos(angle*np.pi/180), 0]],\n        [[0, 0, 1]]\n    ));\n    return rot_matrix","metadata":{"execution":{"iopub.status.busy":"2024-04-22T17:25:59.018186Z","iopub.execute_input":"2024-04-22T17:25:59.018552Z","iopub.status.idle":"2024-04-22T17:25:59.023784Z","shell.execute_reply.started":"2024-04-22T17:25:59.018522Z","shell.execute_reply":"2024-04-22T17:25:59.022959Z"},"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    #extractor1280 = ALIKED(max_num_keypoints=MAX_NUM_KEYPOINTS,detection_threshold=DETECTION_THRESHOLD,resize=1280).eval().to(DEVICE)\n    #extractor1024 = ALIKED(max_num_keypoints=MAX_NUM_KEYPOINTS,detection_threshold=DETECTION_THRESHOLD,resize=1024).eval().to(DEVICE)\n    matcher = LightGlue(features=\"aliked\",depth_confidence=-1, width_confidence=-1).eval().to(DEVICE)\n    rotation = create_model(\"swsl_resnext50_32x4d\").eval().to(DEVICE)\n    \n    rotation_dictionary = {}\n    \n    with h5py.File(\"keypoints.h5\", mode=\"w\") as f_keypoints, \\\n         h5py.File(\"descriptors.h5\", mode=\"w\") as f_descriptors, \\\n         h5py.File(\"matches.h5\", mode=\"w\") as f_matches:\n        \n        exist = np.zeros(len(images_list), dtype=bool)\n        for pair in pairs:\n            key0, key1 = images_list[pair[0]].name, images_list[pair[1]].name\n            image0 = load_image(images_list[pair[0]]).to(DEVICE)\n            image1 = load_image(images_list[pair[1]]).to(DEVICE)\n            \n            # pre-rotate the image\n            if APPLY_ROTATION:\n                # if images_list[pair[0]].parts[-3] in ROTATE_DATASET:\n                [image0, angle0] = rotate_image_v2(image0,rotation)\n                [image1, angle1] = rotate_image_v2(image1,rotation)\n            else:\n                angle0 = 0\n                angle1 = 0\n            rotation_dictionary[key0] = angle0\n            rotation_dictionary[key1] = angle1\n            \n            feats0, feats1, matches01 = overlap_detection(extractor, matcher, image0, image1, MIN_MATCHES_OVERLAP)\n            #feats0, feats1, matches01 = overlap_detection(extractor1280, matcher, image0, image1, MIN_MATCHES_OVERLAP)\n            #feats2, feats3, matches02 = overlap_detection(extractor1024, matcher, image0, image1, MIN_MATCHES_OVERLAP)\n            if not exist[pair[0]]:\n                f_keypoints[key0] = feats0[\"keypoints\"].numpy()\n                f_descriptors[key0] = feats0[\"descriptors\"].numpy()\n                #f_keypoints[key0] = np.concatenate((feats0[\"keypoints\"].numpy(), feats2[\"keypoints\"].numpy()))\n                #f_descriptors[key0] = np.concatenate((feats0[\"descriptors\"].numpy(), feats2[\"descriptors\"].numpy()))\n            if not exist[pair[1]]:\n                f_keypoints[key1] = feats1[\"keypoints\"].numpy()\n                f_descriptors[key1] = feats1[\"descriptors\"].numpy()\n                #f_keypoints[key1] = np.concatenate((feats1[\"keypoints\"].numpy(), feats3[\"keypoints\"].numpy()))\n                #f_descriptors[key1] = np.concatenate((feats1[\"descriptors\"].numpy(), feats3[\"descriptors\"].numpy()))\n            exist[pair[0]],exist[pair[1]] = True,True\n            #if len(np.concatenate((matches01['matches'],matches02['matches']))) >= MIN_MATCHES:\n                #group  = f_matches.require_group(key0)\n                #group.create_dataset(key1, data=np.concatenate((matches01['matches'],matches02['matches'])))\n            if len(matches01['matches']) >= MIN_MATCHES:\n                group  = f_matches.require_group(key0)\n                group.create_dataset(key1, data=matches01['matches'])\n    return rotation_dictionary","metadata":{"execution":{"iopub.status.busy":"2024-04-22T17:26:00.501011Z","iopub.execute_input":"2024-04-22T17:26:00.501378Z","iopub.status.idle":"2024-04-22T17:26:00.515283Z","shell.execute_reply.started":"2024-04-22T17:26:00.501347Z","shell.execute_reply":"2024-04-22T17:26:00.514344Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"","metadata":{}},{"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-04-22T17:26:02.718381Z","iopub.execute_input":"2024-04-22T17:26:02.719063Z","iopub.status.idle":"2024-04-22T17:26:02.725869Z","shell.execute_reply.started":"2024-04-22T17:26:02.719030Z","shell.execute_reply":"2024-04-22T17:26:02.724912Z"},"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.005\nMIN_MATCHES = 100\nROTATE_DATASET = [''] # i.e ['church', 'dioscuri', 'lizard', 'multi-temporal-temple-baalshamin']\nMIN_MATCHES_OVERLAP = 5000\n\n# RANSAC AND SPARSE RECONSTRUCTION\nMIN_MODEL_SIZE = 5\nMAX_NUM_MODELS = 3\n\n# CROSS VALIDATION\nN_SAMPLES = 15\n\nSUBMISSION = True\nAPPLY_ROTATION = True","metadata":{"execution":{"iopub.status.busy":"2024-04-22T17:26:04.494848Z","iopub.execute_input":"2024-04-22T17:26:04.495536Z","iopub.status.idle":"2024-04-22T17:26:04.501023Z","shell.execute_reply.started":"2024-04-22T17:26:04.495503Z","shell.execute_reply":"2024-04-22T17:26:04.500069Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Create plots for report","metadata":{}},{"cell_type":"code","source":"plot = False\nif plot:\n    from lightglue import LightGlue, SuperPoint, DISK, SIFT, ALIKED\n    from lightglue.utils import load_image, rbd\n    from lightglue import viz2d\n    import torch\n\n    device = K.utils.get_cuda_device_if_available(0)\n\n    extractor = ALIKED(max_num_keypoints=4096).eval().to(device)  # load the extractor\n    matcher = LightGlue(features=\"aliked\").eval().to(device)","metadata":{"execution":{"iopub.status.busy":"2024-04-22T17:26:12.055035Z","iopub.execute_input":"2024-04-22T17:26:12.055396Z","iopub.status.idle":"2024-04-22T17:26:12.061275Z","shell.execute_reply.started":"2024-04-22T17:26:12.055367Z","shell.execute_reply":"2024-04-22T17:26:12.060288Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot = False\nif plot:\n    path = \"/kaggle/input/image-matching-challenge-2024/train/\"\n    # church\n    #image0 = load_image(path + \"church/images/00001.png\")\n    #image1 = load_image(path + \"church/images/00010.png\")\n\n    # dioscuri rotation\n    image1 = load_image(path + \"dioscuri/images/3dom_fbk_img_1512.png\")\n    #image0 = load_image(path + \"dioscuri/images/3dom_fbk_img_1549.png\")\n    #image1 = load_image(path + \"dioscuri/images/3dom_fbk_img_1538.png\")\n    image0 = load_image(path + \"dioscuri/images/img_0152.png\") #issue\n\n    rotation = create_model(\"swsl_resnext50_32x4d\").eval().to(DEVICE)\n    [image0, angle0] = rotate_image_v2(image0,rotation)\n    print(angle0)\n    [image1, angle1] = rotate_image_v2(image1,rotation)\n    print(angle1)\n\n    #image0 = load_image(path + \"lizard/images/00001.png\")\n    #image1 = load_image(path + \"lizard/images/00002.png\")\n\n    #image0 = load_image(path + \"lizard/images/00057.png\")\n    #image1 = load_image(path + \"lizard/images/00059.png\")\n\n    # lizard night vs day\n    #image0 = load_image(path + \"lizard/images/00714.png\")\n    #image1 = load_image(path + \"lizard/images/00625.png\")\n\n    # temple different lighting\n    #image0 = load_image(path + \"multi-temporal-temple-baalshamin/images/image29.png\")\n    #image1 = load_image(path + \"multi-temporal-temple-baalshamin/images/dscn2009.png\")\n\n    # foliage\n    #image0 = load_image(path + \"pond/images/00025.png\")\n    #image1 = load_image(path + \"pond/images/00001.png\")\n\n    # transparent\n    #image0 = load_image(path + \"transp_obj_glass_cup/images/00.png\")\n    #image1 = load_image(path + \"transp_obj_glass_cup/images/34.png\")\n\n    feats0 = extractor.extract(image0.to(device))\n    feats1 = extractor.extract(image1.to(device))\n    matches01 = matcher({\"image0\": feats0, \"image1\": feats1})\n    feats0, feats1, matches01 = [\n        rbd(x) for x in [feats0, feats1, matches01]\n    ]  # remove batch dimension\n\n    kpts0, kpts1, matches = feats0[\"keypoints\"], feats1[\"keypoints\"], matches01[\"matches\"]\n    m_kpts0, m_kpts1 = kpts0[matches[..., 0]], kpts1[matches[..., 1]]\n\n    axes = viz2d.plot_images([image0, image1])\n    viz2d.plot_matches(m_kpts0, m_kpts1, color=\"lime\", lw=0.5)\n    viz2d.add_text(0, f'Stop after {matches01[\"stop\"]} layers', fs=20)\n\n    kpc0, kpc1 = viz2d.cm_prune(matches01[\"prune0\"]), viz2d.cm_prune(matches01[\"prune1\"])\n    viz2d.plot_images([image0, image1])\n    viz2d.plot_keypoints([kpts0, kpts1], colors=[kpc0, kpc1], ps=10)\n\n    print(m_kpts0.size())\n    print(m_kpts1.size())\n    print(kpts0.size())\n    print(kpts1.size())","metadata":{"execution":{"iopub.status.busy":"2024-04-22T17:24:02.421388Z","iopub.status.idle":"2024-04-22T17:24:02.421817Z","shell.execute_reply.started":"2024-04-22T17:24:02.421594Z","shell.execute_reply":"2024-04-22T17:24:02.421613Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Run v2","metadata":{}},{"cell_type":"code","source":"def run_v2(data_path,get_pairs,keypoints_matches,ransac_and_sparse_reconstruction,submit=True):\n    results = {}\n    \n    data_dict = parse_sample_submission(data_path)\n    datasets = list(data_dict.keys())\n    \n    for dataset in datasets:\n        if dataset not in results:\n            results[dataset] = {}\n            \n        for scene in data_dict[dataset]:\n            images_dir = data_dict[dataset][scene][0].parent\n            results[dataset][scene] = {}\n            image_paths = data_dict[dataset][scene]\n\n            index_pairs = get_pairs(image_paths)\n            rotation_dictionary = keypoints_matches(image_paths,index_pairs)                \n            maps = ransac_and_sparse_reconstruction(image_paths[0].parent)\n            clear_output(wait=False)\n            \n            path = 'test' if submit else 'train'\n            images_registered  = 0\n            best_idx = 0\n            for idx, rec in maps.items():\n                if len(rec.images) > images_registered:\n                    images_registered = len(rec.images)\n                    best_idx = idx\n                    \n            #print('rotation_dictionary')\n            #print(rotation_dictionary)\n            for k, im in maps[best_idx].images.items():\n                key = Path(IMC_PATH) / path / scene / \"images\" / im.name\n                \n                results[dataset][scene][key] = {}\n                calc_rotation = deepcopy(im.cam_from_world.rotation.matrix())\n                #print('calc_rotation')\n                #print(calc_rotation)\n                inv_rotation = inv_rot_matrix(rotation_dictionary[im.name])\n                #print('inv_rotation')\n                #print(inv_rotation)\n                results[dataset][scene][key][\"R\"] = np.matmul(inv_rotation,calc_rotation)\n                #print('final_rotation')\n                #print(results[dataset][scene][key][\"R\"])\n                \n                results[dataset][scene][key][\"t\"] = deepcopy(np.array(im.cam_from_world.translation))\n\n            create_submission(results, data_dict, Path(IMC_PATH))","metadata":{"execution":{"iopub.status.busy":"2024-04-22T17:26:17.458733Z","iopub.execute_input":"2024-04-22T17:26:17.459687Z","iopub.status.idle":"2024-04-22T17:26:17.470999Z","shell.execute_reply.started":"2024-04-22T17:26:17.459641Z","shell.execute_reply":"2024-04-22T17:26:17.470057Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## CROSS VALIDATION","metadata":{}},{"cell_type":"code","source":"SUBMISSION = False\nAPPLY_ROTATION = False\n#ROTATE_DATASET = ['']\n#ROTATE_DATASET = ['church']\n#ROTATE_DATASET = ['dioscuri']\nN_SAMPLES = 35\n\nif 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        if g[0] == ('dioscuri', 'dioscuri'):\n        #if g[0] == ('church', 'church'):\n        #print(g[0])\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            break\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_v2('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":{"scrolled":true,"execution":{"iopub.status.busy":"2024-04-22T17:33:39.974843Z","iopub.execute_input":"2024-04-22T17:33:39.975220Z","iopub.status.idle":"2024-04-22T17:37:44.140414Z","shell.execute_reply.started":"2024-04-22T17:33:39.975192Z","shell.execute_reply":"2024-04-22T17:37:44.139165Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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)\n    run_v2(data_path, get_pairs, keypoints_matches, ransac_and_sparse_reconstruction)","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2024-04-22T17:24:02.428107Z","iopub.status.idle":"2024-04-22T17:24:02.428844Z","shell.execute_reply.started":"2024-04-22T17:24:02.428602Z","shell.execute_reply":"2024-04-22T17:24:02.428623Z"},"trusted":true},"execution_count":null,"outputs":[]}]}