{"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":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 *\nimport ast\nimport shutil\nfrom pathlib import Path","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2024-06-03T01:54:29.036148Z","iopub.execute_input":"2024-06-03T01:54:29.036498Z","iopub.status.idle":"2024-06-03T01:55:11.251171Z","shell.execute_reply.started":"2024-06-03T01:54:29.036469Z","shell.execute_reply":"2024-06-03T01:55:11.250049Z"},"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-06-03T01:55:11.253354Z","iopub.execute_input":"2024-06-03T01:55:11.253667Z","iopub.status.idle":"2024-06-03T01:55:11.266349Z","shell.execute_reply.started":"2024-06-03T01:55:11.253637Z","shell.execute_reply":"2024-06-03T01:55:11.265386Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# KEYPOINTS EXTRACTOR AND MATCHER","metadata":{}},{"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        for pair in 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            with torch.inference_mode():\n                _, idxs = matcher(desc1, desc2, KF.laf_from_center_scale_ori(kp1[None]), KF.laf_from_center_scale_ori(kp2[None]))\n            if len(idxs): group = f_matches.require_group(key1)\n            if len(idxs) >= MIN_MATCHES: group.create_dataset(key2, data=idxs.detach().cpu().numpy())\n    multiview_consistency_pruning(images_list)","metadata":{"execution":{"iopub.status.busy":"2024-06-03T01:55:11.267531Z","iopub.execute_input":"2024-06-03T01:55:11.267815Z","iopub.status.idle":"2024-06-03T01:55:11.919325Z","shell.execute_reply.started":"2024-06-03T01:55:11.267791Z","shell.execute_reply":"2024-06-03T01:55:11.918218Z"},"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-06-03T01:55:11.921746Z","iopub.execute_input":"2024-06-03T01:55:11.922055Z","iopub.status.idle":"2024-06-03T01:55:11.933032Z","shell.execute_reply.started":"2024-06-03T01:55:11.922030Z","shell.execute_reply":"2024-06-03T01:55:11.932163Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Multiview Consistency Pruning","metadata":{}},{"cell_type":"code","source":"def all_tentative_pairs(images_list):\n    tentative_pairs = []\n    with h5py.File(\"matches.h5\", mode=\"r\") as f_matches:\n        a = list(f_matches.keys())\n        for index_a in a:\n            b = f_matches[index_a]\n            for index_b in b:\n                lista = [imagem.name for imagem in images_list]\n                tentative_pairs.append([lista.index(index_a),lista.index(index_b)])\n    return tentative_pairs\n\ndef create_pairs_matrix (pairs,size):\n    pairs_matrix = np.zeros([size,size], dtype=np.int8)\n    for k in pairs:\n        i,j = k[0],k[1]\n        pairs_matrix[i][j] = 1\n        pairs_matrix[j][i] = 1\n    return pairs_matrix\n\ndef find_triangles_by_pair(pair,pairs_matrix,length):\n    a,b = pair[0],pair[1]\n    triplets = []\n    for k in range(0,length):\n        if a==k or k==b or pairs_matrix[a][k] == 0 or pairs_matrix[k][b] == 0:\n            continue\n        triplets.append([a,k,b])\n    return triplets\n\ndef find_all_triplets(index_pairs,n_images):\n    pairs_matrix = create_pairs_matrix(index_pairs,n_images)\n    all_triplets = []\n    for index_pair in index_pairs:\n        triplets = find_triangles_by_pair(index_pair,pairs_matrix,n_images)\n        if len(triplets) != 0:\n            all_triplets.append(triplets)  \n    return all_triplets \n                        # retorna lista com todos os triangulos para cada par: list[list, list,..., list] \n                        # [ [ [a1,k1,b1],[a1,k2,b1],...,[a1,ki,b1] ] ,...,]\n\ndef partial_assignment_matrix(pair,images_list):\n    a = images_list[pair[0]].name\n    b = images_list[pair[1]].name\n    \n    transposed = False\n    \n    with h5py.File(\"matches.h5\", mode=\"r\") as f_matches:\n        try:\n            M_a = f_matches[a]\n            M_ab = M_a[b].__array__()\n        except:\n            M_b = f_matches[b]\n            M_ba = M_b[a].__array__()\n            transposed = True\n    if not transposed:\n        matches = M_ab\n    else:\n        matches  = np.empty_like(M_ba)\n        matches[:, 0]  = M_ba[:,1]\n        matches[:, 1]  = M_ba[:,0]\n    return matches\n\ndef store_refined(X_ab_all_refined,images_list):\n    with h5py.File(\"matches.h5\", mode=\"r+\") as f_matches:\n        for pair_str,matches_str in X_ab_all_refined.items():\n            pair = ast.literal_eval(pair_str)\n            matches = np.array([ast.literal_eval(m) for m in matches_str])\n            a = images_list[pair[0]].name\n            b = images_list[pair[1]].name\n            del f_matches[f\"{a}/{b}\"]\n            \n            if len(matches) < MIN_MATCHES:\n                continue    \n            f_matches[f\"{a}/{b}\"] = np.array(matches)\n    return\n\ndef multiview_consistency_pruning(images_list):\n    n_images = len(images_list)\n    pairs = all_tentative_pairs(images_list)\n    pairs_matrix = create_pairs_matrix(pairs,n_images)\n    visited_matches = {}\n    all_refined = {}\n    for pair in tqdm(pairs,desc=\"Multiview Pruning\"):\n        all_triplets_from_pair = find_triangles_by_pair(pair,pairs_matrix,n_images)\n        if len(all_triplets_from_pair) == 0: \n            continue\n        first_triplet_from_pair = all_triplets_from_pair[0]\n        t = first_triplet_from_pair\n        a,k,b = t[0],t[1],t[2]\n        \n        if str([a,b]) in visited_matches:\n            X_ab_original = visited_matches[str([a,b])]\n        else:\n            X_ab_original = partial_assignment_matrix([a,b],images_list)\n            visited_matches[str([a,b])] = X_ab_original\n        original_dict = {}\n        for numeric in X_ab_original.tolist():\n            original_dict[str(numeric)] = 1\n        refined_dict = {}\n        for triplet in all_triplets_from_pair:\n            k = triplet[1]\n            \n            if str([a,k]) in visited_matches:\n                partial_aki = visited_matches[str([a,k])]\n            else:\n                partial_aki = partial_assignment_matrix([a,k],images_list)\n                visited_matches[str([a,k])] = partial_aki\n            if str([k,b]) in visited_matches:\n                partial_kib = visited_matches[str([k,b])]\n            else:\n                partial_kib = partial_assignment_matrix([k,b],images_list)\n                visited_matches[str([k,b])] = partial_kib\n            matmul = {}\n            for pair_ak in partial_aki:\n                m,n = pair_ak[0],pair_ak[1]\n                for pair_kb in partial_kib:\n                    if pair_kb[0] == n:\n                        matmul[str([m,pair_kb[1]])] = 1\n                        break\n            X_ab_kihop = matmul\n            if len(X_ab_kihop) == 0:\n                continue\n            for match in X_ab_kihop:\n                if match not in refined_dict:\n                    if match in original_dict:\n                        refined_dict[match]=1\n        all_refined[str([a,b])] = refined_dict\n    store_refined(all_refined,images_list)\n    return","metadata":{"execution":{"iopub.status.busy":"2024-06-03T02:13:14.630814Z","iopub.execute_input":"2024-06-03T02:13:14.631633Z","iopub.status.idle":"2024-06-03T02:13:14.657742Z","shell.execute_reply.started":"2024-06-03T02:13:14.631602Z","shell.execute_reply":"2024-06-03T02:13:14.656788Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# HYPERPARAMETER TUNING","metadata":{}},{"cell_type":"code","source":"# SIMILLIAR PAIRS\nEXHAUSTIVE = False\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\n\n# RANSAC AND SPARSE RECONSTRUCTION\nMIN_MODEL_SIZE = 5\nMAX_NUM_MODELS = 3\n\n# CROSS VALIDATION\nN_SAMPLES = 20\n\nSUBMISSION = True","metadata":{"execution":{"iopub.status.busy":"2024-06-03T01:55:11.962250Z","iopub.execute_input":"2024-06-03T01:55:11.962478Z","iopub.status.idle":"2024-06-03T01:55:11.974545Z","shell.execute_reply.started":"2024-06-03T01:55:11.962457Z","shell.execute_reply":"2024-06-03T01:55:11.973748Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## CROSS VALIDATION","metadata":{}},{"cell_type":"code","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":{"scrolled":true,"execution":{"iopub.status.busy":"2024-06-03T02:13:17.648842Z","iopub.execute_input":"2024-06-03T02:13:17.649226Z","iopub.status.idle":"2024-06-03T02:20:13.759800Z","shell.execute_reply.started":"2024-06-03T02:13:17.649199Z","shell.execute_reply":"2024-06-03T02:20:13.758507Z"},"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)    ","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2024-06-03T02:07:21.314225Z","iopub.status.idle":"2024-06-03T02:07:21.314698Z","shell.execute_reply.started":"2024-06-03T02:07:21.314457Z","shell.execute_reply":"2024-06-03T02:07:21.314477Z"},"trusted":true},"execution_count":null,"outputs":[]}]}