{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# ***Ensemble of the following works***\n- SuperGlue\n    https://www.kaggle.com/code/losveria/superglue-baseline/notebook\n- LoFTR\n    https://www.kaggle.com/code/cbeaud/imc-2022-kornia-score-0-725/notebook\n- DKM\n    https://www.kaggle.com/code/radac98/public-baseline-dkm-0-667\n","metadata":{}},{"cell_type":"code","source":"dry_run = True","metadata":{"execution":{"iopub.status.busy":"2022-07-26T13:22:28.486422Z","iopub.execute_input":"2022-07-26T13:22:28.487133Z","iopub.status.idle":"2022-07-26T13:22:28.491831Z","shell.execute_reply.started":"2022-07-26T13:22:28.487096Z","shell.execute_reply":"2022-07-26T13:22:28.490067Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# *Import dependencies and install Libs*","metadata":{}},{"cell_type":"code","source":"import logging\nimport time\nimport os\nimport numpy as np\nimport cv2\nimport csv\nfrom glob import glob\nimport torch\nimport matplotlib.pyplot as plt\nimport gc\nimport pandas as pd\n\n# the following dependencies are for superglue\nimport random\nfrom collections import namedtuple\nimport sys\nsys.path.append(\"../input/super-glue-pretrained-network\")\nfrom models.matching import Matching\nfrom models.utils import frame2tensor\n\n# the following dependencies are for loftr\n!pip install ../input/kornia-loftr/kornia-0.6.4-py2.py3-none-any.whl\n!pip install ../input/kornia-loftr/kornia_moons-0.1.9-py3-none-any.whl\nimport kornia\nfrom kornia_moons.feature import *\nimport kornia as K\nimport kornia.feature as KF\n\n\n# the following dependencies are for DKM\nfrom PIL import Image\nsys.path.append('/kaggle/input/dkm-dependecies/DKM/')\n!mkdir -p pretrained/checkpoints\n!cp /kaggle/input/dkm-dependecies/pretrained/dkm.pth pretrained/checkpoints/dkm_base_v11.pth\n!pip install -f /kaggle/input/dkm-dependecies/wheels --no-index einops\n!cp -r /kaggle/input/dkm-dependecies/DKM/ /kaggle/working/DKM/\n!cd /kaggle/working/DKM/; pip install -f /kaggle/input/dkm-dependecies/wheels -e . \ntorch.hub.set_dir('/kaggle/working/pretrained/')\nfrom dkm import dkm_base","metadata":{"execution":{"iopub.status.idle":"2022-07-26T13:24:05.989422Z","shell.execute_reply.started":"2022-07-26T13:22:28.575806Z","shell.execute_reply":"2022-07-26T13:24:05.988592Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# *Settings*","metadata":{}},{"cell_type":"code","source":"device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n# src = '/kaggle/input/image-matching-challenge-2022/'\nsrc = '/kaggle/input/matching-test-images/own_test'","metadata":{"execution":{"iopub.status.busy":"2022-07-26T13:24:05.991492Z","iopub.execute_input":"2022-07-26T13:24:05.991792Z","iopub.status.idle":"2022-07-26T13:24:06.022458Z","shell.execute_reply.started":"2022-07-26T13:24:05.991755Z","shell.execute_reply":"2022-07-26T13:24:06.021577Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## *Utils*","metadata":{}},{"cell_type":"code","source":"def BuildCompositeImage(im1, im2, axis=1, margin=0, background=1):\n    '''Convenience function to stack two images with different sizes.'''\n    \n    if background != 0 and background != 1:\n        background = 1\n    if axis != 0 and axis != 1:\n        raise RuntimeError('Axis must be 0 (vertical) or 1 (horizontal')\n    im1 = np.asarray(im1)\n    im2 = np.asarray(im2)\n\n    h1, w1, _ = im1.shape\n    h2, w2, _ = im2.shape\n\n    if axis == 1:\n        composite = np.zeros((max(h1, h2), w1 + w2 + margin, 3), dtype=np.uint8) + 255 * background\n        if h1 > h2:\n            voff1, voff2 = 0, (h1 - h2) // 2\n        else:\n            voff1, voff2 = (h2 - h1) // 2, 0\n        hoff1, hoff2 = 0, w1 + margin\n    else:\n        composite = np.zeros((h1 + h2 + margin, max(w1, w2), 3), dtype=np.uint8) + 255 * background\n        if w1 > w2:\n            hoff1, hoff2 = 0, (w1 - w2) // 2\n        else:\n            hoff1, hoff2 = (w2 - w1) // 2, 0\n        voff1, voff2 = 0, h1 + margin\n    composite[voff1:voff1 + h1, hoff1:hoff1 + w1, :] = im1\n    composite[voff2:voff2 + h2, hoff2:hoff2 + w2, :] = im2\n\n    return (composite, (voff1, voff2), (hoff1, hoff2))\n\n\ndef DrawMatches(im1, im2, kp1, kp2, axis=1, margin=0, background=0, linewidth=2):\n    '''Draw keypoints and matches.'''\n    \n    composite, v_offset, h_offset = BuildCompositeImage(im1, im2, axis, margin, background)\n\n    # Draw all keypoints.\n    for coord_a, coord_b in zip(kp1, kp2):\n        composite = cv2.drawMarker(composite, (int(coord_a[0] + h_offset[0]), int(coord_a[1] + v_offset[0])), color=(255, 0, 0), markerType=cv2.MARKER_CROSS, markerSize=5, thickness=1)\n        composite = cv2.drawMarker(composite, (int(coord_b[0] + h_offset[1]), int(coord_b[1] + v_offset[1])), color=(255, 0, 0), markerType=cv2.MARKER_CROSS, markerSize=5, thickness=1)\n    \n    # Draw matches, and highlight keypoints used in matches.\n#     for idx_a, idx_b in matches:\n#         composite = cv2.drawMarker(composite, (int(kp1[idx_a, 0] + h_offset[0]), int(kp1[idx_a, 1] + v_offset[0])), color=(0, 0, 255), markerType=cv2.MARKER_CROSS, markerSize=12, thickness=1)\n#         composite = cv2.drawMarker(composite, (int(kp2[idx_b, 0] + h_offset[1]), int(kp2[idx_b, 1] + v_offset[1])), color=(0, 0, 255), markerType=cv2.MARKER_CROSS, markerSize=12, thickness=1)\n#         composite = cv2.line(composite,\n#                              tuple([int(kp1[idx_a][0] + h_offset[0]),\n#                                    int(kp1[idx_a][1] + v_offset[0])]),\n#                              tuple([int(kp2[idx_b][0] + h_offset[1]),\n#                                    int(kp2[idx_b][1] + v_offset[1])]), color=(0, 0, 255), thickness=1)\n    return composite","metadata":{"execution":{"iopub.status.busy":"2022-07-26T13:24:06.023772Z","iopub.execute_input":"2022-07-26T13:24:06.024213Z","iopub.status.idle":"2022-07-26T13:24:06.040092Z","shell.execute_reply.started":"2022-07-26T13:24:06.024172Z","shell.execute_reply":"2022-07-26T13:24:06.039403Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_samples = []\nwith open(f'{src}/test.csv') as f:\n    reader = csv.reader(f, delimiter=',')\n    for i, row in enumerate(reader):\n        # Skip header.\n        if i == 0:\n            continue\n        test_samples += [row]\n\n\ndef FlattenMatrix(M, num_digits=8):\n    '''Convenience function to write CSV files.'''\n    return ' '.join([f'{v:.{num_digits}e}' for v in M.flatten()])\n\n\ndef resize_img_loftr(img, max_len, enlarge_scale, variant_scale, device):\n    if max_len == -1:\n        scale = 1\n    else:\n        scale = max(max_len, max(img.shape[0], img.shape[1]) * enlarge_scale) / max(img.shape[0], img.shape[1])\n    w = int(round(img.shape[1] * scale) / 8) * 8\n    h = int(round(img.shape[0] * scale) / 8) * 8\n    \n    isResized = False\n    if w >= h:\n        if int(h * variant_scale) <= w:\n            isResized = True\n            h = int(h * variant_scale / 8) * 8\n    else:\n        if int(w * variant_scale) <= h:\n            isResized = True\n            w = int(w * variant_scale / 8) * 8\n    img_resize = cv2.resize(img, (w, h)) \n    img_resize = K.image_to_tensor(img_resize, False).float() / 255.\n    \n    return img_resize.to(device), (w / img.shape[1], h / img.shape[0]), isResized\n\n\ndef resize_img_superglue(img, max_len, enlarge_scale, variant_scale):\n    if max_len == -1:\n        scale = 1\n    else:\n        scale = max(max_len, max(img.shape[0], img.shape[1]) * enlarge_scale) / max(img.shape[0], img.shape[1])\n    w = int(round(img.shape[1] * scale))\n    h = int(round(img.shape[0] * scale))\n    \n    isResized = False\n    if w >= h:\n        if int(h * variant_scale) <= w:\n            isResized = True\n            h = int(h * variant_scale) \n    else:\n        if int(w * variant_scale) <= h:\n            isResized = True\n            w = int(w * variant_scale)\n    img_resize = cv2.resize(img, (w, h)) \n    return img_resize, (w / img.shape[1], h / img.shape[0]), isResized","metadata":{"execution":{"iopub.status.busy":"2022-07-26T13:24:06.041960Z","iopub.execute_input":"2022-07-26T13:24:06.042365Z","iopub.status.idle":"2022-07-26T13:24:06.061639Z","shell.execute_reply.started":"2022-07-26T13:24:06.042327Z","shell.execute_reply":"2022-07-26T13:24:06.060908Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# *Matchers*","metadata":{}},{"cell_type":"code","source":"# ===========================\n#          SuperGlue\n# ===========================\nresize = [-1, ]\nresize_float = True\nconfig = {\n    \"superpoint\": {\n        \"nms_radius\": 4, \n        \"keypoint_threshold\": 0.005,\n        \"max_keypoints\": 2048\n    },\n    \"superglue\": {\n        \"weights\": \"outdoor\",\n        \"sinkhorn_iterations\": 10,\n        \"match_threshold\": 0.2,\n    }\n}\nmatching_superglue = Matching(config).eval().to(device)\n\n# ===========================\n#          LoFTR\n# ===========================\nmatcher_loftr = KF.LoFTR(pretrained=None)\nmatcher_loftr.load_state_dict(torch.load(\"../input/kornia-loftr/loftr_outdoor.ckpt\")['state_dict'])\nmatcher_loftr = matcher_loftr.to(device).eval()\n\n# ===========================\n#          DKM\n# ===========================\nmodel = dkm_base(pretrained=True, version=\"v11\")","metadata":{"execution":{"iopub.status.busy":"2022-07-26T13:24:06.063217Z","iopub.execute_input":"2022-07-26T13:24:06.064052Z","iopub.status.idle":"2022-07-26T13:24:09.395335Z","shell.execute_reply.started":"2022-07-26T13:24:06.064023Z","shell.execute_reply":"2022-07-26T13:24:09.394503Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Inference","metadata":{}},{"cell_type":"code","source":"F_dict = {}\n\nscales_lens_superglue = [[1.2, 1200, 1.0], [1.2, 1600, 1.6], [0.8, 2000, 2], [1, 2800, 3]]\n\nscales_lens_loftr = [[1.1, 1000, 1.0], [1, 1200, 1.3], [0.9, 1400, 1.6]]\n\nw_h_muts_dkm = [[680 * 510, 1]]\n\nnp.random.seed(42)\n\nDEBUG = True\n# DEBUG = False\n\nif DEBUG == True:\n    import time\n    st = time.time()\n\nmatched_keypoints = {}\nwith torch.no_grad():\n    for i, row in enumerate(test_samples):\n        sample_id, batch_id, image_0_id, image_1_id = row    \n        \n        image_0_BGR = cv2.imread(f'{src}/{batch_id}/{image_0_id}.jpeg') \n        image_1_BGR = cv2.imread(f'{src}/{batch_id}/{image_1_id}.jpeg')\n        \n        image_0_GRAY = cv2.cvtColor(image_0_BGR, cv2.COLOR_BGR2GRAY)\n        image_1_GRAY = cv2.cvtColor(image_1_BGR, cv2.COLOR_BGR2GRAY)\n        \n        # ===========================\n        #           LoFTR\n        # ===========================\n        mkpts0_loftr_all = []\n        mkpts1_loftr_all = []\n        for variant_scale, max_len, enlarge_scale in scales_lens_loftr:\n            \n            image_0_resize, scale_0, isResized_0 = resize_img_loftr(image_0_GRAY, max_len, enlarge_scale, variant_scale, device)\n            image_1_resize, scale_1, isResized_1 = resize_img_loftr(image_1_GRAY, max_len, enlarge_scale, variant_scale, device)\n            \n            if isResized_0 == False or isResized_1 == False: continue\n            \n            input_dict = {\"image0\": image_0_resize, \n                      \"image1\": image_1_resize}\n            correspondences = matcher_loftr(input_dict)\n            confidence = correspondences['confidence'].cpu().numpy()\n            \n            if len(confidence) < 1: continue\n\n            confidence_quantile = np.quantile(confidence, 0.6)\n            idx = np.where(confidence >= confidence_quantile)\n            \n            mkpts0_loftr = correspondences['keypoints0'].cpu().numpy()[idx]\n            mkpts1_loftr = correspondences['keypoints1'].cpu().numpy()[idx]\n            \n            if DEBUG == True:\n                print(\"loftr scale_0\", scale_0)\n                print(\"loftr scale_1\", scale_1)\n\n            mkpts0_loftr = mkpts0_loftr / scale_0\n            mkpts1_loftr = mkpts1_loftr / scale_1\n\n            mkpts0_loftr_all.append(mkpts0_loftr)\n            mkpts1_loftr_all.append(mkpts1_loftr)\n        \n        mkpts0_loftr_all = np.concatenate(mkpts0_loftr_all, axis=0)\n        mkpts1_loftr_all = np.concatenate(mkpts1_loftr_all, axis=0) \n        \n        \n        # ===========================\n        #          SuperGlue\n        # ===========================\n        mkpts0_superglue_all = []\n        mkpts1_superglue_all = []\n        \n        for variant_scale, max_len, enlarge_scale in scales_lens_superglue:\n            image_0, scale_0, isResized_0 = resize_img_superglue(image_0_GRAY, max_len, enlarge_scale, variant_scale)\n            image_1, scale_1, isResized_1 = resize_img_superglue(image_1_GRAY, max_len, enlarge_scale, variant_scale)\n            \n            if isResized_0 == False or isResized_1 == False: break \n            \n            image_0 = frame2tensor(image_0, device)\n            image_1 = frame2tensor(image_1, device)\n\n            pred = matching_superglue({\"image0\": image_0, \"image1\": image_1})\n            pred = {k: v[0].detach().cpu().numpy() for k, v in pred.items()}\n            kpts0, kpts1 = pred[\"keypoints0\"], pred[\"keypoints1\"]\n            matches, conf = pred[\"matches0\"], pred[\"matching_scores0\"]\n            \n            valid = matches > -1\n            mkpts0_superglue = kpts0[valid]\n            mkpts1_superglue = kpts1[matches[valid]]\n\n            if DEBUG == True:\n                print(\"superglue scale_0\", scale_0)\n                print(\"superglue scale_1\", scale_1)\n\n            mkpts0_superglue /= scale_0\n            mkpts1_superglue /= scale_1\n\n            mkpts0_superglue_all.append(mkpts0_superglue)\n            mkpts1_superglue_all.append(mkpts1_superglue)\n            \n        if len(mkpts0_superglue_all) > 0:\n            mkpts0_superglue_all = np.concatenate(mkpts0_superglue_all, axis=0)\n            mkpts1_superglue_all = np.concatenate(mkpts1_superglue_all, axis=0) \n        \n        \n        # ===========================\n        #            DKM\n        # ===========================\n        img0PIL = Image.fromarray(cv2.cvtColor(image_0_BGR, cv2.COLOR_BGR2RGB))\n        img1PIL = Image.fromarray(cv2.cvtColor(image_1_BGR, cv2.COLOR_BGR2RGB))\n        \n        mkpts0_dkm_all = []\n        mkpts1_dkm_all = []\n        \n        for w_h_mut, param in w_h_muts_dkm:\n            \n            ratio = (image_0_BGR.shape[0] + image_1_BGR.shape[0]) / (image_0_BGR.shape[1] + image_1_BGR.shape[1]) * param # 根据图0的高宽比确定计算参数\n            \n            model.w_resized = int(np.sqrt(w_h_mut / ratio))\n            model.h_resized = int(ratio * model.w_resized)\n            \n            dense_matches, dense_certainty = model.match(img0PIL, img1PIL, do_pred_in_og_res=True)\n            dense_certainty = dense_certainty.pow(0.6)\n            \n            sparse_matches, sparse_certainty = model.sample(dense_matches, dense_certainty, max(min(500, (len(mkpts0_loftr_all) + len(mkpts0_superglue_all)) // int(4 * len(w_h_muts_dkm))), 100), 0.01)\n            mkpts0_dkm = sparse_matches[:, :2]\n            mkpts1_dkm = sparse_matches[:, 2:]\n            h, w, c = image_0_BGR.shape\n            mkpts0_dkm[:, 0] = ((mkpts0_dkm[:, 0] + 1) / 2) * w\n            mkpts0_dkm[:, 1] = ((mkpts0_dkm[:, 1] + 1) / 2) * h\n            h, w, c = image_1_BGR.shape\n            mkpts1_dkm[:, 0] = ((mkpts1_dkm[:, 0] + 1) / 2) * w\n            mkpts1_dkm[:, 1] = ((mkpts1_dkm[:, 1] + 1) / 2) * h\n\n            mkpts0_dkm_all.append(mkpts0_dkm)\n            mkpts1_dkm_all.append(mkpts1_dkm)\n\n        mkpts0_dkm_all = np.concatenate(mkpts0_dkm_all, axis=0)\n        mkpts1_dkm_all = np.concatenate(mkpts1_dkm_all, axis=0)\n        \n        # ensemble of all mkpts \n        mkpts0 = []\n        mkpts1 = []\n        \n        if len(mkpts0_loftr_all) > 0:\n            mkpts0.append(mkpts0_loftr_all)\n            mkpts1.append(mkpts1_loftr_all)\n        \n        if len(mkpts0_superglue_all) > 0:\n            mkpts0.append(mkpts0_superglue_all)\n            mkpts1.append(mkpts1_superglue_all)\n        \n        mkpts0.append(mkpts0_dkm_all)\n        mkpts1.append(mkpts1_dkm_all)\n        \n        mkpts0 = np.concatenate(mkpts0, axis=0)\n        mkpts1 = np.concatenate(mkpts1, axis=0)\n        \n        \n        if len(mkpts0) > 8:\n            F, inliers = cv2.findFundamentalMat(mkpts0, mkpts1, cv2.USAC_MAGSAC, 0.15, 0.9999, 20000)\n            F_dict[sample_id] = F\n                    \n            \n        else:\n            F_dict[sample_id] = np.zeros((3, 3))\n            continue\n        \n        if DEBUG == True: \n            print(\"the number of loftr keypoints: \", len(mkpts0_loftr_all))\n            print(\"the number of superglue keypoints: \", len(mkpts0_superglue_all))\n            print(\"the number of dkm keypoints: \", len(mkpts0_dkm_all))\n            print(\"the number of all keypoints: \", len(mkpts0))\n            print(\"Fundamental matrix: \")\n            print(F_dict[sample_id])\n        \n        matched_keypoints[i] = {\n            'img1': image_0_BGR,\n            'img2': image_1_BGR,\n            'mkpts0' : mkpts0,\n            'mkpts1': mkpts1,\n            'F': F,\n            'inliers': inliers\n        }\n        \n   \n        gc.collect()\n\nif DEBUG == True:\n    ed = time.time()\n    print(f\"spend {ed - st:.2f}s\")\n    \nwith open('submission.csv', 'w') as f:\n    f.write('sample_id,fundamental_matrix\\n')\n    for sample_id, F in F_dict.items():\n        f.write(f'{sample_id},{FlattenMatrix(F)}\\n')","metadata":{"execution":{"iopub.status.busy":"2022-07-26T13:31:28.961179Z","iopub.execute_input":"2022-07-26T13:31:28.961445Z","iopub.status.idle":"2022-07-26T13:31:39.932950Z","shell.execute_reply.started":"2022-07-26T13:31:28.961416Z","shell.execute_reply":"2022-07-26T13:31:39.932206Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%matplotlib inline\nVISUALIZE = True\nimport time\n\nimport kornia_moons.feature as KMF\n\n\nif VISUALIZE:\n    st = time.time()\n    for i in range(0,len(matched_keypoints),1):\n        mkpts0, mkpts1, img_np1, img_np2 = matched_keypoints[i]['mkpts0'], matched_keypoints[i]['mkpts1'], matched_keypoints[i]['img1'], matched_keypoints[i]['img2']\n        F, inliers = matched_keypoints[i]['F'], matched_keypoints[i]['inliers']\n\n        gc_st = time.time()\n        gc.collect()\n        nd = time.time()    \n        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))\n        KMF.draw_LAF_matches(\n            KF.laf_from_center_scale_ori(torch.from_numpy(mkpts0).view(1,-1, 2),\n                                         torch.ones(mkpts0.shape[0]).view(1,-1, 1, 1),\n                                         torch.ones(mkpts0.shape[0]).view(1,-1, 1)),\n\n            KF.laf_from_center_scale_ori(torch.from_numpy(mkpts1).view(1,-1, 2),\n                                         torch.ones(mkpts1.shape[0]).view(1,-1, 1, 1),\n                                         torch.ones(mkpts1.shape[0]).view(1,-1, 1)),\n            torch.arange(mkpts0.shape[0]).view(-1,1).repeat(1,2),\n            cv2.cvtColor(img_np1, cv2.COLOR_BGR2RGB),\n            cv2.cvtColor(img_np2, cv2.COLOR_BGR2RGB),\n            inliers,\n            draw_dict={'inlier_color': (0.2, 1, 0.2),\n                       'tentative_color': None, \n                       'feature_color': (0.2, 0.5, 1), 'vertical': False},\n        )\n\ngc.collect()","metadata":{"execution":{"iopub.status.busy":"2022-07-26T13:31:39.934478Z","iopub.execute_input":"2022-07-26T13:31:39.934827Z","iopub.status.idle":"2022-07-26T13:32:07.495920Z","shell.execute_reply.started":"2022-07-26T13:31:39.934790Z","shell.execute_reply":"2022-07-26T13:32:07.494302Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}