{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.10","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":5847506,"sourceType":"datasetVersion","datasetId":3267335},{"sourceId":7884485,"sourceType":"datasetVersion","datasetId":4628051},{"sourceId":8389141,"sourceType":"datasetVersion","datasetId":4989748},{"sourceId":3736,"sourceType":"modelInstanceVersion","modelInstanceId":2663}],"dockerImageVersionId":30498,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## 1. Setup","metadata":{}},{"cell_type":"code","source":"import sys\nimport os\nimport h5py\nimport gc\nfrom glob import glob\nfrom pathlib import Path\nfrom functools import lru_cache\nfrom collections import defaultdict\nfrom fastprogress import progress_bar\nfrom copy import deepcopy\nfrom tqdm import tqdm\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image, ExifTags\nimport cv2\nimport sqlite3\nfrom sklearn.preprocessing import MinMaxScaler\nfrom sklearn.cluster import DBSCAN\nimport torch\nimport torch.nn.functional as F\nimport timm\nfrom timm.data import resolve_data_config\nfrom timm.data.transforms_factory import create_transform\nimport kornia\nimport pycolmap","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-05-12T07:01:13.077683Z","iopub.execute_input":"2024-05-12T07:01:13.078072Z","iopub.status.idle":"2024-05-12T07:01:26.740998Z","shell.execute_reply.started":"2024-05-12T07:01:13.078043Z","shell.execute_reply":"2024-05-12T07:01:26.739446Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# competition_dataset = Path('/kaggle/input/image-matching-challenge-2023')\n# external_dataset = Path('/kaggle/input/image-matching-challenge-2023-dataset')\ncompetition_dataset = Path('/kaggle/input/image-matching-challenge-2024')","metadata":{"execution":{"iopub.status.busy":"2024-05-12T07:01:26.748117Z","iopub.execute_input":"2024-05-12T07:01:26.752262Z","iopub.status.idle":"2024-05-12T07:01:26.761240Z","shell.execute_reply.started":"2024-05-12T07:01:26.752210Z","shell.execute_reply":"2024-05-12T07:01:26.759864Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# sys.path.append(str(external_dataset / 'packages' / 'SuperGluePretrainedNetwork'))\nsys.path.append(str('/kaggle/input/magicleap-supergluepretrainednetwork'))\n# import models\nfrom models.matching import Matching","metadata":{"execution":{"iopub.status.busy":"2024-05-12T07:01:26.767382Z","iopub.execute_input":"2024-05-12T07:01:26.770917Z","iopub.status.idle":"2024-05-12T07:01:26.806750Z","shell.execute_reply.started":"2024-05-12T07:01:26.770869Z","shell.execute_reply":"2024-05-12T07:01:26.805743Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df = pd.read_csv(competition_dataset / 'sample_submission.csv')\n\nif df.shape[0] != 3:\n # Enable submission mode and disable verbose\n submission = True\n verbose = False\n# train_or_test_directory = competition_dataset / 'test'\n train_or_test_directory = competition_dataset\nelse:\n # Disable submission mode and enable verbose\n submission = False\n verbose = True\n train_or_test_directory = competition_dataset / 'train'\n\nif submission is False:\n df = pd.read_csv(competition_dataset / 'train' / 'train_labels.csv')\n # Use only bike scene from the haiper dataset if it's not submission mode\n df = df.loc[df['scene'] == 'bike', :]\n \nprint(f'Dataset Shape: {df.shape}')","metadata":{"execution":{"iopub.status.busy":"2024-05-12T07:01:26.809159Z","iopub.execute_input":"2024-05-12T07:01:26.809842Z","iopub.status.idle":"2024-05-12T07:01:26.825455Z","shell.execute_reply.started":"2024-05-12T07:01:26.809809Z","shell.execute_reply":"2024-05-12T07:01:26.824592Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def load_angle_detection_model():\n \n \"\"\"\n Load angle detection model and processor\n \n Returns\n -------\n model: keras.engine.functional.Functional\n Angle detection model\n \n processor: transformers.ViTFeatureExtractor\n Image processor\n \"\"\"\n \n vit = TFViTModel.from_pretrained(external_dataset / 'models' / 'vit-base-patch16-224')\n \n input_layer = Input(shape=(3, 224, 224))\n x = vit(input_layer)\n y = Dense(1, activation='linear')(x[-1])\n \n model = Model(input_layer, y)\n model.load_weights(external_dataset / 'models' / 'deep-image-orientation-angle-detection' / 'model-vit-ang-loss.h5')\n \n processor = ViTFeatureExtractor.from_pretrained(external_dataset / 'models' / 'vit-base-patch16-224')\n \n return model, processor\n\n\ndef detect_angle(image, model, processor):\n \n \"\"\"\n Detect angle of a given image\n \n Parameters\n ----------\n image: numpy.ndarray of shape (height, width, channel)\n Image array\n \n model: keras.engine.functional.Functional\n Angle detection model\n \n processor: transformers.ViTFeatureExtractor\n Image processor\n \n Returns\n -------\n outputs: float\n \"\"\"\n \n inputs = processor(images=[image], return_tensors='np')['pixel_values']\n outputs = model.predict(inputs, verbose=None)[0][0]\n \n return outputs\n","metadata":{"execution":{"iopub.status.busy":"2024-05-12T07:01:26.827680Z","iopub.execute_input":"2024-05-12T07:01:26.827928Z","iopub.status.idle":"2024-05-12T07:01:26.835595Z","shell.execute_reply.started":"2024-05-12T07:01:26.827906Z","shell.execute_reply":"2024-05-12T07:01:26.834647Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#angle_detection_model, angle_detection_processor = load_angle_detection_model()\n\nfor idx, row in tqdm(df.iterrows(), total=df.shape[0]):\n\n image_path = train_or_test_directory / row['image_path']\n image = np.array(Image.open(image_path))\n \n # Extract image dimensions and memory usage\n df.loc[idx, 'image_height'] = image.shape[0]\n df.loc[idx, 'image_width'] = image.shape[1]\n df.loc[idx, 'memory_usage'] = image.nbytes\n \n # Extract image orientation\n '''\n try:\n with tf.device('/cpu:0'):\n angle = detect_angle(\n image=image,\n model=angle_detection_model,\n processor=angle_detection_processor\n )\n df.loc[idx, 'angle'] = angle\n except:\n df.loc[idx, 'angle'] = 0\n ''' \n\n\ndf['image_id'] = df['image_path'].apply(lambda x: str(x).split('/')[-1])\ndf['memory_usage'] /= (1024 ** 2)\n\ndf['image_height'] = df['image_height'].astype(np.uint16)\ndf['image_width'] = df['image_width'].astype(np.uint16)\ndf['memory_usage'] = df['memory_usage'].astype(np.float32)\n'''\ndf['angle'] = df['angle'].astype(np.float32)\nangles = df[['image_path', 'angle']].set_index('image_path').to_dict()['angle']\ndel angle_detection_model, angle_detection_processor\ngc.collect()\ntorch.cuda.empty_cache()\nK.clear_session()\n'''","metadata":{"execution":{"iopub.status.busy":"2024-05-12T07:01:26.836981Z","iopub.execute_input":"2024-05-12T07:01:26.837414Z","iopub.status.idle":"2024-05-12T07:01:29.160233Z","shell.execute_reply.started":"2024-05-12T07:01:26.837379Z","shell.execute_reply":"2024-05-12T07:01:29.159341Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df","metadata":{"execution":{"iopub.status.busy":"2024-05-12T07:01:29.161530Z","iopub.execute_input":"2024-05-12T07:01:29.161807Z","iopub.status.idle":"2024-05-12T07:01:29.190242Z","shell.execute_reply.started":"2024-05-12T07:01:29.161783Z","shell.execute_reply":"2024-05-12T07:01:29.189359Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 2. Image Utilities","metadata":{}},{"cell_type":"code","source":"def get_global_desc(image_paths, model, device):\n \n config = resolve_data_config({}, model=model)\n transform = create_transform(**config)\n \n global_descs = []\n for i, img_fname_full in tqdm(enumerate(image_paths),total=len(image_paths)):\n key = os.path.splitext(os.path.basename(img_fname_full))[0]\n img = Image.open(img_fname_full).convert('RGB')\n img = transform(img).unsqueeze(0).to(device)\n with torch.no_grad():\n with torch.autocast(device_type=device.type, dtype=torch.float16):\n desc = model.forward_features(img.to(device)).mean(dim=(-1, 2))\n desc = desc.view(1, -1).detach().cpu().float()\n desc = F.normalize(desc, dim=1, p=2)\n global_descs.append(desc)\n global_descs = torch.cat(global_descs, dim=0)\n \n return global_descs\n\n\ndef create_image_pairs_exhaustive(image_paths):\n\n \"\"\"\n Create all possible image pairs from given list of image paths\n\n Parameters\n ----------\n image_paths: list of shape (n_images)\n List of image paths\n\n Returns\n -------\n image_pair_indices: list of shape (n_image_pairs)\n List of tuples of image pair indices\n \"\"\"\n\n image_pair_indices = []\n\n for i in range(len(image_paths)):\n for j in range(i + 1, len(image_paths)):\n image_pair_indices.append((i, j))\n\n return image_pair_indices\n\n\ndef create_image_pairs(image_paths, sim_th=0.6, min_pairs=20, exhaustive_if_less=20, device=torch.device('cuda')):\n \n n_images = len(image_paths)\n\n if n_images <= exhaustive_if_less:\n \n image_pair_indices = []\n # Create all possible image pairs from given list of image paths\n for i in range(len(image_paths)):\n for j in range(i + 1, len(image_paths)):\n image_pair_indices.append((i, j))\n\n return image_pair_indices\n \n else:\n\n model = timm.create_model(\n 'tf_efficientnet_b7',\n checkpoint_path='/kaggle/input/tf-efficientnet/pytorch/tf-efficientnet-b7/1/tf_efficientnet_b7_ra-6c08e654.pth'\n )\n model = model.eval().to(image_matching_device)\n descs = get_global_desc(image_paths, model, device=device)\n dm = torch.cdist(descs, descs, p=2).numpy()\n mask = dm <= sim_th\n total = 0\n image_pair_indices = []\n ar = np.arange(n_images)\n already_there_set = []\n for st_idx in range(n_images - 1):\n mask_idx = mask[st_idx]\n to_match = ar[mask_idx]\n if len(to_match) < min_pairs:\n to_match = np.argsort(dm[st_idx])[:min_pairs] \n for idx in to_match:\n if st_idx == idx:\n continue\n if dm[st_idx, idx] < 1000:\n image_pair_indices.append(tuple(sorted((st_idx, idx.item()))))\n total += 1\n \n image_pair_indices = sorted(list(set(image_pair_indices)))\n \n return image_pair_indices\n\n\ndef resize_with_aspect_ratio(image, longest_edge):\n\n \"\"\"\n Resize image while preserving its aspect ratio\n\n Parameters\n ----------\n image: numpy.ndarray of shape (height, width, 3)\n Image array\n\n longest_edge: int\n Desired number of pixels on the longest edge\n\n Returns\n -------\n image: numpy.ndarray of shape (resized_height, resized_width, 3)\n Resized image array\n \"\"\"\n\n height, width = image.shape[:2]\n scale = longest_edge / max(height, width)\n image = cv2.resize(image, dsize=(int(width * scale), int(height * scale)), interpolation=cv2.INTER_LANCZOS4)\n\n return image\n\n\ndef get_image_tensor(image_path_or_array, resize, resize_shape, resize_longest_edge, scale, grayscale):\n\n \"\"\"\n Load image and preprocess it\n\n Parameters\n ----------\n image_path_or_array: str or numpy.ndarray of shape (height, width, 3)\n Image path or image array\n\n resize: bool\n Whether to resize the image or not\n\n resize_shape: tuple or int\n Tuple of image height and width or number of pixels for both height and width\n\n resize_longest_edge: bool\n Whether to resize the longest edge or not\n\n scale: bool\n Whether to scale image pixel values by max 8-bit pixel value or not\n\n grayscale: bool\n Whether to convert RGB image to grayscale or not\n\n Returns\n -------\n image: torch.Tensor of shape (1, 1 or 3, height, width)\n Image tensor\n \"\"\"\n\n if isinstance(image_path_or_array, Path) or isinstance(image_path_or_array, str):\n # Read image from the given path if image_path_or_array is a path-like string\n image = cv2.imread(str(image_path_or_array))\n image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n else:\n image = image_path_or_array\n\n if resize:\n if resize_longest_edge:\n image = resize_with_aspect_ratio(image=image, longest_edge=resize_shape)\n else:\n resize_shape = (resize_shape, resize_shape) if isinstance(resize_shape, int) else resize_shape\n image = cv2.resize(image, resize_shape, interpolation=cv2.INTER_LANCZOS4)\n\n if scale:\n image = image / 255.\n\n image = kornia.image_to_tensor(image, False).float()\n if grayscale:\n image = kornia.color.rgb_to_grayscale(image)\n\n return image\n\n\ndef crop(image, keypoints, perc_points=0.85, pad=5):\n\n norm_keypoints = MinMaxScaler().fit_transform(keypoints)\n total = len(keypoints)\n best_dist = 1\n best_clusters = None\n best_asm = None\n for eps in [0.01, 0.025, 0.05, 0.1, 0.2]:\n clusters = DBSCAN(eps=eps).fit_predict(norm_keypoints)\n counts = pd.Series(clusters).value_counts().sort_values(ascending=False)\n counts = counts[counts.index > -1]\n if len(counts) == 0:\n continue\n\n cumsums = np.cumsum(counts.values) / total\n dists = np.abs(cumsums - perc_points)\n best_ix = np.argmin(dists)\n\n if dists[best_ix] < best_dist:\n best_dist = dists[best_ix]\n best_clusters = list(counts.head(best_ix + 1).index)\n best_asm = clusters\n\n mask = np.isin(best_asm, best_clusters)\n\n miny = int(np.min(keypoints[mask][:, 1]))\n miny = max(miny - pad, 0)\n\n maxy = int(np.max(keypoints[mask][:, 1]))\n maxy = min(maxy + pad, image.shape[0])\n\n minx = int(np.min(keypoints[mask][:, 0]))\n minx = max(minx - pad, 0)\n\n maxx = int(np.max(keypoints[mask][:, 0]))\n maxx = min(maxx + pad, image.shape[1])\n\n image_cropped = image[miny:maxy + 1, minx:maxx + 1, :]\n keypoints_cropped = np.copy(keypoints)\n keypoints_cropped[:, 0] -= minx\n keypoints_cropped[:, 1] -= miny\n\n keypoints_cropped = keypoints_cropped[(keypoints_cropped[:, 0] > 0) & (keypoints_cropped[:, 1] > 0) & (keypoints_cropped[:, 0] < image_cropped.shape[1]) & (keypoints_cropped[:, 1] < image_cropped.shape[0])]\n\n return image_cropped, keypoints_cropped, minx, miny\n\n\ndef rotate_image(img, angle):\n height, width = img.shape[:2]\n matrix = cv2.getRotationMatrix2D((width / 2 - 0.5, height / 2 - 0.5), angle, 1.0)\n matrix = np.vstack((matrix, [[0, 0, 1]]))\n return cv2.warpPerspective(img, matrix, (width, height))\n\n\ndef rotate_coordinates(img, angle, keypoints):\n height, width = img.shape[:2]\n inv_matrix = cv2.getRotationMatrix2D((width / 2 - 0.5, height / 2 - 0.5), angle, 1.0)\n inv_matrix = np.vstack((inv_matrix, [[0, 0, 1]]))\n return cv2.perspectiveTransform(keypoints[None, :, :], inv_matrix)[0]\n","metadata":{"execution":{"iopub.status.busy":"2024-05-12T07:01:29.191723Z","iopub.execute_input":"2024-05-12T07:01:29.192210Z","iopub.status.idle":"2024-05-12T07:01:29.227278Z","shell.execute_reply.started":"2024-05-12T07:01:29.192163Z","shell.execute_reply":"2024-05-12T07:01:29.226408Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 3. Camera Utilities","metadata":{}},{"cell_type":"code","source":"def get_focal_length(image_path):\n\n \"\"\"\n Get focal length from EXIF or calculate it using prior\n\n Parameters\n ----------\n image_path: str\n Image path\n\n Returns\n -------\n focal_length: float\n Focal length extracted from EXIF or calculated using prior\n \"\"\"\n\n image = Image.open(image_path)\n image_longest_edge = max(image.size)\n\n focal_length = None\n exif = image.getexif()\n \n try:\n if image._getexif() is not None:\n exif = {\n ExifTags.TAGS[k]: v\n for k, v in image._getexif().items()\n if k in ExifTags.TAGS\n }\n else:\n exif = None\n except:\n exif = None\n \n if exif is not None:\n\n focal_length_35mm = None\n\n for tag, value in exif.items():\n if tag == 'FocalLengthIn35mmFilm':\n focal_length_35mm = float(value)\n \n if focal_length_35mm is not None:\n focal_length = focal_length_35mm / 35. * image_longest_edge\n\n if focal_length is None:\n prior_focal_length = 1.2\n focal_length = prior_focal_length * image_longest_edge\n\n return focal_length\n","metadata":{"execution":{"iopub.status.busy":"2024-05-12T07:01:29.228394Z","iopub.execute_input":"2024-05-12T07:01:29.228666Z","iopub.status.idle":"2024-05-12T07:01:29.240530Z","shell.execute_reply.started":"2024-05-12T07:01:29.228642Z","shell.execute_reply":"2024-05-12T07:01:29.239824Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 4. Database Utilities","metadata":{}},{"cell_type":"code","source":"MAX_IMAGE_ID = 2 ** 31 - 1\n\nCREATE_CAMERAS_TABLE = \"\"\"CREATE TABLE IF NOT EXISTS cameras (\n camera_id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,\n model INTEGER NOT NULL,\n width INTEGER NOT NULL,\n height INTEGER NOT NULL,\n params BLOB,\n prior_focal_length INTEGER NOT NULL)\"\"\"\n\nCREATE_DESCRIPTORS_TABLE = \"\"\"CREATE TABLE IF NOT EXISTS descriptors (\n image_id INTEGER PRIMARY KEY NOT NULL,\n rows INTEGER NOT NULL,\n cols INTEGER NOT NULL,\n data BLOB,\n FOREIGN KEY(image_id) REFERENCES images(image_id) ON DELETE CASCADE)\"\"\"\n\nCREATE_IMAGES_TABLE = \"\"\"CREATE TABLE IF NOT EXISTS images (\n image_id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,\n name TEXT NOT NULL UNIQUE,\n camera_id INTEGER NOT NULL,\n prior_qw REAL,\n prior_qx REAL,\n prior_qy REAL,\n prior_qz REAL,\n prior_tx REAL,\n prior_ty REAL,\n prior_tz REAL,\n CONSTRAINT image_id_check CHECK(image_id >= 0 and image_id < {}),\n FOREIGN KEY(camera_id) REFERENCES cameras(camera_id))\n\"\"\".format(MAX_IMAGE_ID)\n\nCREATE_TWO_VIEW_GEOMETRIES_TABLE = \"\"\"\nCREATE TABLE IF NOT EXISTS two_view_geometries (\n pair_id INTEGER PRIMARY KEY NOT NULL,\n rows INTEGER NOT NULL,\n cols INTEGER NOT NULL,\n data BLOB,\n config INTEGER NOT NULL,\n F BLOB,\n E BLOB,\n H BLOB)\n\"\"\"\n\nCREATE_KEYPOINTS_TABLE = \"\"\"CREATE TABLE IF NOT EXISTS keypoints (\n image_id INTEGER PRIMARY KEY NOT NULL,\n rows INTEGER NOT NULL,\n cols INTEGER NOT NULL,\n data BLOB,\n FOREIGN KEY(image_id) REFERENCES images(image_id) ON DELETE CASCADE)\n\"\"\"\n\nCREATE_MATCHES_TABLE = \"\"\"CREATE TABLE IF NOT EXISTS matches (\n pair_id INTEGER PRIMARY KEY NOT NULL,\n rows INTEGER NOT NULL,\n cols INTEGER NOT NULL,\n data BLOB)\"\"\"\n\nCREATE_NAME_INDEX = \\\n \"CREATE UNIQUE INDEX IF NOT EXISTS index_name ON images(name)\"\n\nCREATE_ALL = \"; \".join([\n CREATE_CAMERAS_TABLE,\n CREATE_IMAGES_TABLE,\n CREATE_KEYPOINTS_TABLE,\n CREATE_DESCRIPTORS_TABLE,\n CREATE_MATCHES_TABLE,\n CREATE_TWO_VIEW_GEOMETRIES_TABLE,\n CREATE_NAME_INDEX\n])\n\n\ndef image_ids_to_pair_id(image_id1, image_id2):\n if image_id1 > image_id2:\n image_id1, image_id2 = image_id2, image_id1\n return image_id1 * MAX_IMAGE_ID + image_id2\n\n\ndef pair_id_to_image_ids(pair_id):\n image_id2 = pair_id % MAX_IMAGE_ID\n image_id1 = (pair_id - image_id2) / MAX_IMAGE_ID\n return image_id1, image_id2\n\n\ndef array_to_blob(array):\n return array.tostring()\n\n\ndef blob_to_array(blob, dtype, shape=(-1,)):\n return np.fromstring(blob, dtype=dtype).reshape(*shape)\n\n\nclass COLMAPDatabase(sqlite3.Connection):\n\n @staticmethod\n def connect(database_path):\n return sqlite3.connect(database_path, factory=COLMAPDatabase)\n\n\n def __init__(self, *args, **kwargs):\n super(COLMAPDatabase, self).__init__(*args, **kwargs)\n\n self.create_tables = lambda: self.executescript(CREATE_ALL)\n self.create_cameras_table = \\\n lambda: self.executescript(CREATE_CAMERAS_TABLE)\n self.create_descriptors_table = \\\n lambda: self.executescript(CREATE_DESCRIPTORS_TABLE)\n self.create_images_table = \\\n lambda: self.executescript(CREATE_IMAGES_TABLE)\n self.create_two_view_geometries_table = \\\n lambda: self.executescript(CREATE_TWO_VIEW_GEOMETRIES_TABLE)\n self.create_keypoints_table = \\\n lambda: self.executescript(CREATE_KEYPOINTS_TABLE)\n self.create_matches_table = \\\n lambda: self.executescript(CREATE_MATCHES_TABLE)\n self.create_name_index = lambda: self.executescript(CREATE_NAME_INDEX)\n\n def add_camera(self, model, width, height, params,\n prior_focal_length=False, camera_id=None):\n params = np.asarray(params, np.float64)\n cursor = self.execute(\n \"INSERT INTO cameras VALUES (?, ?, ?, ?, ?, ?)\",\n (camera_id, model, width, height, array_to_blob(params),\n prior_focal_length))\n return cursor.lastrowid\n\n def add_image(self, name, camera_id,\n prior_q=np.zeros(4), prior_t=np.zeros(3), image_id=None):\n cursor = self.execute(\n \"INSERT INTO images VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)\",\n (image_id, name, camera_id, prior_q[0], prior_q[1], prior_q[2],\n prior_q[3], prior_t[0], prior_t[1], prior_t[2]))\n return cursor.lastrowid\n\n def add_keypoints(self, image_id, keypoints):\n assert(len(keypoints.shape) == 2)\n assert(keypoints.shape[1] in [2, 4, 6])\n\n keypoints = np.asarray(keypoints, np.float32)\n self.execute(\n \"INSERT INTO keypoints VALUES (?, ?, ?, ?)\",\n (image_id,) + keypoints.shape + (array_to_blob(keypoints),))\n\n def add_descriptors(self, image_id, descriptors):\n descriptors = np.ascontiguousarray(descriptors, np.uint8)\n self.execute(\n \"INSERT INTO descriptors VALUES (?, ?, ?, ?)\",\n (image_id,) + descriptors.shape + (array_to_blob(descriptors),))\n\n def add_matches(self, image_id1, image_id2, matches):\n assert(len(matches.shape) == 2)\n assert(matches.shape[1] == 2)\n\n if image_id1 > image_id2:\n matches = matches[:,::-1]\n\n pair_id = image_ids_to_pair_id(image_id1, image_id2)\n matches = np.asarray(matches, np.uint32)\n self.execute(\n \"INSERT INTO matches VALUES (?, ?, ?, ?)\",\n (pair_id,) + matches.shape + (array_to_blob(matches),))\n\n def add_two_view_geometry(self, image_id1, image_id2, matches,\n F=np.eye(3), E=np.eye(3), H=np.eye(3), config=2):\n assert(len(matches.shape) == 2)\n assert(matches.shape[1] == 2)\n\n if image_id1 > image_id2:\n matches = matches[:,::-1]\n\n pair_id = image_ids_to_pair_id(image_id1, image_id2)\n matches = np.asarray(matches, np.uint32)\n F = np.asarray(F, dtype=np.float64)\n E = np.asarray(E, dtype=np.float64)\n H = np.asarray(H, dtype=np.float64)\n self.execute(\n \"INSERT INTO two_view_geometries VALUES (?, ?, ?, ?, ?, ?, ?, ?)\",\n (pair_id,) + matches.shape + (array_to_blob(matches), config,\n array_to_blob(F), array_to_blob(E), array_to_blob(H)))\n","metadata":{"execution":{"iopub.status.busy":"2024-05-12T07:01:29.244510Z","iopub.execute_input":"2024-05-12T07:01:29.244798Z","iopub.status.idle":"2024-05-12T07:01:29.269177Z","shell.execute_reply.started":"2024-05-12T07:01:29.244776Z","shell.execute_reply":"2024-05-12T07:01:29.268219Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def create_camera(db, image_path, camera_model):\n \n image = Image.open(image_path)\n width, height = image.size\n\n focal_length = get_focal_length(image_path)\n\n if camera_model == 'simple-pinhole':\n model = 0\n param_arr = np.array([focal_length, width / 2, height / 2])\n if camera_model == 'pinhole':\n model = 1\n param_arr = np.array([focal_length, focal, width / 2, height / 2])\n elif camera_model == 'simple-radial':\n model = 2\n param_arr = np.array([focal_length, width / 2, height / 2, 0.1])\n elif camera_model == 'opencv':\n model = 4\n param_arr = np.array([focal_length, focal_length, width / 2, height / 2, 0., 0., 0., 0.])\n \n return db.add_camera(model, width, height, param_arr)\n\n\ndef add_keypoints(db, h5_path, image_path, camera_model, single_camera=True):\n \n keypoint_f = h5py.File(os.path.join(h5_path, 'keypoints.h5'), 'r')\n camera_id = None\n fname_to_id = {}\n \n for filename in tqdm(list(keypoint_f.keys())):\n \n keypoints = keypoint_f[filename][()]\n\n path = os.path.join(image_path, filename)\n if not os.path.isfile(path):\n raise IOError(f'Invalid image path {path}')\n\n if camera_id is None or not single_camera:\n camera_id = create_camera(db, path, camera_model)\n image_id = db.add_image(filename, camera_id)\n fname_to_id[filename] = image_id\n\n db.add_keypoints(image_id, keypoints)\n\n return fname_to_id\n\n\ndef add_matches(db, h5_path, fname_to_id):\n \n match_file = h5py.File(os.path.join(h5_path, 'matches.h5'), 'r')\n \n added = set()\n n_keys = len(match_file.keys())\n n_total = (n_keys * (n_keys - 1)) // 2\n\n with tqdm(total=n_total) as pbar:\n for key_1 in match_file.keys():\n group = match_file[key_1]\n for key_2 in group.keys():\n id_1 = fname_to_id[key_1]\n id_2 = fname_to_id[key_2]\n\n pair_id = image_ids_to_pair_id(id_1, id_2)\n if pair_id in added:\n continue\n \n matches = group[key_2][()]\n db.add_matches(id_1, id_2, matches)\n added.add(pair_id)\n pbar.update(1)\n \n \ndef get_unique_idxs(A, dim=0):\n \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 \n return first_indices\n","metadata":{"execution":{"iopub.status.busy":"2024-05-12T07:01:29.270248Z","iopub.execute_input":"2024-05-12T07:01:29.270485Z","iopub.status.idle":"2024-05-12T07:01:29.287270Z","shell.execute_reply.started":"2024-05-12T07:01:29.270464Z","shell.execute_reply":"2024-05-12T07:01:29.286467Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 5. Evaluation","metadata":{}},{"cell_type":"code","source":"def array_to_string(a):\n\n \"\"\"\n Flatten given array and convert it to a string with semicolon delimiters\n\n Parameters\n ----------\n a: np.ndarray\n N-dimensional array\n\n Returns\n -------\n s: string\n String form of the given array\n \"\"\"\n\n s = ';'.join([str(x) for x in a.reshape(-1)])\n\n return s\n\n\ndef string_to_array(s):\n\n \"\"\"\n Convert semicolon delimited string to an array\n\n Parameters\n ----------\n s: string\n String form of the array\n\n Returns\n -------\n a: np.ndarray\n N-dimensional array\n \"\"\"\n\n a = np.array(s.split(';')).astype(np.float64)\n\n return a\n\n\ndef rotation_matrix_to_quaternion(rotation_matrix):\n\n \"\"\"\n Convert rotation matrix to quaternion\n\n Parameters\n ----------\n rotation_matrix: numpy.ndarray of shape (3, 3)\n Array of directions of the world-axes in camera coordinates\n\n Returns\n -------\n quaternion: numpy.ndarray of shape (4)\n Array of quaternion\n \"\"\"\n\n r00 = rotation_matrix[0, 0]\n r01 = rotation_matrix[0, 1]\n r02 = rotation_matrix[0, 2]\n r10 = rotation_matrix[1, 0]\n r11 = rotation_matrix[1, 1]\n r12 = rotation_matrix[1, 2]\n r20 = rotation_matrix[2, 0]\n r21 = rotation_matrix[2, 1]\n r22 = rotation_matrix[2, 2]\n\n k = np.array([\n [r00 - r11 - r22, 0.0, 0.0, 0.0],\n [r01 + r10, r11 - r00 - r22, 0.0, 0.0],\n [r02 + r20, r12 + r21, r22 - r00 - r11, 0.0],\n [r21 - r12, r02 - r20, r10 - r01, r00 + r11 + r22]\n ])\n k /= 3.0\n\n # Quaternion is the eigenvector of k that corresponds to the largest eigenvalue\n w, v = np.linalg.eigh(k)\n quaternion = v[[3, 0, 1, 2], np.argmax(w)]\n\n if quaternion[0] < 0:\n np.negative(quaternion, quaternion)\n\n return quaternion\n\n\ndef pose_difference(r1, t1, r2, t2):\n\n \"\"\"\n Calculate relative pose difference from given rotation matrices and translation vectors\n\n Parameters\n ----------\n r1: numpy.ndarray of shape (3, 3)\n First rotation matrix\n\n t1: numpy.ndarray of shape (3)\n First translation vector\n\n r2: numpy.ndarray of shape (3, 3)\n Second rotation matrix\n\n t2: numpy.ndarray of shape (3)\n Second translation vector\n\n Returns\n -------\n rotation_difference: float\n Rotation difference in terms of degrees from the first image\n\n translation_difference: float\n Translation difference in terms of meters from the first image\n \"\"\"\n\n rotation_difference = np.dot(r2, r1.T)\n translation_difference = t2 - np.dot(rotation_difference, t1)\n\n return rotation_difference, translation_difference\n\n\ndef rotation_and_translation_error(q_ground_truth, t_ground_truth, q_prediction, t_prediction, epsilon=1e-15):\n\n \"\"\"\n Calculate rotation and translation error\n\n Parameters\n ----------\n q_ground_truth: numpy.ndarray of shape (4)\n Array of quaternion derived from ground truth rotation matrix\n\n t_ground_truth: numpy.ndarray of shape (3)\n Array of ground truth translation vector\n\n q_prediction: numpy.ndarray of shape (4)\n Array of quaternion derived from estimated rotation matrix\n\n t_prediction: numpy.ndarray of shape (3)\n Array of estimated translation vector\n\n epsilon: float\n A small number for preventing zero division\n\n Returns\n -------\n rotation_error: float\n Rotation error in terms of degrees\n\n translation_error: float\n Translation error in terms of meters\n \"\"\"\n\n q_ground_truth_norm = q_ground_truth / (np.linalg.norm(q_ground_truth) + epsilon)\n q_prediction_norm = q_prediction / (np.linalg.norm(q_prediction) + epsilon)\n loss_q = np.maximum(epsilon, (1.0 - np.sum(q_prediction_norm * q_ground_truth_norm) ** 2))\n\n rotation_error = np.degrees(np.arccos(1 - (2 * loss_q)))\n\n scaling_factor = np.linalg.norm(t_ground_truth)\n t_prediction = scaling_factor * (t_prediction / (np.linalg.norm(t_prediction) + epsilon))\n translation_error = min(\n np.linalg.norm(t_ground_truth - t_prediction),\n np.linalg.norm(t_ground_truth + t_prediction)\n )\n\n return rotation_error, translation_error\n\n\ndef mean_average_accuracy(rotation_errors, translation_errors, rotation_error_thresholds, translation_error_thresholds):\n\n \"\"\"\n Calculate mean average accuracies over a set of thresholds for rotation and translation\n\n Parameters\n ----------\n rotation_errors: list of shape (n_pairs)\n List of rotation errors\n\n translation_errors: list of shape (n_pairs)\n List of translation errors\n\n rotation_error_thresholds: numpy.ndarray of shape (10)\n Array of rotation error thresholds\n\n translation_error_thresholds: numpy.ndarray of shape (10)\n Array of translation error thresholds\n\n Returns\n -------\n maa: float\n Mean average accuracy calculated on both rotation and translation errors\n\n rotation_maa: float\n Mean average accuracy calculated on rotation errors\n\n translation_maa: float\n Mean average accuracy calculated on translation errors\n \"\"\"\n\n accuracies, rotation_accuracies, translation_accuracies = [], [], []\n\n for rotation_error_threshold, translation_error_threshold in zip(rotation_error_thresholds, translation_error_thresholds):\n\n # Calculate whether the errors are less than specified thresholds or not\n rotation_accuracy = (rotation_errors <= rotation_error_threshold)\n translation_accuracy = (translation_errors <= translation_error_threshold)\n accuracy = rotation_accuracy & translation_accuracy\n\n accuracies.append(accuracy.astype(np.float32).mean())\n rotation_accuracies.append(rotation_accuracy.astype(np.float32).mean())\n translation_accuracies.append(translation_accuracy.astype(np.float32).mean())\n\n maa = np.array(accuracies).mean()\n rotation_maa = np.array(rotation_accuracies).mean()\n translation_maa = np.array(translation_accuracies).mean()\n\n return maa, rotation_maa, translation_maa\n\n\ndef evaluate(df, verbose=False):\n\n \"\"\"\n Calculate mean average accuracies over a set of thresholds for rotation and translation\n\n Parameters\n ----------\n df: pandas.DataFrame\n Dataframe with dataset, scene, rotation_matrix, translation_vector, rotation_matrix_prediction and translation_vector_prediction columns\n\n verbose: bool\n Whether to print scores or not\n\n Returns\n -------\n df_scores: pandas.DataFrame\n Dataframe of scores\n \"\"\"\n\n rotation_error_thresholds = {\n **{('haiper', scene): np.linspace(1, 10, 10) for scene in ['bike', 'chairs', 'fountain']},\n **{('heritage', scene): np.linspace(1, 10, 10) for scene in ['cyprus', 'dioscuri']},\n **{('heritage', 'wall'): np.linspace(0.2, 10, 10)},\n **{('urban', 'kyiv-puppet-theater'): np.linspace(1, 10, 10)},\n }\n translation_error_thresholds = {\n **{('haiper', scene): np.geomspace(0.05, 0.5, 10) for scene in ['bike', 'chairs', 'fountain']},\n **{('heritage', scene): np.geomspace(0.1, 2, 10) for scene in ['cyprus', 'dioscuri']},\n **{('heritage', 'wall'): np.geomspace(0.05, 1, 10)},\n **{('urban', 'kyiv-puppet-theater'): np.geomspace(0.5, 5, 10)},\n }\n df_scores = pd.DataFrame(columns=['dataset', 'scene', 'image_pairs', 'maa', 'rotation_maa', 'translation_maa'])\n\n for (dataset, scene), df_scene in tqdm(df.groupby(['dataset', 'scene'])):\n\n scene_rotation_errors = []\n scene_translation_errors = []\n\n for i in range(df_scene.shape[0]):\n for j in range(i + 1, df_scene.shape[0]):\n\n rotation_matrix_difference_ground_truth, translation_vector_difference_ground_truth = pose_difference(\n r1=string_to_array((df_scene.iloc[i]['rotation_matrix'])).reshape(3, 3),\n t1=string_to_array((df_scene.iloc[i]['translation_vector'])),\n r2=string_to_array((df_scene.iloc[j]['rotation_matrix'])).reshape(3, 3),\n t2=string_to_array((df_scene.iloc[j]['translation_vector'])),\n )\n quaternion_ground_truth = rotation_matrix_to_quaternion(rotation_matrix=rotation_matrix_difference_ground_truth)\n\n rotation_matrix_difference_prediction, translation_vector_difference_prediction = pose_difference(\n r1=string_to_array((df_scene.iloc[i]['rotation_matrix_prediction'])).reshape(3, 3),\n t1=string_to_array((df_scene.iloc[i]['translation_vector_prediction'])),\n r2=string_to_array((df_scene.iloc[j]['rotation_matrix_prediction'])).reshape(3, 3),\n t2=string_to_array((df_scene.iloc[j]['translation_vector_prediction'])),\n )\n quaternion_prediction = rotation_matrix_to_quaternion(rotation_matrix=rotation_matrix_difference_prediction)\n\n rotation_error, translation_error = rotation_and_translation_error(\n q_ground_truth=quaternion_ground_truth,\n t_ground_truth=translation_vector_difference_ground_truth,\n q_prediction=quaternion_prediction,\n t_prediction=translation_vector_difference_prediction,\n epsilon=1e-15\n )\n scene_rotation_errors.append(rotation_error)\n scene_translation_errors.append(translation_error)\n\n scene_maa, scene_rotation_maa, scene_translation_maa = mean_average_accuracy(\n rotation_errors=scene_rotation_errors,\n translation_errors=scene_translation_errors,\n rotation_error_thresholds=rotation_error_thresholds[(dataset, scene)],\n translation_error_thresholds=translation_error_thresholds[(dataset, scene)]\n )\n\n if verbose:\n settings.logger.info(\n f'''\n Dataset: {dataset} - Scene: {scene}\n Number of image pairs: {len(scene_rotation_errors)}\n mAA: {scene_maa:.6f} - rotation mAA: {scene_rotation_maa:.6f} - translation mAA: {scene_translation_maa:.6f}\n '''\n )\n\n df_scores = pd.concat((\n df_scores,\n pd.DataFrame(\n data=[[dataset, scene, len(scene_rotation_errors), scene_maa, scene_rotation_maa, scene_translation_maa]],\n columns=['dataset', 'scene', 'image_pairs', 'maa', 'rotation_maa', 'translation_maa']\n )\n ), axis=0)\n\n return df_scores\n","metadata":{"execution":{"iopub.status.busy":"2024-05-12T07:01:29.288445Z","iopub.execute_input":"2024-05-12T07:01:29.288704Z","iopub.status.idle":"2024-05-12T07:01:29.323379Z","shell.execute_reply.started":"2024-05-12T07:01:29.288682Z","shell.execute_reply":"2024-05-12T07:01:29.322483Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 6. Models","metadata":{}},{"cell_type":"code","source":"def loftr_match_images(image1, image2, model, device, amp, transforms, confidence_threshold, top_k):\n\n \"\"\"\n Match given two images with each other using LoFTR model\n\n Parameters\n ----------\n image1: numpy.ndarray of shape (3, height, width)\n Array of first image\n\n image2: numpy.ndarray of shape (3, height, width)\n Array of second image\n\n model: torch.nn.Module\n LoFTR Model\n\n device: torch.device\n Location of the image1, image2 and the model\n\n amp: bool\n Whether to use auto mixed precision or not\n\n transforms: dict\n Dictionary of transform parameters\n\n confidence_threshold: float or int\n Confidence threshold to filter out low confidence matches\n\n top_k: int\n Number of matches to take\n\n Returns\n -------\n outputs: dict\n Model outputs\n \"\"\"\n\n image1_raw_height, image1_raw_width = image1.shape[:2]\n image1 = get_image_tensor(\n image_path_or_array=image1,\n resize=transforms['resize'],\n resize_shape=transforms['resize_shape'],\n resize_longest_edge=transforms['resize_longest_edge'],\n scale=transforms['scale'],\n grayscale=transforms['grayscale']\n )\n image1 = image1.to(device)\n image1_transformed_height, image1_transformed_width = image1.shape[2:]\n\n image2_raw_height, image2_raw_width = image2.shape[:2]\n image2 = get_image_tensor(\n image_path_or_array=image2,\n resize=transforms['resize'],\n resize_shape=transforms['resize_shape'],\n resize_longest_edge=transforms['resize_longest_edge'],\n scale=transforms['scale'],\n grayscale=transforms['grayscale']\n )\n image2 = image2.to(device)\n image2_transformed_height, image2_transformed_width = image2.shape[2:]\n\n with torch.no_grad():\n if amp:\n with torch.autocast(device_type=device.type, dtype=torch.float16):\n outputs = model({'image0': image1, 'image1': image2})\n else:\n outputs = model({'image0': image1, 'image1': image2})\n\n for k in outputs.keys():\n outputs[k] = outputs[k].detach().cpu().numpy()\n\n if confidence_threshold is not None:\n if isinstance(confidence_threshold, float):\n # Select matched keypoints with above given confidence threshold\n confidence_mask = outputs['confidence'] >= confidence_threshold\n elif isinstance(confidence_threshold, int):\n # Select keypoints dynamically based on confidence distribution\n confidence_mean, confidence_std = outputs['confidence'].mean(), outputs['confidence'].std()\n confidence_mask = outputs['confidence'] >= (confidence_mean + (confidence_std * confidence_threshold))\n else:\n raise ValueError(f'Invalid confidence_threshold {confidence_threshold}')\n\n for k in outputs.keys():\n outputs[k] = outputs[k][confidence_mask]\n\n if top_k is not None:\n # Select top-k keypoints based on their confidences\n sorting_idx = outputs['matching_scores0'].argsort()[-top_k:]\n for k in outputs.keys():\n outputs[k] = outputs[k][sorting_idx]\n\n outputs['keypoints0'][:, 0] *= image1_raw_width / image1_transformed_width\n outputs['keypoints0'][:, 1] *= image1_raw_height / image1_transformed_height\n outputs['keypoints1'][:, 0] *= image2_raw_width / image2_transformed_width\n outputs['keypoints1'][:, 1] *= image2_raw_height / image2_transformed_height\n\n return outputs\n","metadata":{"execution":{"iopub.status.busy":"2024-05-12T07:01:29.324360Z","iopub.execute_input":"2024-05-12T07:01:29.324604Z","iopub.status.idle":"2024-05-12T07:01:29.340561Z","shell.execute_reply.started":"2024-05-12T07:01:29.324582Z","shell.execute_reply":"2024-05-12T07:01:29.339679Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def superglue_match_images(image1, image2, model, device, amp, transforms, score_threshold, top_k):\n\n \"\"\"\n Match given two images with each other using SuperGlue model\n\n Parameters\n ----------\n image1: numpy.ndarray of shape (3, height, width)\n Array of first image\n\n image2: numpy.ndarray of shape (3, height, width)\n Array of second image\n\n model: torch.nn.Module\n SuperGlue Model\n\n device: torch.device\n Location of the image1, image2 and the model\n\n amp: bool\n Whether to use auto mixed precision or not\n\n transforms: dict\n Dictionary of transform parameters\n\n score_threshold: float, int or None\n Confidence threshold\n\n top_k: int or None\n Number of keypoints to take\n\n Returns\n -------\n outputs: dict\n Model outputs\n \"\"\"\n\n image1_raw_height, image1_raw_width = image1.shape[:2]\n image1 = get_image_tensor(\n image_path_or_array=image1,\n resize=transforms[1]['resize'],\n resize_shape=transforms[1]['resize_shape'],\n resize_longest_edge=transforms[1]['resize_longest_edge'],\n scale=transforms[1]['scale'],\n grayscale=transforms[1]['grayscale']\n )\n image1 = image1.to(device)\n image1_transformed_height, image1_transformed_width = image1.shape[2:]\n\n image2_raw_height, image2_raw_width = image2.shape[:2]\n image2 = get_image_tensor(\n image_path_or_array=image2,\n resize=transforms[2]['resize'],\n resize_shape=transforms[2]['resize_shape'],\n resize_longest_edge=transforms[2]['resize_longest_edge'],\n scale=transforms[2]['scale'],\n grayscale=transforms[2]['grayscale']\n )\n image2 = image2.to(device)\n image2_transformed_height, image2_transformed_width = image2.shape[2:]\n\n with torch.no_grad():\n if amp:\n with torch.autocast(device_type=device.type, dtype=torch.float16):\n outputs = model({'image0': image1, 'image1': image2})\n else:\n outputs = model({'image0': image1, 'image1': image2})\n\n for k in outputs.keys():\n if k == 'descriptors0' or k == 'descriptors1':\n outputs[k] = outputs[k][0].detach().cpu().numpy().T\n else:\n outputs[k] = outputs[k][0].detach().cpu().numpy()\n\n matches_mask = outputs['matches0'] > -1\n\n for k in ['keypoints1', 'scores1', 'descriptors1', 'matches1', 'matching_scores1']:\n outputs[k] = outputs[k][outputs['matches0'][matches_mask]]\n\n for k in ['keypoints0', 'scores0', 'descriptors0', 'matches0', 'matching_scores0']:\n outputs[k] = outputs[k][matches_mask]\n\n if score_threshold is not None:\n if isinstance(score_threshold, float):\n # Select matched keypoints with above given score threshold\n score_mask = outputs['matching_scores0'] >= score_threshold\n elif isinstance(score_threshold, int):\n # Select keypoints dynamically based on score distribution\n score_mean, score_std = outputs['matching_scores0'].mean(), outputs['matching_scores0'].std()\n score_mask = outputs['matching_scores0'] >= (score_mean + (score_std * score_threshold))\n else:\n raise ValueError(f'Invalid score_threshold {score_threshold}')\n\n for k in outputs.keys():\n outputs[k] = outputs[k][score_mask]\n\n if top_k is not None:\n # Select top-k keypoints based on their scores\n sorting_idx = outputs['matching_scores0'].argsort()[-top_k:]\n for k in outputs.keys():\n outputs[k] = outputs[k][sorting_idx]\n\n outputs['keypoints0'][:, 0] *= image1_raw_width / image1_transformed_width\n outputs['keypoints0'][:, 1] *= image1_raw_height / image1_transformed_height\n outputs['keypoints1'][:, 0] *= image2_raw_width / image2_transformed_width\n outputs['keypoints1'][:, 1] *= image2_raw_height / image2_transformed_height\n\n return outputs\n","metadata":{"execution":{"iopub.status.busy":"2024-05-12T07:01:29.341580Z","iopub.execute_input":"2024-05-12T07:01:29.342120Z","iopub.status.idle":"2024-05-12T07:01:29.360089Z","shell.execute_reply.started":"2024-05-12T07:01:29.342095Z","shell.execute_reply":"2024-05-12T07:01:29.359309Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 7. Inference","metadata":{}},{"cell_type":"code","source":"image_matching_device = torch.device('cuda')\n'''\n# Load LoFTR model with specified configurations\nloftr_model = LoFTR(\n pretrained=None,\n config={\n 'backbone_type': 'ResNetFPN',\n 'resolution': (8, 2),\n 'fine_window_size': 5,\n 'fine_concat_coarse_feat': True,\n 'resnetfpn': {\n 'initial_dim': 128,\n 'block_dims': [128, 196, 256]\n },\n 'coarse': {\n 'd_model': 256,\n 'd_ffn': 256,\n 'nhead': 8,\n 'layer_names': ['self', 'cross', 'self', 'cross', 'self', 'cross', 'self', 'cross'],\n 'attention': 'linear',\n 'temp_bug_fix': False,\n },\n 'match_coarse': {\n 'thr': 0.2,\n 'border_rm': 2,\n 'match_type': 'dual_softmax',\n 'dsmax_temperature': 0.1,\n 'skh_iters': 3,\n 'skh_init_bin_score': 1.0,\n 'skh_prefilter': True,\n 'train_coarse_percent': 0.4,\n 'train_pad_num_gt_min': 200,\n },\n 'fine': {\n 'd_model': 128,\n 'd_ffn': 128,\n 'nhead': 8,\n 'layer_names': ['self', 'cross'],\n 'attention': 'linear'\n }\n }\n)\nloftr_model.load_state_dict(torch.load(external_dataset / 'models' / 'loftr' / 'loftr_outdoor.ckpt')['state_dict'], strict=False)\nloftr_model = loftr_model.eval().to(image_matching_device)\n'''\n# Load SuperPoint and SuperGlue model with specified configurations\nsuperglue_model = Matching(config={\n 'superpoint': {\n 'descriptor_dim': 256,\n 'nms_radius': 4,\n 'keypoint_threshold': 0.01,\n 'max_keypoints': -1,\n 'remove_borders': 4\n },\n 'superglue': {\n 'descriptor_dim': 256,\n 'weights': 'outdoor',\n 'keypoint_encoder': [32, 64, 128, 256],\n 'sinkhorn_iterations': 100,\n 'match_threshold': 0.2\n }\n})\nsuperglue_model = superglue_model.eval().to(image_matching_device)","metadata":{"execution":{"iopub.status.busy":"2024-05-12T07:01:29.361405Z","iopub.execute_input":"2024-05-12T07:01:29.361641Z","iopub.status.idle":"2024-05-12T07:01:30.352826Z","shell.execute_reply.started":"2024-05-12T07:01:29.361620Z","shell.execute_reply":"2024-05-12T07:01:30.351864Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def match(read_image_function, image_paths, image_pair_indices, feature_dir, superglue_model, stage2=False):\n \n with h5py.File(f'{feature_dir}/matches_loftr.h5', mode='w') as f_match:\n for pair_idx in progress_bar(image_pair_indices):\n \n idx1, idx2 = pair_idx\n fname1, fname2 = image_paths[idx1], image_paths[idx2]\n key1, key2 = fname1.split('/')[-1], fname2.split('/')[-1]\n \n image1 = read_image_function(fname1)\n image2 = read_image_function(fname2)\n \n keypoints1 = []\n keypoints2 = []\n \n # Largest SuperGlue size that can be used on Kaggle GPU is 2560\n superglue_longest_edge_limit = 2560\n if (np.max(image1.shape[:2]) > superglue_longest_edge_limit) and (np.max(image2.shape[:2]) > superglue_longest_edge_limit):\n # Both of the images have longest edges greater than 2560\n first_stage_superglue_transforms = {\n 1: {\n 'resize': True,\n 'resize_shape': superglue_longest_edge_limit,\n 'resize_longest_edge': True,\n 'scale': True,\n 'grayscale': True\n },\n 2: {\n 'resize': True,\n 'resize_shape': superglue_longest_edge_limit,\n 'resize_longest_edge': True,\n 'scale': True,\n 'grayscale': True\n }\n }\n elif (np.max(image1.shape[:2]) > superglue_longest_edge_limit) and (np.max(image2.shape[:2]) <= superglue_longest_edge_limit):\n # First image's longest edge is greater than 2560\n first_stage_superglue_transforms = {\n 1: {\n 'resize': True,\n 'resize_shape': superglue_longest_edge_limit,\n 'resize_longest_edge': True,\n 'scale': True,\n 'grayscale': True\n },\n 2: {\n 'resize': False,\n 'resize_shape': None,\n 'resize_longest_edge': None,\n 'scale': True,\n 'grayscale': True\n }\n }\n elif (np.max(image1.shape[:2]) <= superglue_longest_edge_limit) and (np.max(image2.shape[:2]) > superglue_longest_edge_limit):\n # Second image's longest edge is greater than 2560\n first_stage_superglue_transforms = {\n 1: {\n 'resize': False,\n 'resize_shape': None,\n 'resize_longest_edge': None,\n 'scale': True,\n 'grayscale': True\n },\n 2: {\n 'resize': True,\n 'resize_shape': superglue_longest_edge_limit,\n 'resize_longest_edge': True,\n 'scale': True,\n 'grayscale': True\n }\n }\n else:\n # Neither of the image's longest edge is greater than 2560\n first_stage_superglue_transforms = {\n 1: {\n 'resize': False,\n 'resize_shape': None,\n 'resize_longest_edge': None,\n 'scale': True,\n 'grayscale': True\n },\n 2: {\n 'resize': False,\n 'resize_shape': None,\n 'resize_longest_edge': None,\n 'scale': True,\n 'grayscale': True\n }\n }\n \n first_stage_superglue_outputs = superglue_match_images(\n image1=image1,\n image2=image2,\n model=superglue_model,\n device=image_matching_device,\n amp=True,\n transforms=first_stage_superglue_transforms,\n score_threshold=None,\n top_k=None\n )\n \n keypoints1.append(first_stage_superglue_outputs['keypoints0'])\n keypoints2.append(first_stage_superglue_outputs['keypoints1'])\n del first_stage_superglue_outputs\n \n keypoints1 = np.concatenate(keypoints1, axis=0)\n keypoints2 = np.concatenate(keypoints2, axis=0)\n \n if stage2:\n if (keypoints1.shape[0] > 10) & (keypoints1.shape[0] < 200):\n \n image1_cropped, keypoints1_cropped, x_offset1, y_offset1 = crop(image1, keypoints1)\n image2_cropped, keypoints2_cropped, x_offset2, y_offset2 = crop(image2, keypoints2)\n \n # Largest SuperGlue size that can be used on Kaggle GPU is 2560\n superglue_longest_edge_limit = 2560\n if (np.max(image1_cropped.shape[:2]) > superglue_longest_edge_limit) and (np.max(image2_cropped.shape[:2]) > superglue_longest_edge_limit):\n # Both of the images have longest edges greater than 2560\n second_stage_superglue_transforms = {\n 1: {\n 'resize': True,\n 'resize_shape': superglue_longest_edge_limit,\n 'resize_longest_edge': True,\n 'scale': True,\n 'grayscale': True\n },\n 2: {\n 'resize': True,\n 'resize_shape': superglue_longest_edge_limit,\n 'resize_longest_edge': True,\n 'scale': True,\n 'grayscale': True\n }\n }\n elif (np.max(image1_cropped.shape[:2]) > superglue_longest_edge_limit) and (np.max(image2_cropped.shape[:2]) <= superglue_longest_edge_limit):\n # First image's longest edge is greater than 2560\n second_stage_superglue_transforms = {\n 1: {\n 'resize': True,\n 'resize_shape': superglue_longest_edge_limit,\n 'resize_longest_edge': True,\n 'scale': True,\n 'grayscale': True\n },\n 2: {\n 'resize': False,\n 'resize_shape': None,\n 'resize_longest_edge': None,\n 'scale': True,\n 'grayscale': True\n }\n }\n elif (np.max(image1_cropped.shape[:2]) <= superglue_longest_edge_limit) and (np.max(image2_cropped.shape[:2]) > superglue_longest_edge_limit):\n # Second image's longest edge is greater than 2560\n second_stage_superglue_transforms = {\n 1: {\n 'resize': False,\n 'resize_shape': None,\n 'resize_longest_edge': None,\n 'scale': True,\n 'grayscale': True\n },\n 2: {\n 'resize': True,\n 'resize_shape': superglue_longest_edge_limit,\n 'resize_longest_edge': True,\n 'scale': True,\n 'grayscale': True\n }\n }\n else:\n # Neither of the image's longest edge is greater than 2560\n second_stage_superglue_transforms = {\n 1: {\n 'resize': False,\n 'resize_shape': None,\n 'resize_longest_edge': None,\n 'scale': True,\n 'grayscale': True\n },\n 2: {\n 'resize': False,\n 'resize_shape': None,\n 'resize_longest_edge': None,\n 'scale': True,\n 'grayscale': True\n }\n }\n\n second_stage_superglue_outputs = superglue_match_images(\n image1=image1_cropped,\n image2=image2_cropped,\n model=superglue_model,\n device=image_matching_device,\n amp=True,\n transforms=second_stage_superglue_transforms,\n score_threshold=None,\n top_k=200\n )\n\n second_stage_superglue_outputs['keypoints0'][:, 0] += x_offset1\n second_stage_superglue_outputs['keypoints0'][:, 1] += y_offset1\n second_stage_superglue_outputs['keypoints1'][:, 0] += x_offset2\n second_stage_superglue_outputs['keypoints1'][:, 1] += y_offset2\n\n keypoints1 = np.concatenate([keypoints1, second_stage_superglue_outputs['keypoints0']], axis=0)\n keypoints2 = np.concatenate([keypoints2, second_stage_superglue_outputs['keypoints1']], axis=0)\n \n n_matches = len(keypoints1)\n group = f_match.require_group(key1)\n if n_matches >= 15:\n group.create_dataset(key2, data=np.concatenate([keypoints1, keypoints2], axis=1))\n\n kpts = defaultdict(list)\n match_indexes = defaultdict(dict)\n total_kpts=defaultdict(int)\n \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 \n unique_kpts = {}\n unique_match_idxs = {}\n out_match = defaultdict(dict)\n \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 \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([\n unique_kpts[k1][m2[:,0]],\n unique_kpts[k2][m2[:,1]]\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 \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 \n return\n\n\ndef import_into_colmap(img_dir, feature_dir='.featureout', database_path='colmap.db'):\n \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, 'simple-radial', single_camera)\n add_matches(db, feature_dir, fname_to_id,)\n db.commit()\n \n return\n","metadata":{"execution":{"iopub.status.busy":"2024-05-12T07:01:30.354396Z","iopub.execute_input":"2024-05-12T07:01:30.354755Z","iopub.status.idle":"2024-05-12T07:01:30.396560Z","shell.execute_reply.started":"2024-05-12T07:01:30.354725Z","shell.execute_reply":"2024-05-12T07:01:30.395787Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"reconstruction_root_directory = Path('./inference')\nreconstruction_root_directory.mkdir(parents=True, exist_ok=True)\n\ndatasets = [\n directory for directory in os.listdir(train_or_test_directory)\n if (train_or_test_directory / directory).is_dir() and (directory in df['dataset'].unique())\n]\n\nfor dataset in datasets:\n \n dataset_directory = train_or_test_directory / dataset\n scenes = [\n directory for directory in os.listdir(dataset_directory)\n if (dataset_directory / directory).is_dir() and (directory in df['scene'].unique())\n ]\n \n for scene in scenes:\n \n df_scene = df.loc[df['scene'] == scene, :]\n \n scene_directory = dataset_directory / scene\n scene_image_directory = scene_directory / 'images'\n image_paths = sorted(glob(str(scene_image_directory / '*')))\n scene_image_count = len(image_paths)\n \n scene_reconstruction_directory = reconstruction_root_directory / dataset / scene\n scene_reconstruction_directory.mkdir(parents=True, exist_ok=True) \n database_path = scene_reconstruction_directory / 'colmap.db'\n if os.path.isfile(database_path):\n os.remove(database_path)\n \n scene_mean_memory_usage = df_scene['memory_usage'].mean()\n scene_mean_memory_usage_limit = 16\n if scene_mean_memory_usage < scene_mean_memory_usage_limit:\n \n # Create image pair indices from given image paths\n scene_image_pair_indices = create_image_pairs(\n image_paths,\n sim_th=0.85,\n min_pairs=120,\n exhaustive_if_less=200,\n device=image_matching_device\n )\n \n @lru_cache(maxsize=2)\n def read_image(image_path):\n return cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB)\n \n match(\n read_image,\n image_paths,\n scene_image_pair_indices,\n feature_dir=scene_reconstruction_directory,\n superglue_model=superglue_model,\n stage2=False\n )\n import_into_colmap(scene_image_directory, feature_dir=scene_reconstruction_directory, database_path=database_path)\n else:\n sift_extraction_options = pycolmap.SiftExtractionOptions()\n sift_extraction_options.num_threads = -1\n sift_extraction_options.max_image_size = 1400\n sift_extraction_options.max_num_features = 8192\n sift_extraction_options.estimate_affine_shape = False\n sift_extraction_options.upright = False\n sift_extraction_options.normalization = 'L2'\n pycolmap.extract_features(\n database_path=database_path,\n image_path=scene_image_directory,\n image_list=[image_path.split('/')[-1] for image_path in image_paths],\n sift_options=sift_extraction_options,\n device=pycolmap.Device('cpu'),\n verbose=verbose\n )\n \n pycolmap.match_exhaustive(database_path)\n \n output_path = scene_reconstruction_directory / 'colmap_rec'\n output_path.mkdir(parents=True, exist_ok=True) \n\n incremental_mapper_options = pycolmap.IncrementalMapperOptions()\n incremental_mapper_options.min_model_size = 3\n incremental_mapper_options.min_num_matches = 5\n reconstructions = pycolmap.incremental_mapping(\n database_path=database_path,\n image_path=scene_image_directory,\n output_path=output_path,\n options=incremental_mapper_options\n )\n \n if len(reconstructions) > 0:\n\n best_registered_image_count = 0\n best_reconstruction_idx = None\n\n for reconstruction_idx in reconstructions.keys():\n if reconstructions[reconstruction_idx].num_reg_images() > best_registered_image_count:\n best_reconstruction_idx = reconstruction_idx\n best_registered_image_count = reconstructions[reconstruction_idx].num_reg_images()\n\n best_reconstruction = reconstructions[best_reconstruction_idx]\n else:\n best_registered_image_count = 0\n best_reconstruction_idx = None\n best_reconstruction = None\n\n if verbose:\n print(\n f'''\n Dataset: {dataset} - Scene: {scene}\n Reconstruction count: {len(reconstructions)}\n Best reconstruction registered image count: {best_registered_image_count}/{scene_image_count}\n '''\n )\n\n if best_reconstruction is not None:\n registered_images = {image.name: image for image in best_reconstruction.images.values()}\n else:\n registered_images = {}\n \n for idx, row in df.loc[(df['dataset'] == dataset) & (df['scene'] == scene)].iterrows():\n if row['image_id'] in registered_images:\n rotation_matrix_prediction = registered_images[row['image_id']].rotmat()\n translation_vector_prediction = registered_images[row['image_id']].tvec\n df.loc[idx, 'rotation_matrix'] = ';'.join([str(x) for x in rotation_matrix_prediction.reshape(-1)])\n df.loc[idx, 'translation_vector'] = ';'.join([str(x) for x in translation_vector_prediction.reshape(-1)])\n else:\n df.loc[idx, 'rotation_matrix'] = np.nan\n df.loc[idx, 'translation_vector'] = np.nan\n\n # Fill unregistered images rotation matrices with the prediction mean or zeros\n scene_rotation_matrix_predictions = df.loc[df['scene'] == scene, 'rotation_matrix'].dropna().apply(lambda x: np.array(str(x).split(';'), dtype=np.float64).reshape(1, 3, 3)).values\n if scene_rotation_matrix_predictions.shape[0] == 0:\n rotation_matrix_fill_value = np.zeros((3, 3))\n else:\n rotation_matrix_fill_value = np.mean(np.concatenate(scene_rotation_matrix_predictions, axis=0), axis=0)\n df.loc[(df['scene'] == scene) & (df['rotation_matrix'].isnull()), 'rotation_matrix'] = ';'.join([str(x) for x in rotation_matrix_fill_value.reshape(-1)])\n\n # Fill unregistered images translation vectors with the prediction mean or zeros\n scene_translation_vector_predictions = df.loc[df['scene'] == scene, 'translation_vector'].dropna().apply(lambda x: np.array(str(x).split(';'), dtype=np.float64).reshape(1, 3)).values\n if scene_translation_vector_predictions.shape[0] == 0:\n translation_vector_fill_value = np.zeros((3, 1))\n else:\n translation_vector_fill_value = np.mean(np.concatenate(scene_translation_vector_predictions, axis=0), axis=0)\n df.loc[(df['scene'] == scene) & (df['translation_vector'].isnull()), 'translation_vector'] = ';'.join([str(x) for x in translation_vector_fill_value.reshape(-1)])\n","metadata":{"execution":{"iopub.status.busy":"2024-05-12T07:01:30.397699Z","iopub.execute_input":"2024-05-12T07:01:30.397965Z","iopub.status.idle":"2024-05-12T07:01:30.428694Z","shell.execute_reply.started":"2024-05-12T07:01:30.397942Z","shell.execute_reply":"2024-05-12T07:01:30.427898Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 8. Submission","metadata":{}},{"cell_type":"code","source":"df_submission = df.loc[:, ['image_path', 'dataset', 'scene', 'rotation_matrix', 'translation_vector']]\ndf_submission.to_csv('submission.csv', index=False)","metadata":{"execution":{"iopub.status.busy":"2024-05-12T07:01:30.429696Z","iopub.execute_input":"2024-05-12T07:01:30.430287Z","iopub.status.idle":"2024-05-12T07:01:30.442827Z","shell.execute_reply.started":"2024-05-12T07:01:30.430262Z","shell.execute_reply":"2024-05-12T07:01:30.441955Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!cat submission.csv","metadata":{"execution":{"iopub.status.busy":"2024-05-12T07:01:30.443919Z","iopub.execute_input":"2024-05-12T07:01:30.444278Z","iopub.status.idle":"2024-05-12T07:01:31.415574Z","shell.execute_reply.started":"2024-05-12T07:01:30.444247Z","shell.execute_reply":"2024-05-12T07:01:31.414412Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}