{"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":"## Baseline submission\n\nA notebook to generate a valid submission. Implements three local feature/matcher methods: LoFTR, DISK, and KeyNetAffNetHardNet.\n\nRemember to enable a GPU accelerator and disable internet access, then press \"submit\" on the right pane.","metadata":{}},{"cell_type":"code","source":"SUBMIT = True\nSET = \"test\" if SUBMIT else \"train\"\nNUM_FEATS = 2048 if SUBMIT else 1000\nRESIZE_SMALL_EDGE_TO=800\nif SUBMIT:\n    DESCRIPTOR = \"HARDNET\"\nelse:\n    DESCRIPTOR = \"HARDNET\"\nVERBOSE = False\nCOMPOSITE_DESCRIPTOR = True\nUPRIGHT = False","metadata":{"execution":{"iopub.status.busy":"2023-06-11T09:03:07.008589Z","iopub.execute_input":"2023-06-11T09:03:07.009659Z","iopub.status.idle":"2023-06-11T09:03:07.040157Z","shell.execute_reply.started":"2023-06-11T09:03:07.009596Z","shell.execute_reply":"2023-06-11T09:03:07.039145Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# General utilities\nimport os\nfrom tqdm import tqdm\nfrom time import time\nfrom fastprogress import progress_bar\nimport gc\nimport numpy as np\nimport pandas as pd\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\n\n\nfrom imc_2023_eval import eval_submission\nfrom get_bounding_boxes import get_crop\nfrom imc_2023_helper_functions import *\nimport pandas as pd\nfrom itertools import product","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-06-11T09:03:07.042492Z","iopub.execute_input":"2023-06-11T09:03:07.042948Z","iopub.status.idle":"2023-06-11T09:03:12.460179Z","shell.execute_reply.started":"2023-06-11T09:03:07.042902Z","shell.execute_reply":"2023-06-11T09:03:12.458911Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print('Kornia version', K.__version__)\nprint('Pycolmap version', pycolmap.__version__)\n\nLOCAL_FEATURE = 'SUPERPOINT'#'LoFTR'#'KeyNetAffNetHardNet'\nif torch.cuda.is_available():\n    device = torch.device('cuda')\nelse:\n    device=torch.device('cpu')\n# Can be LoFTR, KeyNetAffNetHardNet, or DISK","metadata":{"execution":{"iopub.status.busy":"2023-06-11T09:03:12.462674Z","iopub.execute_input":"2023-06-11T09:03:12.463028Z","iopub.status.idle":"2023-06-11T09:03:12.522639Z","shell.execute_reply.started":"2023-06-11T09:03:12.462977Z","shell.execute_reply":"2023-06-11T09:03:12.521246Z"},"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":"2023-06-11T09:03:12.524782Z","iopub.execute_input":"2023-06-11T09:03:12.525612Z","iopub.status.idle":"2023-06-11T09:03:12.534080Z","shell.execute_reply.started":"2023-06-11T09:03:12.525568Z","shell.execute_reply":"2023-06-11T09:03:12.532776Z"},"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\n        # ORI MODULE\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        \n        # DETECTOR\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',map_location=device)['state_dict']\n        detector.model.load_state_dict(kn_weights)\n        \n        affnet_weights = torch.load('/kaggle/input/kornia-local-feature-weights/AffNet.pth',map_location=device)['state_dict']\n        detector.aff.load_state_dict(affnet_weights)\n        \n        # DESCRIPTOR\n        if DESCRIPTOR==\"HARDNET\":\n            patch_descriptor_module = KF.HardNet(False).eval()\n            weights = torch.load('/kaggle/input/kornia-local-feature-weights/HardNetLib.pth',map_location=device)['state_dict']  \n        elif DESCRIPTOR==\"HYNET\":\n            patch_descriptor_module = KF.HyNet(False).eval()\n            weights = torch.load('/kaggle/input/kornia-local-feature-weights/HyNet_LIB.pth',map_location=device)#['state_dict']\n        elif DESCRIPTOR==\"SOSNET\":\n            patch_descriptor_module = KF.SOSNet(False).eval()\n            weights = torch.load('/kaggle/input/kornia-local-feature-weights/sosnet_32x32_liberty.pth',map_location=device)#['state_dict']\n        patch_descriptor_module.load_state_dict(weights)\n        descriptor = KF.LAFDescriptor(patch_descriptor_module, patch_size=32, grayscale_descriptor=True).to(device)\n        super().__init__(detector, descriptor, scale_laf)\n","metadata":{"execution":{"iopub.status.busy":"2023-06-11T09:03:12.537825Z","iopub.execute_input":"2023-06-11T09:03:12.538677Z","iopub.status.idle":"2023-06-11T09:03:12.552335Z","shell.execute_reply.started":"2023-06-11T09:03:12.538635Z","shell.execute_reply":"2023-06-11T09:03:12.551444Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class SuperPointNet(torch.nn.Module):\n    \"\"\" Pytorch definition of SuperPoint Network. \"\"\"\n    def __init__(self):\n        super(SuperPointNet, self).__init__()\n        self.relu = torch.nn.ReLU(inplace=True)\n        self.pool = torch.nn.MaxPool2d(kernel_size=2, stride=2)\n        c1, c2, c3, c4, c5, d1 = 64, 64, 128, 128, 256, 256\n        # Shared Encoder.\n        self.conv1a = torch.nn.Conv2d(1, c1, kernel_size=3, stride=1, padding=1)\n        self.conv1b = torch.nn.Conv2d(c1, c1, kernel_size=3, stride=1, padding=1)\n        self.conv2a = torch.nn.Conv2d(c1, c2, kernel_size=3, stride=1, padding=1)\n        self.conv2b = torch.nn.Conv2d(c2, c2, kernel_size=3, stride=1, padding=1)\n        self.conv3a = torch.nn.Conv2d(c2, c3, kernel_size=3, stride=1, padding=1)\n        self.conv3b = torch.nn.Conv2d(c3, c3, kernel_size=3, stride=1, padding=1)\n        self.conv4a = torch.nn.Conv2d(c3, c4, kernel_size=3, stride=1, padding=1)\n        self.conv4b = torch.nn.Conv2d(c4, c4, kernel_size=3, stride=1, padding=1)\n        # Detector Head.\n        self.convPa = torch.nn.Conv2d(c4, c5, kernel_size=3, stride=1, padding=1)\n        self.convPb = torch.nn.Conv2d(c5, 65, kernel_size=1, stride=1, padding=0)\n        # Descriptor Head.\n        self.convDa = torch.nn.Conv2d(c4, c5, kernel_size=3, stride=1, padding=1)\n        self.convDb = torch.nn.Conv2d(c5, d1, kernel_size=1, stride=1, padding=0)\n\n    def forward(self, x):\n        \"\"\" Forward pass that jointly computes unprocessed point and descriptor\n        tensors.\n        Input\n          x: Image pytorch tensor shaped N x 1 x H x W.\n        Output\n          semi: Output point pytorch tensor shaped N x 65 x H/8 x W/8.\n          desc: Output descriptor pytorch tensor shaped N x 256 x H/8 x W/8.\n        \"\"\"\n        # Shared Encoder.\n        x = self.relu(self.conv1a(x))\n        x = self.relu(self.conv1b(x))\n        x = self.pool(x)\n        x = self.relu(self.conv2a(x))\n        x = self.relu(self.conv2b(x))\n        x = self.pool(x)\n        x = self.relu(self.conv3a(x))\n        x = self.relu(self.conv3b(x))\n        x = self.pool(x)\n        x = self.relu(self.conv4a(x))\n        x = self.relu(self.conv4b(x))\n        # Detector Head.\n        cPa = self.relu(self.convPa(x))\n        semi = self.convPb(cPa)\n        # Descriptor Head.\n        cDa = self.relu(self.convDa(x))\n        desc = self.convDb(cDa)\n        dn = torch.norm(desc, p=2, dim=1) # Compute the norm.\n        desc = desc.div(torch.unsqueeze(dn, 1)) # Divide by norm to normalize.\n        return semi, desc\n\nclass SuperPointFrontend(object):\n    \"\"\" Wrapper around pytorch net to help with pre and post image processing. \"\"\"\n    def __init__(self, weights_path, nms_dist, conf_thresh, nn_thresh,\n               cuda=False):\n        self.name = 'SuperPoint'\n        self.cuda = cuda\n        self.nms_dist = nms_dist\n        self.conf_thresh = conf_thresh\n        self.nn_thresh = nn_thresh # L2 descriptor distance for good match.\n        self.cell = 8 # Size of each output cell. Keep this fixed.\n        self.border_remove = 4 # Remove points this close to the border.\n\n        # Load the network in inference mode.\n        self.net = SuperPointNet()\n        if cuda:\n          # Train on GPU, deploy on GPU.\n            self.net.load_state_dict(torch.load(weights_path))\n            self.net = self.net.cuda()\n        else:\n          # Train on GPU, deploy on CPU.\n            self.net.load_state_dict(torch.load(weights_path,\n                                   map_location=lambda storage, loc: storage))\n            self.net.eval()\n\n    def nms_fast(self, in_corners, H, W, dist_thresh):\n        \"\"\"\n        Run a faster approximate Non-Max-Suppression on numpy corners shaped:\n          3xN [x_i,y_i,conf_i]^T\n\n        Algo summary: Create a grid sized HxW. Assign each corner location a 1, rest\n        are zeros. Iterate through all the 1's and convert them either to -1 or 0.\n        Suppress points by setting nearby values to 0.\n\n        Grid Value Legend:\n        -1 : Kept.\n         0 : Empty or suppressed.\n         1 : To be processed (converted to either kept or supressed).\n\n        NOTE: The NMS first rounds points to integers, so NMS distance might not\n        be exactly dist_thresh. It also assumes points are within image boundaries.\n\n        Inputs\n          in_corners - 3xN numpy array with corners [x_i, y_i, confidence_i]^T.\n          H - Image height.\n          W - Image width.\n          dist_thresh - Distance to suppress, measured as an infinty norm distance.\n        Returns\n          nmsed_corners - 3xN numpy matrix with surviving corners.\n          nmsed_inds - N length numpy vector with surviving corner indices.\n        \"\"\"\n        grid = np.zeros((H, W)).astype(int) # Track NMS data.\n        inds = np.zeros((H, W)).astype(int) # Store indices of points.\n        # Sort by confidence and round to nearest int.\n        inds1 = np.argsort(-in_corners[2,:])\n        corners = in_corners[:,inds1]\n        rcorners = corners[:2,:].round().astype(int) # Rounded corners.\n        # Check for edge case of 0 or 1 corners.\n        if rcorners.shape[1] == 0:\n            return np.zeros((3,0)).astype(int), np.zeros(0).astype(int)\n        if rcorners.shape[1] == 1:\n            out = np.vstack((rcorners, in_corners[2])).reshape(3,1)\n            return out, np.zeros((1)).astype(int)\n        # Initialize the grid.\n        for i, rc in enumerate(rcorners.T):\n            grid[rcorners[1,i], rcorners[0,i]] = 1\n            inds[rcorners[1,i], rcorners[0,i]] = i\n        # Pad the border of the grid, so that we can NMS points near the border.\n        pad = dist_thresh\n        grid = np.pad(grid, ((pad,pad), (pad,pad)), mode='constant')\n        # Iterate through points, highest to lowest conf, suppress neighborhood.\n        count = 0\n        for i, rc in enumerate(rcorners.T):\n          # Account for top and left padding.\n            pt = (rc[0]+pad, rc[1]+pad)\n            if grid[pt[1], pt[0]] == 1: # If not yet suppressed.\n                grid[pt[1]-pad:pt[1]+pad+1, pt[0]-pad:pt[0]+pad+1] = 0\n                grid[pt[1], pt[0]] = -1\n            count += 1\n        # Get all surviving -1's and return sorted array of remaining corners.\n        keepy, keepx = np.where(grid==-1)\n        keepy, keepx = keepy - pad, keepx - pad\n        inds_keep = inds[keepy, keepx]\n        out = corners[:, inds_keep]\n        values = out[-1, :]\n        inds2 = np.argsort(-values)\n        out = out[:, inds2]\n        out_inds = inds1[inds_keep[inds2]]\n        return out, out_inds\n\n    def run(self, img):\n        \"\"\" Process a numpy image to extract points and descriptors.\n        Input\n        img - HxW numpy float32 input image in range [0,1].\n        Output\n        corners - 3xN numpy array with corners [x_i, y_i, confidence_i]^T.\n        desc - 256xN numpy array of corresponding unit normalized descriptors.\n        heatmap - HxW numpy heatmap in range [0,1] of point confidences.\n        \"\"\"\n        assert img.ndim == 2, 'Image must be grayscale.'\n        assert img.dtype == np.float32, 'Image must be float32.'\n        H, W = img.shape[0], img.shape[1]\n        inp = img.copy()\n        inp = (inp.reshape(1, H, W))\n        inp = torch.from_numpy(inp)\n        inp = torch.autograd.Variable(inp).view(1, 1, H, W)\n        if self.cuda:\n            inp = inp.cuda()\n        # Forward pass of network.\n        outs = self.net.forward(inp)\n        semi, coarse_desc = outs[0], outs[1]\n        # Convert pytorch -> numpy.\n        semi = semi.data.cpu().numpy().squeeze()\n        # --- Process points.\n        dense = np.exp(semi) # Softmax.\n        dense = dense / (np.sum(dense, axis=0)+.00001) # Should sum to 1.\n        # Remove dustbin.\n        nodust = dense[:-1, :, :]\n        # Reshape to get full resolution heatmap.\n        Hc = int(H / self.cell)\n        Wc = int(W / self.cell)\n        nodust = nodust.transpose(1, 2, 0)\n        heatmap = np.reshape(nodust, [Hc, Wc, self.cell, self.cell])\n        heatmap = np.transpose(heatmap, [0, 2, 1, 3])\n        heatmap = np.reshape(heatmap, [Hc*self.cell, Wc*self.cell])\n        xs, ys = np.where(heatmap >= self.conf_thresh) # Confidence threshold.\n        if len(xs) == 0:\n            return np.zeros((3, 0)), None, None\n        pts = np.zeros((3, len(xs))) # Populate point data sized 3xN.\n        pts[0, :] = ys\n        pts[1, :] = xs\n        pts[2, :] = heatmap[xs, ys]\n        pts, _ = self.nms_fast(pts, H, W, dist_thresh=self.nms_dist) # Apply NMS.\n        inds = np.argsort(pts[2,:])\n        pts = pts[:,inds[::-1]] # Sort by confidence.\n        # Remove points along border.\n        bord = self.border_remove\n        toremoveW = np.logical_or(pts[0, :] < bord, pts[0, :] >= (W-bord))\n        toremoveH = np.logical_or(pts[1, :] < bord, pts[1, :] >= (H-bord))\n        toremove = np.logical_or(toremoveW, toremoveH)\n        pts = pts[:, ~toremove]\n        # --- Process descriptor.\n        D = coarse_desc.shape[1]\n        if pts.shape[1] == 0:\n            desc = np.zeros((D, 0))\n        else:\n            # Interpolate into descriptor map using 2D point locations.\n            samp_pts = torch.from_numpy(pts[:2, :].copy())\n            samp_pts[0, :] = (samp_pts[0, :] / (float(W)/2.)) - 1.\n            samp_pts[1, :] = (samp_pts[1, :] / (float(H)/2.)) - 1.\n            samp_pts = samp_pts.transpose(0, 1).contiguous()\n            samp_pts = samp_pts.view(1, 1, -1, 2)\n            samp_pts = samp_pts.float()\n            if self.cuda:\n                samp_pts = samp_pts.cuda()\n            desc = torch.nn.functional.grid_sample(coarse_desc, samp_pts)\n            desc = desc.data.cpu().numpy().reshape(D, -1)\n            desc /= np.linalg.norm(desc, axis=0)[np.newaxis, :]\n        return pts, desc, heatmap\n\n\nsp_path = '/kaggle/input/superpoint-magicleap/superpoint_v1.pth'\nfrontend = SuperPointFrontend(sp_path, nms_dist=4, conf_thresh=0.015, nn_thresh=0.7, cuda=False)","metadata":{"execution":{"iopub.status.busy":"2023-06-11T09:03:12.554101Z","iopub.execute_input":"2023-06-11T09:03:12.554841Z","iopub.status.idle":"2023-06-11T09:03:12.716989Z","shell.execute_reply.started":"2023-06-11T09:03:12.554800Z","shell.execute_reply":"2023-06-11T09:03:12.715905Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"src = '/kaggle/input/image-matching-challenge-2023'\n\n# Get data from csv.\n\ndata_dict = {}\nif SUBMIT:\n    sample_submission_file = f'{src}/sample_submission.csv'\nelse:\n    sample_submission_file = f'{src}/train/train_labels.csv'\n    \nsample_df = pd.read_csv(sample_submission_file)\n\n\nfor i,j in sample_df.iterrows():\n    dataset = j[\"dataset\"]\n    scene = j[\"scene\"]\n    \n    data_dict.setdefault(dataset,{}).setdefault(scene,[]).append(j[\"image_path\"])\n\nif False:\n    for dataset,value in data_dict.items():\n        for scene,image_dir in value.items():\n            scene_dir = f\"{src}/{SET}/{dataset}/{scene}\"\n            if \"images_full\" in os.listdir(scene_dir):\n\n                image_dir = f\"{scene_dir}/images_full\"\n                for image in os.listdir(image_dir):\n                    data_dict[dataset][scene].append(os.path.join(image_dir,image))\n\nif not SUBMIT:\n    a = data_dict[\"heritage\"].pop(\"dioscuri\")\n#     b = data_dict.pop(\"haiper\")\nfor 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":"2023-06-11T09:03:12.718523Z","iopub.execute_input":"2023-06-11T09:03:12.718914Z","iopub.status.idle":"2023-06-11T09:03:12.766193Z","shell.execute_reply.started":"2023-06-11T09:03:12.718872Z","shell.execute_reply":"2023-06-11T09:03:12.765038Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not SUBMIT:\n    DESCRIPTOR = \"HYNET\"\n    HYNET = KeyNetAffNetHardNet(4000, True, device).to(device).eval()\n    DESCRIPTOR = \"HARDNET\"","metadata":{"execution":{"iopub.status.busy":"2023-06-11T09:03:12.767620Z","iopub.execute_input":"2023-06-11T09:03:12.768789Z","iopub.status.idle":"2023-06-11T09:03:15.494580Z","shell.execute_reply.started":"2023-06-11T09:03:12.768748Z","shell.execute_reply":"2023-06-11T09:03:15.493526Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not SUBMIT:\n    img_fnames = [f'{src}/{SET}/{x}' for x in data_dict[\"haiper\"][\"bike\"]]\n\n    DESCRIPTOR = \"HYNET\"\n    HYNET = KeyNetAffNetHardNet(1000, True, device).to(device).eval()\n    DESCRIPTOR = \"HARDNET\"\n    model_ = KF.SIFTFeature(2000)\n    crop = False\n    if crop:\n        features = [HYNET]#[model_]#[HARDNET1k]\n    else:\n        features = [HYNET]\n\n    for img_path in progress_bar(img_fnames[1:2]):\n        img_fname = img_path.split('/')[-1]\n        with torch.inference_mode():\n            timg = load_torch_image(img_path, device=device)\n            if crop:\n                timg=timg[:,:,x1:x2,y1:y2]\n            H, W = timg.shape[2:]\n            print(H,W)\n\n            timg_resized = timg#K.geometry.resize(timg, None, antialias=True)\n            h, w = timg_resized.shape[2:]\n            begin=True\n            for feature in features:\n                _lafs, _resps, _descs = feature(K.color.rgb_to_grayscale(timg_resized))\n                if begin:\n                    lafs,resps,descs = _lafs, _resps, _descs\n                    begin = False\n                else:\n                    lafs, resps,descs = torch.cat((lafs,_lafs),1),torch.cat((resps,_resps),1),torch.cat((descs,_descs),1)\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            lafs = lafs.detach().cpu().numpy()\n    print(len(kpts))","metadata":{"execution":{"iopub.status.busy":"2023-06-11T09:03:15.496813Z","iopub.execute_input":"2023-06-11T09:03:15.497649Z","iopub.status.idle":"2023-06-11T09:03:20.115558Z","shell.execute_reply.started":"2023-06-11T09:03:15.497600Z","shell.execute_reply":"2023-06-11T09:03:20.114395Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# laACCESS_FASTfs[0,0]\n# array([[ 1.5561136e+01,  0.0000000e+00,  4.2082513e+02],\n#        [-9.6237969e-01,  1.5556797e+01,  1.0098564e+03]], dtype=float32)","metadata":{"execution":{"iopub.status.busy":"2023-06-11T09:03:20.117437Z","iopub.execute_input":"2023-06-11T09:03:20.118216Z","iopub.status.idle":"2023-06-11T09:03:20.122882Z","shell.execute_reply.started":"2023-06-11T09:03:20.118173Z","shell.execute_reply":"2023-06-11T09:03:20.121817Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# with torch.inference_mode():\n#     a,b,c = frontend.run(K.color.rgb_to_grayscale(timg).detach().cpu().numpy()[0][0])","metadata":{"execution":{"iopub.status.busy":"2023-06-11T09:03:20.128893Z","iopub.execute_input":"2023-06-11T09:03:20.129253Z","iopub.status.idle":"2023-06-11T09:03:30.279607Z","shell.execute_reply.started":"2023-06-11T09:03:20.129210Z","shell.execute_reply":"2023-06-11T09:03:30.278317Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# with torch.inference_mode():\n#     l = KF.laf_from_center_scale_ori(torch.tensor(a[:2].transpose()).unsqueeze(0))","metadata":{"execution":{"iopub.status.busy":"2023-06-11T09:03:30.281554Z","iopub.execute_input":"2023-06-11T09:03:30.282028Z","iopub.status.idle":"2023-06-11T09:03:30.288768Z","shell.execute_reply.started":"2023-06-11T09:03:30.281961Z","shell.execute_reply":"2023-06-11T09:03:30.286081Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# l.shape","metadata":{"execution":{"iopub.status.busy":"2023-06-11T09:03:30.290746Z","iopub.execute_input":"2023-06-11T09:03:30.292364Z","iopub.status.idle":"2023-06-11T09:03:30.302174Z","shell.execute_reply.started":"2023-06-11T09:03:30.292321Z","shell.execute_reply":"2023-06-11T09:03:30.301077Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not SUBMIT:\n    import cv2\n    import matplotlib.pyplot as plt\n\n    imageread = cv2.imread(img_path)\n\n    # input image is converted to gray scale image\n    imagegray = cv2.cvtColor(imageread, cv2.COLOR_BGR2GRAY)\n    if crop:\n        imagegray = imagegray[x1:x2,y1:y2]\n    pts = cv2.KeyPoint_convert(kpts)\n\n    # drawKeypoints function is used to draw keypoints\n    output_image = cv2.drawKeypoints(imagegray, pts, 0, (255, 0, 0),\n                                     flags=cv2.DRAW_MATCHES_FLAGS_NOT_DRAW_SINGLE_POINTS)\n\n    # displaying the image with keypoints as the\n    # output on the screen\n    fig = plt.figure(figsize=(10,10))\n    ax = fig.gca()\n    ax.set_xticks(np.arange(0, W, 100))\n    ax.set_yticks(np.arange(0, H, 100))\n    plt.imshow(output_image)#[x1:x2,y1:y2])\n    plt.grid(10)\n    # plotting image\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2023-06-11T09:03:30.304493Z","iopub.execute_input":"2023-06-11T09:03:30.305020Z","iopub.status.idle":"2023-06-11T09:03:31.187785Z","shell.execute_reply.started":"2023-06-11T09:03:30.304956Z","shell.execute_reply":"2023-06-11T09:03:31.186659Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not SUBMIT:\n    import cv2\n    import matplotlib.pyplot as plt\n\n    imageread = cv2.imread(img_path)\n\n    # input image is converted to gray scale image\n    imagegray = cv2.cvtColor(imageread, cv2.COLOR_BGR2GRAY)\n    if crop:\n        imagegray = imagegray[x1:x2,y1:y2]\n    pts = cv2.KeyPoint_convert(kpts.astype(np.float32))\n\n    # drawKeypoints function is used to draw keypoints\n    output_image = cv2.drawKeypoints(imagegray, pts, 0, (255, 0, 0),\n                                     flags=cv2.DRAW_MATCHES_FLAGS_NOT_DRAW_SINGLE_POINTS)\n\n    # displaying the image with keypoints as the\n    # output on the screen\n    fig = plt.figure(figsize=(10,10))\n    ax = fig.gca()\n    ax.set_xticks(np.arange(0, W, 100))\n    ax.set_yticks(np.arange(0, H, 100))\n    plt.imshow(output_image)#[x1:x2,y1:y2])\n    plt.grid(10)\n    # plotting image\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2023-06-11T09:03:31.188936Z","iopub.execute_input":"2023-06-11T09:03:31.190324Z","iopub.status.idle":"2023-06-11T09:03:32.040341Z","shell.execute_reply.started":"2023-06-11T09:03:31.190277Z","shell.execute_reply":"2023-06-11T09:03:32.039381Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not SUBMIT:\n    data = kpts.astype(int)\n    data = np.unique(data,axis=0)   \n    len(data),len(data)/len(kpts)\n\n    boxes = get_crop(data,H,W)#v1 -> get corners too\n    [box for box in boxes if box[1]>3][0][0]\n\n    data = kpts.astype(int)\n    data = np.unique(data,axis=0)   \n    boxes = get_crop(data,H,W)#v1 -> get corners too\n    [box for box in boxes if box[1]>3][0][0]","metadata":{"execution":{"iopub.status.busy":"2023-06-11T09:03:32.041958Z","iopub.execute_input":"2023-06-11T09:03:32.042649Z","iopub.status.idle":"2023-06-11T09:03:32.072982Z","shell.execute_reply.started":"2023-06-11T09:03:32.042610Z","shell.execute_reply":"2023-06-11T09:03:32.071849Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_feats(feature_name,feature,timg,num_feats,resize_small_edge_to=None):\n    with torch.inference_mode():\n        \n        H, W = timg.shape[2:]\n        if resize_small_edge_to:\n            timg = K.geometry.resize(timg, resize_small_edge_to, antialias=True)\n        h, w = timg.shape[2:]\n        if feature_name==\"DISK\":\n            feats = feature(timg, num_feats, pad_if_not_divisible=True)[0]\n            kpts, descs = feats.keypoints, torch.unsqueeze(feats.descriptors,0)\n            lafs = KF.laf_from_center_scale_ori(kpts[None], torch.ones(1, len(kpts), 1, 1, device=device))\n        else:\n            lafs, _, descs = feature(K.color.rgb_to_grayscale(timg))\n        \n        lafs[:,:,0,:] *= float(W) / float(w)\n        lafs[:,:,1,:] *= float(H) / float(h)\n        kpts = KF.get_laf_center(lafs).reshape(-1, 2)\n        desc_dim = descs.shape[-1]\n        descs = descs.reshape(-1, desc_dim)\n        return lafs.detach().cpu().numpy(),descs.detach().cpu().numpy(),kpts.detach().cpu().numpy()\n        # return {\n        #     \"laf\":lafs.detach().cpu().numpy(),\n        #     \"desc\":descs.detach().cpu().numpy(),\n        #     \"kpts\":kpts.detach().cpu().numpy()\n        # }","metadata":{"execution":{"iopub.status.busy":"2023-06-11T09:03:32.074442Z","iopub.execute_input":"2023-06-11T09:03:32.075033Z","iopub.status.idle":"2023-06-11T09:03:32.087567Z","shell.execute_reply.started":"2023-06-11T09:03:32.074977Z","shell.execute_reply":"2023-06-11T09:03:32.086374Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# timg = cv2.imread(img_path)\n# H, W = timg.shape[2:]\n","metadata":{"execution":{"iopub.status.busy":"2023-06-11T09:03:32.089231Z","iopub.execute_input":"2023-06-11T09:03:32.089606Z","iopub.status.idle":"2023-06-11T09:03:32.096967Z","shell.execute_reply.started":"2023-06-11T09:03:32.089569Z","shell.execute_reply":"2023-06-11T09:03:32.095893Z"},"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    \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    HARDNET = KeyNetAffNetHardNet(num_feats, upright, device,).to(device)\n    HARDNET.eval()\n    \n    if LOCAL_FEATURE == 'KeyNetAffNetHardNet':   \n        if COMPOSITE_DESCRIPTOR:\n            features_detectors = {\n#                 \"DISK\":disk,\n                \"KAHN\":HARDNET,\n                \"HYNET\":HYNET,\n    #             \"sift\":KF.SIFTFeature(num_features=num_feats)\n            }\n        else:\n            features_detectors = {DESCRIPTOR:KeyNetAffNetHardNet(num_feats, upright, device).to(device).eval()}\n    else:\n        if COMPOSITE_DESCRIPTOR:\n            features_detectors = {\n                \"KAHN\":HARDNET,\n                \"HYNET\":HYNET,\n            }\n        else:\n            features_detectors = {DESCRIPTOR:KeyNetAffNetHardNet(1000, 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            if LOCAL_FEATURE==\"SUPERPOINT\":\n                with torch.inference_mode():\n                    #maybe try using cv2\n                    timg = load_torch_image(img_path, device=device)\n                    H, W = timg.shape[2:]\n                    if resize_small_edge_to:\n                        timg = K.geometry.resize(timg, resize_small_edge_to, antialias=True)\n                    h, w = timg_resized.shape[2:]\n                    timg = K.color.rgb_to_grayscale(timg)\n                    kpts,descs,_ = frontend.run(timg.detach().cpu().numpy()[0][0])\n                    idxs = kpts[2,:]>0.1\n                    kpts = kpts[:2,t].transpose().detach().cpu().numpy()\n                    descs = descs[:,t].transpose().detach().cpu().numpy()\n                    lafs = KF.laf_from_center_scale_ori(torch.tensor(kpts).unsqueeze(0))\n                    lafs = lafs.detach().cpu().numpy()\n            elif LOCAL_FEATURE in ['KeyNetAffNetHardNet','COMBINATION']:\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                    h, w = timg_resized.shape[2:]\n\n                    if LOCAL_FEATURE in ['KeyNetAffNetHardNet','COMBINATION']:\n                        begin = True\n                        \n                        try:\n                            for feature_name,feature in features_detectors.items():\n                                if feature_name==\"DISK\":\n                                    feats_ = feature(timg_resized, num_feats, pad_if_not_divisible=True)[0]\n                                    _kps, _descs = feats_.keypoints, torch.unsqueeze(feats_.descriptors,0)\n                                    _lafs = KF.laf_from_center_scale_ori(_kps[None], torch.ones(1, len(_kps), 1, 1, device=device))\n                                else:\n                                    _lafs, _, _descs = feature(K.color.rgb_to_grayscale(timg_resized))\n                                if begin:\n                                    lafs,descs = _lafs, _descs\n                                    begin = False\n                                else:\n\n                                    lafs,descs = torch.cat((lafs,_lafs),1),torch.cat((descs,_descs),1)\n                        except Exception as e:\n                            print(\"ERORORORORORORORO\",e)                                    \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                    \n            if LOCAL_FEATURE == 'CROP_TECH':\n                try:\n                    timg = load_torch_image(img_path, device=device)\n                    H, W = timg.shape[2:]\n                    lafs,descs,kpts = get_feats(\"DISK\",disk,timg,1000,None)\n                    data = kpts.astype(int)\n                    data = np.unique(data,axis=0) \n                    boxes = get_crop(data,H,W)\n                    boxes = [box for box in boxes if box[1]>3]\n                    for box in boxes:\n                        x1,y1,x2,y2 = box[0]\n                        crop_timg = timg[:,:,x1:x2,y1:y2]\n                        for feature_name,feature in features_detectors.items():\n                            for num_feats,resize_to in product([1000],[None]):\n                                _lafs,_descs,_kpts = get_feats(feature_name,feature,crop_timg,num_feats,resize_to)\n                                _lafs[:,:,0,2] +=x1\n                                _lafs[:,:,1,2] +=y1\n                                _kpts[:,0]+=x1\n                                _kpts[:,1]+=y1\n\n                                lafs,descs,kpts = np.concatenate((lafs,_lafs),1),np.concatenate((descs,_descs),0),np.concatenate((kpts,_kpts),0)\n                except Exception as e:\n                    print(\"ERRROR\",e)\n                f_laf[key] = lafs\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\n\ndef match_loftr1(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    print(\"f1\")\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    print(\"f2\")\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    print(\"f3\")\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    print(\"f4\")\n    with h5py.File(f'{feature_dir}/keypoints.h5', mode='a') as f_kp:\n        print(\"k\")\n        for k, kpts1 in unique_kpts.items():\n            print(\"k1\")\n            f_kp[k] = kpts1\n            print(\"k2\")\n    print(\"f5\")\n    with h5py.File(f'{feature_dir}/matches.h5', mode='a') 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\n\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","metadata":{"execution":{"iopub.status.busy":"2023-06-11T09:03:32.098922Z","iopub.execute_input":"2023-06-11T09:03:32.099298Z","iopub.status.idle":"2023-06-11T09:03:32.208274Z","shell.execute_reply.started":"2023-06-11T09:03:32.099269Z","shell.execute_reply":"2023-06-11T09:03:32.207059Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def main_func():\n    gc.collect()\n    datasets = []\n    for dataset in data_dict:\n        datasets.append(dataset)\n\n    for dataset in datasets:\n        print(dataset)\n        if dataset not in out_results:\n            out_results[dataset] = {}\n        for scene in data_dict[dataset]:\n            print(scene)\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 = f'{src}/{SET}/{dataset}/{scene}/images'\n            if not os.path.exists(img_dir):\n                continue\n            # Wrap the meaty part in a try-except block.\n            try:\n                out_results[dataset][scene] = {}\n                img_fnames = [f'{src}/{SET}/{x}' for x in data_dict[dataset][scene]]\n                print (f\"Got {len(img_fnames)} images\")\n                feature_dir = f'featureout/{dataset}_{scene}'\n                if not os.path.isdir(feature_dir):\n                    os.makedirs(feature_dir, exist_ok=True)\n                t=time()\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                \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 in ['KeyNetAffNetHardNet','CROP_TECH']:\n                    detect_features(img_fnames, \n                                NUM_FEATS,\n                                feature_dir=feature_dir,\n                                upright=UPRIGHT,\n                                device=device,\n                                resize_small_edge_to=RESIZE_SMALL_EDGE_TO\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,matching_alg='smnn')\n                elif LOCAL_FEATURE=='COMBINATION':\n                    print(\"bleh\")\n                    detect_features(img_fnames, \n                                NUM_FEATS,\n                                feature_dir=feature_dir,\n                                upright=UPRIGHT,\n                                device=device,\n                                resize_small_edge_to=RESIZE_SMALL_EDGE_TO\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,matching_alg='smnn')\n                    print(\"buh\")\n                    try:\n                        match_loftr1(img_fnames, index_pairs, feature_dir=feature_dir, device=device, resize_to_=RESIZE_SMALL_EDGE_TO)\n                    except Exception as e:\n                        print(e)\n                    print(\"bruh\")\n                else:\n                    match_loftr(img_fnames, index_pairs, feature_dir=feature_dir, device=device, resize_to_=(800, 1000))\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,verbose=VERBOSE)\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                #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                if isinstance(maps, dict):\n                    for idx1, rec in maps.items():\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                    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(data_dict[dataset][scene])} images')\n                create_submission(out_results, data_dict)\n                gc.collect()\n            except:\n                pass","metadata":{"execution":{"iopub.status.busy":"2023-06-11T09:03:32.209985Z","iopub.execute_input":"2023-06-11T09:03:32.210764Z","iopub.status.idle":"2023-06-11T09:03:32.240060Z","shell.execute_reply.started":"2023-06-11T09:03:32.210709Z","shell.execute_reply":"2023-06-11T09:03:32.239066Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def assert_func(s:str)->bool:\n    s = s.lower()\n    if \"n\" in s or \";;\" in s:\n        return False\n    return True","metadata":{"execution":{"iopub.status.busy":"2023-06-11T09:03:32.243610Z","iopub.execute_input":"2023-06-11T09:03:32.243968Z","iopub.status.idle":"2023-06-11T09:03:32.254418Z","shell.execute_reply.started":"2023-06-11T09:03:32.243932Z","shell.execute_reply":"2023-06-11T09:03:32.251072Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Function to create a submission file.\ndef create_submission(data_dict):\n    \n    main_func()\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 \"full\" not in image:\n                        try:\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\n                            s1 = arr_to_str(R)\n                            s2 = arr_to_str(T)\n                            assert assert_func(s1) and assert_func(s2)\n\n                        except:\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')","metadata":{"execution":{"iopub.status.busy":"2023-06-11T09:03:32.256697Z","iopub.execute_input":"2023-06-11T09:03:32.257088Z","iopub.status.idle":"2023-06-11T09:03:32.269584Z","shell.execute_reply.started":"2023-06-11T09:03:32.257051Z","shell.execute_reply":"2023-06-11T09:03:32.268713Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"out_results = {}\ntimings = {\"shortlisting\":[],\n               \"feature_detection\": [],\n               \"feature_matching\":[],\n               \"RANSAC\": [],\n               \"Reconstruction\": []}\nif SUBMIT:\n    create_submission(data_dict)","metadata":{"execution":{"iopub.status.busy":"2023-06-11T09:03:32.272669Z","iopub.execute_input":"2023-06-11T09:03:32.273724Z","iopub.status.idle":"2023-06-11T09:03:32.283445Z","shell.execute_reply.started":"2023-06-11T09:03:32.273682Z","shell.execute_reply":"2023-06-11T09:03:32.282393Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not SUBMIT:\n# Get data from csv.\n    data_dict = {}\n    if SUBMIT:\n        sample_submission_file = f'{src}/sample_submission.csv'\n    else:\n        sample_submission_file = f'{src}/train/train_labels.csv'\n\n    sample_df = pd.read_csv(sample_submission_file)\n\n\n    for i,j in sample_df.iterrows():\n        dataset = j[\"dataset\"]\n        scene = j[\"scene\"]\n\n        data_dict.setdefault(dataset,{}).setdefault(scene,[]).append(j[\"image_path\"])\n\n    if False:\n        for dataset,value in data_dict.items():\n            for scene,image_dir in value.items():\n                scene_dir = f\"{src}/{SET}/{dataset}/{scene}\"\n                if \"images_full\" in os.listdir(scene_dir):\n\n                    image_dir = f\"{scene_dir}/images_full\"\n                    for image in os.listdir(image_dir):\n                        data_dict[dataset][scene].append(os.path.join(image_dir,image))\n\n    if not SUBMIT:\n        a = data_dict[\"heritage\"].pop(\"dioscuri\")\n#         b = data_dict.pop(\"haiper\")\n#         b = data_dict.pop(\"heritage\")\n    for dataset in data_dict:\n        for scene in data_dict[dataset]:\n            print(f'{dataset} / {scene} -> {len(data_dict[dataset][scene])} images')\n    \n        \n    rotation_thresholds_degrees_dict = {\n        **{('haiper', scene): np.linspace(0.2, 10, 10) for scene in ['bike', 'chairs', 'fountain']},\n        **{('heritage', scene): np.linspace(0.2, 10, 10) for scene in ['cyprus', 'dioscuri']},\n        **{('heritage', 'wall'): np.linspace(0.2, 10, 10)},\n        **{('urban', 'kyiv-puppet-theater'): np.linspace(0.2, 10, 10)},\n    }\n\n    translation_thresholds_meters_dict = {\n        **{('haiper', scene): np.geomspace(0.05, 1, 10) for scene in ['bike', 'chairs', 'fountain']},\n        **{('heritage', scene): np.geomspace(0.05, 1, 10) for scene in ['cyprus', 'dioscuri']},\n        **{('heritage', 'wall'): np.geomspace(0.05, 1, 10)},\n        **{('urban', 'kyiv-puppet-theater'): np.geomspace(0.5, 1, 10)},\n    }\n    \n    df = pd.read_csv(\"/kaggle/input/image-matching-challenge-2023/train/train_labels.csv\")\n    df = df[['image_path','dataset', 'scene',  'rotation_matrix',\n           'translation_vector']]\n    df = df[df[\"scene\"]!=\"dioscuri\"]\n#     df = df[df[\"dataset\"]!=\"heritage\"]\n#     df = df[df[\"dataset\"]!=\"haiper\"]\n    df.to_csv(\"train.csv\",index=False)\n    \n    ","metadata":{"execution":{"iopub.status.busy":"2023-06-11T09:03:32.285066Z","iopub.execute_input":"2023-06-11T09:03:32.285656Z","iopub.status.idle":"2023-06-11T09:03:32.345058Z","shell.execute_reply.started":"2023-06-11T09:03:32.285590Z","shell.execute_reply":"2023-06-11T09:03:32.343984Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not SUBMIT:\n    expr = {}\n    import itertools\n    descriptors = [\"HARDNET\"]\n    feats = [1000]\n    sizes = [800]\n    composite_status = [True]\n    upright_ = [False]\n    LOCAL_FEATURE = 'SUPERPOINT'\n\n    for DESCRIPTOR,NUM_FEATS,RESIZE_SMALL_EDGE_TO,COMPOSITE_DESCRIPTOR,UPRIGHT in itertools.product(*[descriptors,feats,sizes,composite_status,upright_]):\n        out_results = {}\n        timings = {\"shortlisting\":[],\n                       \"feature_detection\": [],\n                       \"feature_matching\":[],\n                       \"RANSAC\": [],\n                       \"Reconstruction\": []}\n        t = time()\n        create_submission(data_dict)\n        t = time()-t\n        expr[(DESCRIPTOR,NUM_FEATS,RESIZE_SMALL_EDGE_TO,f\"COMPOSITE_DESCRIPTOR_{str(COMPOSITE_DESCRIPTOR)}\",f\"UPRIGHT_{str(UPRIGHT)}\")] = (eval_submission(submission_csv_path='submission.csv',\n               ground_truth_csv_path='train.csv',\n               rotation_thresholds_degrees_dict=rotation_thresholds_degrees_dict,\n               translation_thresholds_meters_dict=translation_thresholds_meters_dict,\n               verbose=True),t)\n        \n    print(expr)","metadata":{"execution":{"iopub.status.busy":"2023-06-11T09:03:32.346686Z","iopub.execute_input":"2023-06-11T09:03:32.347152Z"},"collapsed":true,"jupyter":{"outputs_hidden":true},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}