{"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":"","metadata":{}},{"cell_type":"markdown","source":"origin：https://www.kaggle.com/code/eduardtrulls/imc-2023-submission-example?scriptVersionId=129165789","metadata":{}},{"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 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","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-05-24T01:19:11.891605Z","iopub.execute_input":"2023-05-24T01:19:11.892334Z","iopub.status.idle":"2023-05-24T01:19:11.900381Z","shell.execute_reply.started":"2023-05-24T01:19:11.892291Z","shell.execute_reply":"2023-05-24T01:19:11.899253Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"os：Python的标准库，提供了与操作系统交互的功能，例如文件和目录操作。\n\ntqdm：一个用于在循环中显示进度条的库。\n\ntime：Python的标准库，提供了与时间相关的功能，如计时和时间戳。\n\nfastprogress：用于显示快速进度条的库。\n\ngc：Python的标准库，提供了垃圾回收功能，用于自动管理内存。\n\nnumpy：一个用于科学计算的Python库，提供了高性能的多维数组对象和各种计算函数。\n\nh5py：一个用于读写HDF5文件的库，HDF5是一种数据存储格式。\n\nIPython.display：IPython的库，提供了在交互式环境中显示内容的功能。\n\ncollections.defaultdict：一个字典子类，提供了在访问不存在的键时返回默认值的功能。\n\ncopy.deepcopy：用于创建对象的深拷贝，包括嵌套对象的拷贝。\n\ncv2：OpenCV库的Python接口，用于计算机视觉任务，如图像处理和计算机视觉算法。\n\ntorch：PyTorch深度学习框架的主要库，提供了张量操作和自动微分功能。\n\ntorch.nn.functional：PyTorch的子模块，提供了一些函数形式的神经网络操作，如激活函数和损失函数。\n\nkornia：一个用于计算机视觉任务的PyTorch库，提供了各种几何变换和图像处理操作。\n\nkornia.feature：kornia的子模块，提供了一些特征提取相关的函数和模块。\n\nPIL.Image：Python Imaging Library（PIL）的模块，用于图像处理和操作。\n\ntimm：一个用于图像分类和计算机视觉任务的PyTorch库，提供了各种预训练模型和训练工具。\n\ntimm.data.resolve_data_config：timm库的函数，用于解析数据配置。\n\ntimm.data.transforms_factory.create_transform：timm库的函数，用于创建数据转换。\n\npycolmap：用于三维重建的Colmap库的Python接口，Colmap是一个开源的视觉SLAM系统。","metadata":{}},{"cell_type":"code","source":"print('Kornia version', K.__version__)\nprint('Pycolmap version', pycolmap.__version__)\n\nLOCAL_FEATURE = 'KeyNetAffNetHardNet'\ndevice=torch.device('cuda')\n# Can be LoFTR, KeyNetAffNetHardNet, or DISK\n\n#这两行代码分别给变量LOCAL_FEATURE和device赋值。\n#LOCAL_FEATURE被赋值为字符串'KeyNetAffNetHardNet'，表示选择使用KeyNetAffNetHardNet作为本地特征算法。\n#device被赋值为torch.device('cuda')，表示将使用CUDA加速来执行计算。","metadata":{"execution":{"iopub.status.busy":"2023-05-24T01:19:11.902324Z","iopub.execute_input":"2023-05-24T01:19:11.903431Z","iopub.status.idle":"2023-05-24T01:19:11.914418Z","shell.execute_reply.started":"2023-05-24T01:19:11.903389Z","shell.execute_reply":"2023-05-24T01:19:11.913169Z"},"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-05-24T01:19:11.917364Z","iopub.execute_input":"2023-05-24T01:19:11.917816Z","iopub.status.idle":"2023-05-24T01:19:11.925294Z","shell.execute_reply.started":"2023-05-24T01:19:11.917778Z","shell.execute_reply":"2023-05-24T01:19:11.924084Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"函数首先使用OpenCV的cv2.imread函数加载图像文件，并使用K.image_to_tensor函数将图像转换为PyTorch张量。False参数表示图像不进行归一化。然后，将图像的像素值从BGR颜色空间转换为RGB颜色空间，使用K.color.bgr_to_rgb函数实现。最后，如果提供了device参数，将图像张量移动到指定的设备上，然后返回图像张量。\n\n这两个函数分别用于将数组转换为字符串和加载图像并转换为PyTorch张量，为后续的数据处理和图像操作提供了方便的工具函数。","metadata":{}},{"cell_type":"code","source":"# We will use ViT global descriptor to get matching shortlists.\ndef get_global_desc(fnames, model,\n device = torch.device('cpu')):\n model = model.eval()\n model= model.to(device)\n config = resolve_data_config({}, model=model)\n transform = create_transform(**config)\n global_descs_convnext=[]\n for i, img_fname_full in tqdm(enumerate(fnames),total= len(fnames)):\n key = os.path.splitext(os.path.basename(img_fname_full))[0]\n img = Image.open(img_fname_full).convert('RGB')\n timg = transform(img).unsqueeze(0).to(device)\n with torch.no_grad():\n desc = model.forward_features(timg.to(device)).mean(dim=(-1,2))#\n #print (desc.shape)\n desc = desc.view(1, -1)\n desc_norm = F.normalize(desc, dim=1, p=2)\n #print (desc_norm)\n global_descs_convnext.append(desc_norm.detach().cpu())\n global_descs_all = torch.cat(global_descs_convnext, dim=0)\n return global_descs_all\n\n#该函数用于提取图像的全局描述符。它接受一个图像文件名列表fnames，一个模型model，和一个可选的device参数，默认为torch.device('cpu')。\n#首先，将模型设置为评估模式，并将其移动到指定的设备上。然后，使用resolve_data_config函数解析模型配置，\n#并使用create_transform函数创建数据转换。然后，对于每个图像文件名，打开图像文件并将其转换为RGB格式，然后应用数据转换并将图像转换为张量。\n#之后，使用模型提取特征，并计算特征的均值，得到全局描述符。最后，将所有图像的全局描述符拼接成一个张量并返回。\n\ndef get_img_pairs_exhaustive(img_fnames):\n index_pairs = []\n for i in range(len(img_fnames)):\n for j in range(i+1, len(img_fnames)):\n index_pairs.append((i,j))\n return index_pairs\n#该函数用于生成图像之间的完全匹配的索引对。它接受一个图像文件名列表img_fnames作为输入。\n#函数通过两层循环遍历所有图像文件名的组合，生成索引对。对于索引i和j，其中i= 3\n\nMAX_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 if IS_PYTHON3:\n return array.tostring()\n else:\n return np.getbuffer(array)\n\n\ndef blob_to_array(blob, dtype, shape=(-1,)):\n if IS_PYTHON3:\n return np.fromstring(blob, dtype=dtype).reshape(*shape)\n else:\n return np.frombuffer(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)))","metadata":{"execution":{"iopub.status.busy":"2023-05-24T01:19:12.179360Z","iopub.execute_input":"2023-05-24T01:19:12.179721Z","iopub.status.idle":"2023-05-24T01:19:12.215656Z","shell.execute_reply.started":"2023-05-24T01:19:12.179661Z","shell.execute_reply":"2023-05-24T01:19:12.214525Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"这段代码是一个用于操作COLMAP数据库的脚本，下面对其进行注释和解释：\n\nCOLMAPDatabase 类是一个继承自 sqlite3.Connection 的自定义类，用于操作 COLMAP 数据库。它包含了一系列方法来添加相机、图像、关键点、描述符、匹配等数据到数据库中。\n\nIS_PYTHON3 变量用于检查是否运行在 Python 3 及以上的版本上。\n\nMAX_IMAGE_ID 变量定义了图像ID的最大值，用于在生成匹配对ID时进行计算。\n\nCREATE_CAMERAS_TABLE、CREATE_DESCRIPTORS_TABLE、CREATE_IMAGES_TABLE、CREATE_TWO_VIEW_GEOMETRIES_TABLE、CREATE_KEYPOINTS_TABLE、CREATE_MATCHES_TABLE 和 CREATE_NAME_INDEX 是用于创建数据库表的 SQL 语句。\n\nimage_ids_to_pair_id 函数用于将两个图像ID转换为匹配对ID。\n\npair_id_to_image_ids 函数用于将匹配对ID转换为两个图像ID。\n\narray_to_blob 函数用于将数组转换为二进制数据。\n\nblob_to_array 函数用于将二进制数据转换为数组。\n\nadd_camera 方法用于向数据库中添加相机信息。\n\nadd_image 方法用于向数据库中添加图像信息。\n\nadd_keypoints 方法用于向数据库中添加关键点信息。\n\nadd_descriptors 方法用于向数据库中添加描述符信息。\n\nadd_matches 方法用于向数据库中添加匹配信息。\n\nadd_two_view_geometry 方法用于向数据库中添加两视图几何信息。\n\n该脚本提供了一组用于操作 COLMAP 数据库的方法，用于添加相机、图像、关键点、描述符、匹配等数据到数据库中，并提供了一些辅助函数用于数据转换和ID转换。","metadata":{}},{"cell_type":"code","source":"# Code to interface DISK with Colmap.\n# Forked from https://github.com/cvlab-epfl/disk/blob/37f1f7e971cea3055bb5ccfc4cf28bfd643fa339/colmap/h5_to_db.py\n\n# Copyright [2020] [Michał Tyszkiewicz, Pascal Fua, Eduard Trulls]\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os, argparse, h5py, warnings\nimport numpy as np\nfrom tqdm import tqdm\nfrom PIL import Image, ExifTags\n\n\ndef get_focal(image_path, err_on_default=False):\n image = Image.open(image_path)\n max_size = max(image.size)\n\n exif = image.getexif()\n focal = None\n if exif is not None:\n focal_35mm = None\n # https://github.com/colmap/colmap/blob/d3a29e203ab69e91eda938d6e56e1c7339d62a99/src/util/bitmap.cc#L299\n for tag, value in exif.items():\n focal_35mm = None\n if ExifTags.TAGS.get(tag, None) == 'FocalLengthIn35mmFilm':\n focal_35mm = float(value)\n break\n\n if focal_35mm is not None:\n focal = focal_35mm / 35. * max_size\n \n if focal is None:\n if err_on_default:\n raise RuntimeError(\"Failed to find focal length\")\n\n # failed to find it in exif, use prior\n FOCAL_PRIOR = 1.2\n focal = FOCAL_PRIOR * max_size\n\n return focal\n#get_focal 函数用于从图像文件中获取焦距信息。它接受图像文件路径 image_path 和一个可选的 err_on_default 参数，默认为 False。\n#函数打开图像文件，并获取其尺寸信息。然后，尝试从图像的EXIF数据中获取焦距信息。如果成功获取到焦距，根据相对于35mm胶片的焦距进行换算，\n#得到相对于图像尺寸的焦距。如果无法获取焦距信息，根据默认的先验值计算一个焦距。最后，返回计算得到的焦距。\n\n\ndef create_camera(db, image_path, camera_model):\n image = Image.open(image_path)\n width, height = image.size\n\n focal = get_focal(image_path)\n\n if camera_model == 'simple-pinhole':\n model = 0 # simple pinhole\n param_arr = np.array([focal, width / 2, height / 2])\n if camera_model == 'pinhole':\n model = 1 # pinhole\n param_arr = np.array([focal, focal, width / 2, height / 2])\n elif camera_model == 'simple-radial':\n model = 2 # simple radial\n param_arr = np.array([focal, width / 2, height / 2, 0.1])\n elif camera_model == 'opencv':\n model = 4 # opencv\n param_arr = np.array([focal, focal, width / 2, height / 2, 0., 0., 0., 0.])\n \n return db.add_camera(model, width, height, param_arr)\n#create_camera 函数用于向数据库中添加相机信息。它接受一个数据库对象 db，图像文件路径 image_path 和相机模型 camera_model。\n#函数打开图像文件，并获取图像的宽度和高度。然后，调用 get_focal 函数获取图像的焦距。根据相机模型类型，选择相应的相机模型ID和参数数组。\n#最后，调用数据库的 add_camera 方法添加相机信息，并返回相机ID。\n\ndef add_keypoints(db, h5_path, image_path, img_ext, camera_model, single_camera = True):\n keypoint_f = h5py.File(os.path.join(h5_path, 'keypoints.h5'), 'r')\n\n camera_id = None\n fname_to_id = {}\n for filename in tqdm(list(keypoint_f.keys())):\n keypoints = keypoint_f[filename][()]\n\n fname_with_ext = filename# + img_ext\n path = os.path.join(image_path, fname_with_ext)\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(fname_with_ext, camera_id)\n fname_to_id[filename] = image_id\n\n db.add_keypoints(image_id, keypoints)\n\n return fname_to_id\n#add_keypoints 函数用于向数据库中添加关键点信息。它接受一个数据库对象 db，HDF5文件路径 h5_path，图像文件路径 image_path，\n#图像文件扩展名 img_ext，相机模型 camera_model 和一个可选的 single_camera 参数，默认为 True。\n\n#函数打开关键点的HDF5文件，并遍历文件中的每个关键点数据。对于每个关键点数据，构建关键点文件名（包括扩展名），\n#并与图像文件路径进行拼接得到图像文件的完整路径。如果图像文件不存在，抛出异常。\n\n#如果是第一次添加关键点数据或者不使用单一相机模型，调用 create_camera 函数创建相机，并返回相机ID。然后，\n#调用数据库的 add_image 方法添加图像信息，并将图像文件名与图像ID进行映射保存。最后，调用数据库的 add_keypoints 方法添加关键点信息。\n\n#返回关键点文件名到图像ID的映射字典 fname_to_id。\n\ndef add_matches(db, h5_path, fname_to_id):\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 warnings.warn(f'Pair {pair_id} ({id_1}, {id_2}) already added!')\n continue\n \n matches = group[key_2][()]\n db.add_matches(id_1, id_2, matches)\n\n added.add(pair_id)\n\n pbar.update(1)\n# 这段代码用于将DISK的结果与COLMAP进行接口对接，下面对其进行注释和解释：\n\ndef get_focal(image_path, err_on_default=False):\n image = Image.open(image_path)\n max_size = max(image.size)\n\n exif = image.getexif()\n focal = None\n if exif is not None:\n focal_35mm = None\n for tag, value in exif.items():\n focal_35mm = None\n if ExifTags.TAGS.get(tag, None) == 'FocalLengthIn35mmFilm':\n focal_35mm = float(value)\n break\n\n if focal_35mm is not None:\n focal = focal_35mm / 35. * max_size\n\n if focal is None:\n if err_on_default:\n raise RuntimeError(\"Failed to find focal length\")\n\n FOCAL_PRIOR = 1.2\n focal = FOCAL_PRIOR * max_size\n\n return focal\n\n#get_focal 函数用于从图像文件中获取焦距信息。它接受图像文件路径 image_path 和一个可选的 err_on_default 参数，默认为 False。\n\n#函数打开图像文件，并获取其尺寸信息。然后，尝试从图像的EXIF数据中获取焦距信息。如果成功获取到焦距，根据相对于35mm胶片的焦距进行换算，得到相对于图像尺寸的焦距。如果无法获取焦距信息，根据默认的先验值计算一个焦距。最后，返回计算得到的焦距。\n\ndef create_camera(db, image_path, camera_model):\n image = Image.open(image_path)\n width, height = image.size\n\n focal = get_focal(image_path)\n\n if camera_model == 'simple-pinhole':\n model = 0\n param_arr = np.array([focal, width / 2, height / 2])\n if camera_model == 'pinhole':\n model = 1\n param_arr = np.array([focal, focal, width / 2, height / 2])\n elif camera_model == 'simple-radial':\n model = 2\n param_arr = np.array([focal, width / 2, height / 2, 0.1])\n elif camera_model == 'opencv':\n model = 4\n param_arr = np.array([focal, focal, width / 2, height / 2, 0., 0., 0., 0.])\n\n return db.add_camera(model, width, height, param_arr)\n#create_camera 函数用于向数据库中添加相机信息。它接受一个数据库对象 db，图像文件路径 image_path 和相机模型 camera_model。\n\n#函数打开图像文件，并获取图像的宽度和高度。然后，调用 get_focal 函数获取图像的焦距。根据相机模型类型，选择相应的相机模型ID和参数数组。最后，调用数据库的 add_camera 方法添加相机信息，并返回相机ID。\n\ndef add_keypoints(db, h5_path, image_path, img_ext, camera_model, single_camera=True):\n keypoint_f = h5py.File(os.path.join(h5_path, 'keypoints.h5'), 'r')\n\n camera_id = None\n fname_to_id = {}\n for filename in tqdm(list(keypoint_f.keys())):\n keypoints = keypoint_f[filename][()]\n\n fname_with_ext = filename # + img_ext\n path = os.path.join(image_path, fname_with_ext)\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(fname_with_ext, 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#add_keypoints 函数用于向数据库中添加关键点信息。它接受一个数据库对象 db，HDF5文件路径 h5_path，图像文件路径 image_path，图像文件扩展名 img_ext，相机模型 camera_model 和一个可选的 single_camera 参数，默认为 True。\n\n#函数打开关键点的HDF5文件，并遍历文件中的每个关键点数据。对于每个关键点数据，构建关键点文件名（包括扩展名），并与图像文件路径进行拼接得到图像文件的完整路径。如果图像文件不存在，抛出异常。\n\n#如果是第一次添加关键点数据或者不使用单一相机模型，调用 create_camera 函数创建相机，并返回相机ID。然后，调用数据库的 add_image 方法添加图像信息，并将图像文件名与图像ID进行映射保存。最后，调用数据库的 add_keypoints 方法添加关键点信息。\n\n#返回关键点文件名到图像ID的映射字典 fname_to_id。\n\n\ndef add_matches(db, h5_path, fname_to_id):\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 warnings.warn(f'Pair {pair_id} ({id_1}, {id_2}) already added!')\n continue\n\n matches = group[key_2][()]\n db.add_matches(id_1, id_2, matches)\n\n added.add(pair_id)\n\n pbar.update(1)\n#add_matches 函数用于向数据库中添加匹配信息。它接受一个数据库对象 db，HDF5文件路径 h5_path 和图像文件名到图像ID的映射字典 fname_to_id。\n\n#函数打开匹配的HDF5文件，并遍历文件中的每对匹配数据。对于每对匹配数据，获取对应的图像ID。使用图像ID计算匹配对ID，\n#并检查是否已经添加过该匹配对。如果已经添加过，则发出警告并跳过。\n\n#如果匹配对没有被添加过，则获取匹配数据，并调用数据库的 add_matches 方法添加匹配信息。\n\n#最后，更新进度条并完成匹配添加。\n\n#该代码用于将DISK的关键点和匹配结果添加到COLMAP的数据库中，以便进行后续的三维重建。\n\n","metadata":{"execution":{"iopub.status.busy":"2023-05-24T01:19:12.218424Z","iopub.execute_input":"2023-05-24T01:19:12.219363Z","iopub.status.idle":"2023-05-24T01:19:12.262364Z","shell.execute_reply.started":"2023-05-24T01:19:12.219332Z","shell.execute_reply":"2023-05-24T01:19:12.261329Z"},"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 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 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')['state_dict']\n detector.model.load_state_dict(kn_weights)\n affnet_weights = torch.load('/kaggle/input/kornia-local-feature-weights/AffNet.pth')['state_dict']\n detector.aff.load_state_dict(affnet_weights)\n \n hardnet = KF.HardNet(False).eval()\n hn_weights = torch.load('/kaggle/input/kornia-local-feature-weights/HardNetLib.pth')['state_dict']\n hardnet.load_state_dict(hn_weights)\n descriptor = KF.LAFDescriptor(hardnet, patch_size=32, grayscale_descriptor=True).to(device)\n super().__init__(detector, descriptor, scale_laf)","metadata":{"execution":{"iopub.status.busy":"2023-05-24T01:19:12.263874Z","iopub.execute_input":"2023-05-24T01:19:12.264973Z","iopub.status.idle":"2023-05-24T01:19:12.279596Z","shell.execute_reply.started":"2023-05-24T01:19:12.264897Z","shell.execute_reply":"2023-05-24T01:19:12.278511Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"这段代码定义了一个名为KeyNetAffNetHardNet的类，继承自KF.LocalFeature类。它实现了KeyNet检测器 + AffNet + HardNet描述符的功能。\n\n构造函数__init__接受一些参数，包括num_features（特征数量）、upright（是否使用直立（upright）的检测器，默认为False）、device（设备，默认为torch.device('cpu')）和scale_laf（LAF（局部仿射帧）的尺度，默认为1.0）。\n\n在构造函数中，首先根据upright参数选择使用KF.PassLAF()或KF.LAFOrienter(angle_detector=KF.OriNet(False)).eval()作为角度检测器。如果不是直立的检测器，加载预训练的角度检测器模型的权重。\n\n然后，创建KF.KeyNetDetector对象，设置其参数，并加载预训练的KeyNet模型和AffNet模型的权重。\n\n接下来，创建KF.HardNet对象，并加载预训练的HardNet模型的权重。\n\n最后，创建KF.LAFDescriptor对象，设置其参数，并将HardNet模型作为描述符。\n\n该类封装了KeyNet检测器、AffNet形状估计器和HardNet描述符，方便使用并提供局部特征提取功能。","metadata":{}},{"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 if LOCAL_FEATURE == 'DISK':\n # Load DISK from Kaggle models so it can run when the notebook is offline.\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 if LOCAL_FEATURE == 'KeyNetAffNetHardNet':\n feature = KeyNetAffNetHardNet(num_feats, 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 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 print(f'Resized {timg.shape} to {timg_resized.shape} (resize_small_edge_to={resize_small_edge_to})')\n h, w = timg_resized.shape[2:]\n if LOCAL_FEATURE == 'DISK':\n features = disk(timg_resized, num_feats, pad_if_not_divisible=True)[0]\n kps1, descs = features.keypoints, features.descriptors\n \n lafs = KF.laf_from_center_scale_ori(kps1[None], torch.ones(1, len(kps1), 1, 1, device=device))\n if LOCAL_FEATURE == 'KeyNetAffNetHardNet':\n lafs, resps, descs = feature(K.color.rgb_to_grayscale(timg_resized))\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 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\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-05-24T01:19:12.448697Z","iopub.execute_input":"2023-05-24T01:19:12.449036Z","iopub.status.idle":"2023-05-24T01:19:12.511526Z","shell.execute_reply.started":"2023-05-24T01:19:12.449003Z","shell.execute_reply":"2023-05-24T01:19:12.510436Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"这段代码实现了一系列函数用于进行特征检测、特征匹配和将特征导入到COLMAP数据库中。\n\ndetect_features函数用于检测图像中的特征。它接受一系列图像文件名作为输入，然后使用指定的局部特征提取算法（如DISK或KeyNetAffNetHardNet）提取每个图像的局部特征。提取的特征包括局部仿射帧（LAF）、关键点坐标和描述符。提取的特征将保存在指定的特征目录中的HDF5文件中。\n\nget_unique_idxs函数用于获取张量中唯一元素的第一次出现的索引。\n\nmatch_features函数用于特征匹配。它接受图像文件名列表和图像对的索引对作为输入。对于每对图像，它从特征目录中加载对应的局部特征，并使用指定的匹配算法（如SMNN或Adalam）进行特征匹配。匹配的结果将保存在特征目录中的HDF5文件中。\n\nmatch_loftr函数是特定于LoFTR模型的特征匹配函数。它使用LoFTR模型进行特征匹配，并将匹配结果保存在特征目录中的HDF5文件中。\n\nimport_into_colmap函数用于将特征导入到COLMAP数据库中。它连接到数据库，创建必要的表格，并将特征从特征目录中的HDF5文件导入到数据库中。\n\n这些函数可以按照特定的顺序调用，以便进行特征检测、特征匹配和导入到COLMAP数据库的完整流程","metadata":{}},{"cell_type":"code","source":"src = '/kaggle/input/image-matching-challenge-2023'","metadata":{"execution":{"iopub.status.busy":"2023-05-24T01:19:12.514010Z","iopub.execute_input":"2023-05-24T01:19:12.514405Z","iopub.status.idle":"2023-05-24T01:19:12.524732Z","shell.execute_reply.started":"2023-05-24T01:19:12.514366Z","shell.execute_reply":"2023-05-24T01:19:12.523670Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Get data from csv.\n\ndata_dict = {}\nwith open(f'{src}/sample_submission.csv', 'r') as f:\n for i, l in enumerate(f):\n # Skip header.\n if l and i > 0:\n image, dataset, scene, _, _ = l.strip().split(',')\n if dataset not in data_dict:\n data_dict[dataset] = {}\n if scene not in data_dict[dataset]:\n data_dict[dataset][scene] = []\n data_dict[dataset][scene].append(image)","metadata":{"execution":{"iopub.status.busy":"2023-05-24T01:19:12.527294Z","iopub.execute_input":"2023-05-24T01:19:12.528239Z","iopub.status.idle":"2023-05-24T01:19:12.540830Z","shell.execute_reply.started":"2023-05-24T01:19:12.528198Z","shell.execute_reply":"2023-05-24T01:19:12.539725Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"这段代码从CSV文件中获取数据，并将其存储在data_dict字典中。","metadata":{}},{"cell_type":"code","source":"for 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-05-24T01:19:12.542541Z","iopub.execute_input":"2023-05-24T01:19:12.542932Z","iopub.status.idle":"2023-05-24T01:19:12.549538Z","shell.execute_reply.started":"2023-05-24T01:19:12.542894Z","shell.execute_reply":"2023-05-24T01:19:12.548360Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"该代码段用于遍历data_dict字典，打印每个数据集和场景中图像的数量。它会输出每个数据集和场景的名称以及对应的图像数量。这样可以对数据集和场景中的图像数量进行可视化和统计。","metadata":{}},{"cell_type":"code","source":"out_results = {}\ntimings = {\"shortlisting\":[],\n \"feature_detection\": [],\n \"feature_matching\":[],\n \"RANSAC\": [],\n \"Reconstruction\": []}","metadata":{"execution":{"iopub.status.busy":"2023-05-24T01:19:12.552808Z","iopub.execute_input":"2023-05-24T01:19:12.553600Z","iopub.status.idle":"2023-05-24T01:19:12.559057Z","shell.execute_reply.started":"2023-05-24T01:19:12.553562Z","shell.execute_reply":"2023-05-24T01:19:12.557750Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"out_results是一个空的字典，用于存储输出结果。\n\ntimings是一个包含不同阶段时间统计的字典。它包含以下键：\n\n\"shortlisting\"：用于存储图像筛选阶段的时间。\n\"feature_detection\"：用于存储特征检测阶段的时间。\n\"feature_matching\"：用于存储特征匹配阶段的时间。\n\"RANSAC\"：用于存储RANSAC阶段的时间。\n\"Reconstruction\"：用于存储重建阶段的时间。\n这些键对应的值是空列表，用于存储每个阶段的时间统计。","metadata":{}},{"cell_type":"code","source":"# Function to create a submission file.\ndef create_submission(out_results, data_dict):\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 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 f.write(f'{image},{dataset},{scene},{arr_to_str(R)},{arr_to_str(T)}\\n')","metadata":{"execution":{"iopub.status.busy":"2023-05-24T01:19:12.560803Z","iopub.execute_input":"2023-05-24T01:19:12.561607Z","iopub.status.idle":"2023-05-24T01:19:12.571869Z","shell.execute_reply.started":"2023-05-24T01:19:12.561569Z","shell.execute_reply":"2023-05-24T01:19:12.570832Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"create_submission是一个用于创建提交文件的函数。\n\n它接受两个参数：\n\nout_results：存储了匹配和重建结果的字典。\ndata_dict：存储了待处理图像的数据字典。\n函数会打开一个名为submission.csv的文件，并写入CSV文件的头部信息。然后，它会遍历data_dict中的每个数据集和场景，并检查out_results中是否有对应的结果。如果有结果，则从结果中提取旋转矩阵和平移向量。如果没有结果，则使用单位矩阵和零向量作为默认值。最后，将图像路径、数据集、场景、旋转矩阵和平移向量写入CSV文件中。\n\n请注意，arr_to_str是一个将NumPy数组转换为字符串的辅助函数。","metadata":{}},{"cell_type":"code","source":"gc.collect()\ndatasets = []\nfor dataset in data_dict:\n datasets.append(dataset)\n\nfor 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}/test/{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}/test/{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.45, # should be strict\n min_pairs =35, # we select at least min_pairs PER IMAGE with biggest similarity\n exhaustive_if_less = 20,\n device=device)\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 != 'LoFTR':\n detect_features(img_fnames, \n 2048,\n feature_dir=feature_dir,\n upright=True,\n device=device,\n resize_small_edge_to=600\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)\n else:\n match_loftr(img_fnames, index_pairs, feature_dir=feature_dir, device=device, resize_to_=(600, 800))\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)\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 print(maps)\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 print (\"Looking for the best reconstruction\")\n if isinstance(maps, dict):\n for idx1, rec in maps.items():\n print (idx1, rec.summary())\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 print (maps[best_idx].summary())\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-05-24T01:19:12.573400Z","iopub.execute_input":"2023-05-24T01:19:12.574122Z","iopub.status.idle":"2023-05-24T01:19:12.754585Z","shell.execute_reply.started":"2023-05-24T01:19:12.574084Z","shell.execute_reply":"2023-05-24T01:19:12.753395Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"该代码段中的循环遍历数据集和场景，并在每个场景中执行以下操作：\n\n检查输出结果字典out_results中是否存在当前数据集和场景的条目，如果不存在，则创建空字典作为条目。\n获取图像文件路径列表img_fnames。\n创建特征存储目录feature_dir。\n进行图像对的筛选，得到用于匹配的图像对索引index_pairs。\n检测特征并进行匹配。\n如果使用的是非LoFTR本地特征，则先进行特征检测，然后进行特征匹配。\n如果使用的是LoFTR本地特征，则直接进行特征匹配。\n创建Colmap数据库并将特征导入。\n执行RANSAC算法进行几何一致性估计。\n执行重建过程，生成稀疏点云。\n查找生成的重建中包含最多图像的模型，并提取该模型中每个图像的旋转矩阵和平移向量。\n将图像的旋转矩阵和平移向量存储到out_results字典中。\n创建提交文件。\n请注意，代码中使用了try-except块来捕获任何可能的错误，并在出现错误时继续执行下一个场景的处理。这样可以确保即使在处理某些场景时出现问题，代码也能继续运行并处理其他场景。\n\n另外，gc.collect()语句用于手动触发垃圾回收，以释放内存。timings字典用于记录每个步骤的运行时间。","metadata":{}},{"cell_type":"code","source":"create_submission(out_results, data_dict)","metadata":{"execution":{"iopub.status.busy":"2023-05-24T01:19:12.756534Z","iopub.execute_input":"2023-05-24T01:19:12.757204Z","iopub.status.idle":"2023-05-24T01:19:12.764786Z","shell.execute_reply.started":"2023-05-24T01:19:12.757164Z","shell.execute_reply":"2023-05-24T01:19:12.763615Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"create_submission函数用于创建提交文件，将out_results中的结果以CSV格式写入到submission.csv文件中。\n\n函数的输入参数包括out_results和data_dict。\n\nout_results是一个字典，包含了每个数据集和场景中的图像的旋转矩阵和平移向量信息。\n\ndata_dict是一个字典，包含了每个数据集和场景中的图像文件路径信息。\n\n函数会遍历data_dict中的每个数据集和场景，并在submission.csv文件中写入每个图像的信息，包括图像路径、数据集、场景、旋转矩阵和平移向量。\n\n函数的输出是一个submission.csv文件，其中包含了所有图像的信息，可以用于提交比赛结果。","metadata":{}}]}