{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":91498,"databundleVersionId":11655853,"sourceType":"competition"},{"sourceId":7884485,"sourceType":"datasetVersion","datasetId":4628051},{"sourceId":8026384,"sourceType":"datasetVersion","datasetId":4726252},{"sourceId":11924468,"sourceType":"datasetVersion","datasetId":6988459},{"sourceId":176463227,"sourceType":"kernelVersion"},{"sourceId":4534,"sourceType":"modelInstanceVersion","modelInstanceId":3326,"modelId":986},{"sourceId":17191,"sourceType":"modelInstanceVersion","modelInstanceId":14317,"modelId":21716},{"sourceId":17555,"sourceType":"modelInstanceVersion","modelInstanceId":14611,"modelId":22086}],"dockerImageVersionId":30919,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Example submission\n\nImage Matching Challenge 2025: https://www.kaggle.com/competitions/image-matching-challenge-2025\n\nThis notebook creates a simple submission using ALIKED and LightGlue, plus DINO for shortlisting, on GPU. Adapted from [last year](https://www.kaggle.com/code/oldufo/imc-2024-submission-example).\n\nRemember to select an accelerator on the sidebar to the right, and to disable internet access when submitting a notebook to the competition.","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19"}},{"cell_type":"code","source":"# IMPORTANT \n#Install dependencies and copy model weights to run the notebook without internet access when submitting to the competition.\n\n!pip install --no-index /kaggle/input/imc2024-packages-lightglue-rerun-kornia/* --no-deps\n!mkdir -p /root/.cache/torch/hub/checkpoints\n!cp /kaggle/input/aliked/pytorch/aliked-n16/1/aliked-n16.pth /root/.cache/torch/hub/checkpoints/\n!cp /kaggle/input/lightglue/pytorch/aliked/1/aliked_lightglue.pth /root/.cache/torch/hub/checkpoints/\n!cp /kaggle/input/lightglue/pytorch/aliked/1/aliked_lightglue.pth /root/.cache/torch/hub/checkpoints/aliked_lightglue_v0-1_arxiv-pth","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-18T08:43:22.430203Z","iopub.execute_input":"2025-06-18T08:43:22.430561Z","iopub.status.idle":"2025-06-18T08:43:27.652454Z","shell.execute_reply.started":"2025-06-18T08:43:22.430532Z","shell.execute_reply":"2025-06-18T08:43:27.651472Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!python -m pip install --no-index --find-links=/kaggle/input/pkg-check-orientation/ check_orientation==0.0.5\n\n!mkdir -p /root/.cache/torch/hub/checkpoints/","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-31T03:17:44.491912Z","iopub.execute_input":"2025-05-31T03:17:44.492318Z","iopub.status.idle":"2025-05-31T03:17:50.926647Z","shell.execute_reply.started":"2025-05-31T03:17:44.492284Z","shell.execute_reply":"2025-05-31T03:17:50.925604Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!cp /kaggle/input/pkg-check-orientation/2020-11-16_resnext50_32x4d.zip /root/.cache/torch/hub/checkpoints/","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-31T03:17:50.92863Z","iopub.execute_input":"2025-05-31T03:17:50.928969Z","iopub.status.idle":"2025-05-31T03:17:51.688381Z","shell.execute_reply.started":"2025-05-31T03:17:50.928938Z","shell.execute_reply":"2025-05-31T03:17:51.687515Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import sys\nimport os\nfrom tqdm import tqdm\nfrom time import time, sleep\nimport gc\nimport numpy as np\nimport h5py\nimport dataclasses\nimport pandas as pd\nfrom IPython.display import clear_output\nfrom collections import defaultdict\nfrom copy import deepcopy\nfrom PIL import Image\n\nimport cv2\nimport torch\nimport torch.nn.functional as F\nimport kornia as K\nimport kornia.feature as KF\n\nimport torch\nfrom lightglue import match_pair\nfrom lightglue import ALIKED, LightGlue\nfrom lightglue.utils import load_image, rbd\n\nfrom transformers import AutoImageProcessor, AutoModel\n\n# IMPORTANT Utilities: importing data into colmap and competition metric\nimport pycolmap\nsys.path.append('/kaggle/input/imc25-utils')\nfrom database import *\nfrom h5_to_db import *\nimport metric\n\nprint(\"============ Fine tuning Parametesr ==============\")\n\n# Do not forget to select an accelerator on the sidebar to the right.\ndevice = K.utils.get_cuda_device_if_available(0)\nprint(f'{device=}')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-31T03:17:51.689669Z","iopub.execute_input":"2025-05-31T03:17:51.689894Z","iopub.status.idle":"2025-05-31T03:18:13.271602Z","shell.execute_reply.started":"2025-05-31T03:17:51.689875Z","shell.execute_reply":"2025-05-31T03:18:13.27079Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 该函数的作用是加载指定路径的图像文件，并将其转换为 PyTorch 的张量格式。\n# 图像加载时会被处理为 32 位 RGB 格式。\n# 函数支持将图像加载到指定的设备（如 CPU 或 GPU）。\ndef load_torch_image(fname, device=torch.device('cpu')):\n    \"\"\"\n    加载图像文件并返回 PyTorch 张量形式的图像。\n\n    参数:\n        fname (str): 图像文件路径。\n        device (torch.device): 指定加载图像时的设备（默认为 CPU）。\n\n    返回:\n        torch.Tensor: 一个 4 维张量，形状为 (1, C, H, W)，\n                      其中 C 是通道数，H 和 W 分别是高度和宽度。\n    \"\"\"\n    # 使用 Kornia 的 `load_image` 函数加载图像\n    # 参数解释:\n    # - fname: 图像文件路径。\n    # - K.io.ImageLoadType.RGB32: 指定图像加载为 32 位深度的 RGB 格式 (3 通道)。\n    # - device: 指定加载图像的设备 (CPU 或 GPU)。\n    img = K.io.load_image(fname, K.io.ImageLoadType.RGB32, device=device)[None, ...]\n\n    # 在图像张量的第 0 维增加一个新的维度，表示 batch_size = 1\n    # Kornia 加载的图像通常是 3 维张量 (C, H, W)，\n    # 增加维度后变为 4 维张量 (1, C, H, W)，以符合神经网络的输入规范。\n    \n    return img\n    # 返回 4 维张量 (1, C, H, W)。\n\n\n# 这段代码定义了一个函数 get_global_desc，用于从一组图像文件路径中提取全局特征描述符（global descriptors）。\n# 它使用了一个预训练的视觉模型（可能是 DINOv2）来处理输入图像并生成特征表示。\ndef get_global_desc(fnames, device=torch.device('cpu')):\n    \"\"\"\n    提取图像的全局特征描述符。\n\n    参数:\n        fnames (list of str): 图像文件路径列表。\n        device (torch.device): 指定模型运行的设备（默认为 CPU）。\n\n    返回:\n        torch.Tensor: 所有图像的全局特征描述符，形状为 (N, D)，\n                      其中 N 是图像数量，D 是特征维度。\n    \"\"\"\n\n    # 加载预训练的图像处理器和模型\n    processor = AutoImageProcessor.from_pretrained('/kaggle/input/dinov2/pytorch/base/1')\n    model = AutoModel.from_pretrained('/kaggle/input/dinov2/pytorch/base/1')\n\n    # 设置模型为推理模式并放置到指定设备上\n    model = model.eval()\n    model = model.to(device)\n\n    # 用于存储所有图像的全局特征描述符\n    global_descs_dinov2 = []\n\n    # 遍历所有图像文件\n    for i, img_fname_full in tqdm(enumerate(fnames), total=len(fnames)):\n        \"\"\"\n        循环每张图像：\n        - 加载图像\n        - 使用预处理器和模型生成全局特征\n        \"\"\"\n        # 获取图像文件的名称（去掉路径和扩展名）\n        key = os.path.splitext(os.path.basename(img_fname_full))[0]\n\n        # 加载图像为 PyTorch 张量 (1, C, H, W)\n        timg = load_torch_image(img_fname_full)\n\n        # 在推理模式下执行模型前向传播\n        with torch.inference_mode():\n            # 使用预处理器处理图像，并将结果移动到指定设备\n            inputs = processor(images=timg, return_tensors=\"pt\", do_rescale=False).to(device)\n\n            # 前向传播，提取模型的输出\n            outputs = model(**inputs)\n\n            # 提取全局特征：\n            # - `outputs.last_hidden_state` 是模型的最后一层输出\n            # - `[:, 1:]` 去掉 [CLS] token 的特征，仅保留其他特征\n            # - `.max(dim=1)` 在 tokens 维度上取最大值（生成全局特征）\n            # - `F.normalize` 对特征进行 L2 归一化\n            dino_mac = F.normalize(outputs.last_hidden_state[:, 1:].max(dim=1)[0], dim=1, p=2)\n\n        # 将特征从 GPU 移动到 CPU，并添加到列表中\n        global_descs_dinov2.append(dino_mac.detach().cpu())\n\n    # 将所有特征按维度 0 连接成一个张量，形状为 (N, D)\n    global_descs_dinov2 = torch.cat(global_descs_dinov2, dim=0)\n\n    # 返回完整的特征描述符张量\n    return global_descs_dinov2\n\n\n# 这段代码定义了一个函数 get_img_pairs_exhaustive，其作用是从输入的图像文件列表中生成所有可能的图像对（不重复且有序）。\ndef get_img_pairs_exhaustive(img_fnames):\n    \"\"\"\n    生成输入图像文件列表的所有可能索引对。\n\n    参数:\n        img_fnames (list): 图像文件名列表。\n\n    返回:\n        list of tuple: 包含所有可能的图像索引对 (i, j)，其中 0 <= i < j < len(img_fnames)。\n    \"\"\"\n    # 用于存储所有图像对索引的列表\n    index_pairs = []\n\n    # 遍历每个图像索引\n    for i in range(len(img_fnames)):\n        # 从 i+1 开始，避免重复和自比较\n        for j in range(i + 1, len(img_fnames)):\n            # 将索引对 (i, j) 添加到结果列表\n            index_pairs.append((i, j))\n\n    # 返回生成的索引对列表\n    return index_pairs\n\n\n# 这段代码定义了一个函数 get_image_pairs_shortlist，用于从图像文件名列表中生成图像对的索引列表，目标是根据图像间的相似性选择具有一定相似度的图像对，避免穷举所有可能的图像对。\ndef get_image_pairs_shortlist(fnames,\n                              sim_th=0.6,  # 相似度阈值，越小越严格\n                              min_pairs=30,  # 每张图像至少匹配的最小对数\n                              exhaustive_if_less=20,  # 图像数量小于此值时直接穷举\n                              device=torch.device('cpu')):  # 指定计算设备\n    \"\"\"\n    根据图像的全局描述符生成一组相似的图像对索引。\n\n    参数:\n        fnames (list of str): 图像文件名列表。\n        sim_th (float): 相似性阈值，欧氏距离小于此值的图像对会被选中。\n        min_pairs (int): 每张图像至少匹配的对数。\n        exhaustive_if_less (int): 如果图像数量小于此值，则使用穷举法生成图像对。\n        device (torch.device): 指定用 CPU 或 GPU 进行计算。\n\n    返回:\n        list of tuple: 图像索引对列表，按顺序排列。\n    \"\"\"\n    # 获取图像数量\n    num_imgs = len(fnames)\n\n    # 如果图像数量小于等于 `exhaustive_if_less`，直接使用穷举法生成图像对\n    if num_imgs <= exhaustive_if_less:\n        return get_img_pairs_exhaustive(fnames)\n\n    # 提取图像的全局描述符\n    descs = get_global_desc(fnames, device=device)\n\n    # 计算描述符间的欧氏距离矩阵\n    dm = torch.cdist(descs, descs, p=2).detach().cpu().numpy()\n\n    \"\"\"\n    距离矩阵解释：\n    - dm 是一个二维矩阵，形状为 (num_imgs, num_imgs)。\n    - dm[i, j] 表示第 i 张图像和第 j 张图像的欧氏距离。\n    \"\"\"\n\n    # 创建一个布尔掩码矩阵，标记相似度满足阈值的图像对\n    mask = dm <= sim_th\n\n    # 初始化变量\n    total = 0  # 总匹配对数\n    matching_list = []  # 存储最终的匹配图像索引对\n    ar = np.arange(num_imgs)  # 图像索引的数组\n    already_there_set = []  # 已处理的图像索引（未实际使用）\n\n    # 遍历距离矩阵的每一行（即对每张图像寻找匹配的图像）\n    for st_idx in range(num_imgs - 1):\n        # 获取与当前图像 st_idx 距离小于阈值的图像索引\n        mask_idx = mask[st_idx]\n        to_match = ar[mask_idx]\n\n        # 如果满足条件的图像对数少于 `min_pairs`，补充最近的图像对\n        if len(to_match) < min_pairs:\n            to_match = np.argsort(dm[st_idx])[:min_pairs]\n\n        # 遍历匹配的图像索引\n        for idx in to_match:\n            # 跳过与自身的匹配\n            if st_idx == idx:\n                continue\n\n            # 距离小于 1000 的图像对认为有效（防止异常值）\n            if dm[st_idx, idx] < 1000:\n                # 将图像对加入匹配列表，并保证索引对按升序排列\n                matching_list.append(tuple(sorted((st_idx, idx.item()))))\n                total += 1\n\n    # 去重并排序匹配列表，确保唯一性\n    matching_list = sorted(list(set(matching_list)))\n\n    # 返回最终的匹配图像对\n    return matching_list\n\n\n# 这段代码定义了一个函数 convert_coord，用于对二维坐标进行旋转变换，适应于图像的旋转操作。\ndef convert_coord(coords, width, height, rotation):\n    \"\"\"\n    坐标旋转变换\n\n    参数:\n        coords: ndarray, 形状为(N, 2)，每行为(x, y)\n        width: int, 原图宽度\n        height: int, 原图高度\n        rotation: int, 旋转方式，0, 1, 2, 3 分别表示0°、90°、180°、270°\n\n    返回:\n        ndarray, 变换后坐标，形状(N, 2)\n    \"\"\"\n    if rotation == 0:\n        # 不变换\n        return coords.copy()\n    if rotation == 1:\n        # 顺时针90°\n        x_new = width - 1 - coords[:, 1]\n        y_new = coords[:, 0]\n    elif rotation == 2:\n        # 顺时针180°\n        x_new = width - 1 - coords[:, 0]\n        y_new = height - 1 - coords[:, 1]\n    elif rotation == 3:\n        # 顺时针270°\n        x_new = coords[:, 1]\n        y_new = height - 1 - coords[:, 0]\n    else:\n        raise ValueError(\"rotation must be one of 0, 1, 2, 3.\")\n    return np.stack((x_new, y_new), axis=1)\n        \n\n\n# 这段代码定义了一个函数 detect_aliked，用于检测图像的关键点和描述符，并将结果保存到指定的文件夹中。\ndef detect_aliked(img_fnames,\n                  rots,\n                  feature_dir='.featureout',\n                  num_features=4096,\n                  resize_to=2048,\n                  device=torch.device('cpu')):\n    \"\"\"\n    使用 ALIKED 模型检测图像关键点和描述符，支持旋转图像。\n\n    参数:\n        img_fnames (list of str): 图像文件路径列表。\n        rots (list of int): 每张图像的旋转角度，以 90° 为单位（0, 1, 2, 3 分别表示 0°、90°、180°、270°）。\n        feature_dir (str): 保存特征输出的路径（默认当前目录下的 .featureout）。\n        num_features (int): 每张图像的最大关键点数量。\n        resize_to (int): 将图像调整到的最大尺寸（长边）。\n        device (torch.device): 运行设备，默认为 CPU。\n\n    返回:\n        None\n    \"\"\"\n    # 使用 float32 数据类型（避免 ALIKED 的 float16 问题）\n    dtype = torch.float32\n\n    # 初始化 ALIKED 特征提取器\n    extractor = ALIKED(\n        max_num_keypoints=num_features,  # 最大关键点数量\n        detection_threshold=0.1,        # 检测阈值\n        resize=resize_to                # 图像调整大小\n    ).eval().to(device, dtype)\n\n    # 如果特征输出目录不存在，则创建\n    if not os.path.isdir(feature_dir):\n        os.makedirs(feature_dir)\n\n    # 使用 h5py 打开三个输出文件\n    with h5py.File(f'{feature_dir}/keypoints.h5', mode='w') as f_kp, \\\n         h5py.File(f'{feature_dir}/keypoints_rot.h5', mode='w') as f_kp_rot, \\\n         h5py.File(f'{feature_dir}/descriptors.h5', mode='w') as f_desc:\n\n        # 遍历图像及其旋转角度\n        for img_path, rotated_i in tqdm(zip(img_fnames, rots)):\n            \"\"\"\n            对每张图像：\n            - 加载原始图像\n            - 按指定角度旋转\n            - 提取关键点和描述符\n            - 保存旋转后关键点、转换回原图的关键点和描述符\n            \"\"\"\n            # 提取图像文件名作为键\n            img_fname = img_path.split('/')[-1]\n            key = img_fname\n\n            # 禁用梯度计算，提高推理效率\n            with torch.inference_mode():\n                # 加载图像并转换为指定数据类型\n                img_tensor = load_torch_image(img_path, device=device).to(dtype)\n                height, width = img_tensor.shape[2], img_tensor.shape[3]  # 获取图像尺寸\n\n                # 按指定角度旋转图像\n                rotated_tensor = torch.rot90(img_tensor, rotated_i, dims=[2, 3])\n\n                # 使用 ALIKED 提取特征\n                features = extractor.extract(rotated_tensor)\n                keypoints = features['keypoints'].view(-1, 2).detach().cpu().numpy()  # 提取关键点\n                descriptors = features['descriptors'].view(keypoints.shape[0], -1).detach().cpu().numpy()  # 提取描述符\n\n                # 保存旋转后的关键点\n                f_kp_rot[key] = keypoints\n\n                # 将旋转后的关键点转换到原图坐标\n                converted_keypoints = convert_coord(keypoints, width, height, rotated_i)\n                f_kp[key] = converted_keypoints\n\n                # 保存描述符\n                f_desc[key] = descriptors\n\n\n\n# 这段代码定义了一个函数 match_with_lightglue，用于使用 LightGlue 匹配两组图像的特征点，并将匹配结果保存到 .h5 文件中。\ndef match_with_lightglue(img_fnames,\n                         index_pairs,\n                         feature_dir='.featureout',\n                         device=torch.device('cpu'),\n                         min_matches=20,\n                         verbose=True):\n    \"\"\"\n    使用 LightGlue 匹配图像对的特征点。\n\n    参数:\n        img_fnames (list of str): 图像文件路径列表。\n        index_pairs (list of tuple): 图像对索引列表，每个元素是 (idx1, idx2)。\n        feature_dir (str): 存储特征点和匹配结果的路径，默认是 .featureout。\n        device (torch.device): 运行设备，默认为 CPU。\n        min_matches (int): 最小匹配点数量，低于此数量的图像对将被忽略。\n        verbose (bool): 是否打印匹配信息。\n\n    返回:\n        None\n    \"\"\"\n    # 初始化 LightGlue 匹配器\n    lg_matcher = KF.LightGlueMatcher(\n        \"aliked\",  # 使用 ALIKED 特征描述符\n        {\n            \"width_confidence\": -1,  # 禁用宽度置信度\n            \"depth_confidence\": -1,  # 禁用深度置信度\n            \"mp\": True if 'cuda' in str(device) else False  # 启用多线程 (MP) 模式，仅在 CUDA 上启用\n        }\n    ).eval().to(device)  # 设置为评估模式，并转移到指定设备\n\n    # 打开特征存储文件\n    with h5py.File(f'{feature_dir}/keypoints.h5', mode='r') as f_kp, \\\n         h5py.File(f'{feature_dir}/keypoints_rot.h5', mode='r') as f_kp_rot, \\\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        # 遍历每对图像索引\n        for pair_idx in tqdm(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\n            # 加载旋转后的关键点和描述符数据\n            kp1_rotated = torch.from_numpy(f_kp_rot[key1][...]).to(device)  # 图像1的旋转后关键点\n            kp2_rotated = torch.from_numpy(f_kp_rot[key2][...]).to(device)  # 图像2的旋转后关键点\n            desc1 = torch.from_numpy(f_desc[key1][...]).to(device)  # 图像1的描述符\n            desc2 = torch.from_numpy(f_desc[key2][...]).to(device)  # 图像2的描述符\n\n            # 使用 LightGlue 进行特征点匹配\n            with torch.inference_mode():\n                dists, idxs = lg_matcher(\n                    desc1, desc2,\n                    KF.laf_from_center_scale_ori(kp1_rotated[None]),  # 图像1的 LAF（局部仿射框架）\n                    KF.laf_from_center_scale_ori(kp2_rotated[None])   # 图像2的 LAF\n                )\n\n            # 如果没有找到匹配点，跳过该图像对\n            if len(idxs) == 0:\n                continue\n\n            # 获取匹配点数量\n            n_matches = len(idxs)\n\n            # 打印匹配信息\n            if verbose:\n                print(f'{key1}-{key2}: {n_matches} matches')\n\n            # 创建匹配结果组\n            group = f_match.require_group(key1)\n\n            # 如果匹配点数量超过最小阈值，保存匹配结果\n            if n_matches >= min_matches:\n                group.create_dataset(key2, data=idxs.detach().cpu().numpy().reshape(-1, 2))\n\n\n# 这段代码定义了一个函数 import_into_colmap，用于将图像特征和匹配结果导入 COLMAP 数据库，为后续的三维重建或图像配准任务做准备。\ndef import_into_colmap(img_dir, feature_dir='.featureout', database_path='colmap.db'):\n    \"\"\"\n    将图像特征和匹配信息导入 COLMAP 数据库。\n\n    参数:\n        img_dir (str): 包含图像文件的目录路径。\n        feature_dir (str): 存储图像特征和匹配结果的目录。\n        database_path (str): 存储 COLMAP 数据库的路径，默认是 'colmap.db'。\n\n    返回:\n        None\n    \"\"\"\n    # 连接到 COLMAP 数据库\n    db = COLMAPDatabase.connect(database_path)\n\n    # 创建所需的表（如果尚未存在）\n    db.create_tables()\n\n    # 相机模型设置\n    single_camera = False  # 是否假设所有图像都使用同一个相机模型\n\n    # 将关键点信息添加到数据库\n    fname_to_id = add_keypoints(\n        db,  # 数据库连接\n        feature_dir,  # 特征存储目录\n        img_dir,  # 图像目录\n        camera_params='',  # 相机参数（未使用）\n        camera_model='simple-pinhole',  # 相机模型（COLMAP 支持的模型）\n        single_camera=single_camera  # 是否所有图像使用单一相机\n    )\n\n    # 将匹配结果添加到数据库\n    add_matches(\n        db,  # 数据库连接\n        feature_dir,  # 特征存储目录\n        fname_to_id  # 图像文件名到数据库 ID 的映射\n    )\n\n    # 提交更改到数据库\n    db.commit()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-31T03:18:13.272421Z","iopub.execute_input":"2025-05-31T03:18:13.273025Z","iopub.status.idle":"2025-05-31T03:18:13.294945Z","shell.execute_reply.started":"2025-05-31T03:18:13.273Z","shell.execute_reply":"2025-05-31T03:18:13.294116Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Collect vital info from the dataset\n\n@dataclasses.dataclass\nclass Prediction:\n    image_id: str | None  # A unique identifier for the row -- unused otherwise. Used only on the hidden test set.\n    dataset: str\n    filename: str\n    cluster_index: int | None = None\n    rotation: np.ndarray | None = None\n    translation: np.ndarray | None = None\n\n# Set is_train=True to run the notebook on the training data.\n# Set is_train=False if submitting an entry to the competition (test data is hidden, and different from what you see on the \"test\" folder).\nis_train = False\ndata_dir = '/kaggle/input/image-matching-challenge-2025'\nworkdir = '/kaggle/working/result/'\nos.makedirs(workdir, exist_ok=True)\n\nif is_train:\n    sample_submission_csv = os.path.join(data_dir, 'train_labels.csv')\nelse:\n    sample_submission_csv = os.path.join(data_dir, 'sample_submission.csv')\n\nsamples = {}\ncompetition_data = pd.read_csv(sample_submission_csv)\nfor _, row in competition_data.iterrows():\n    # Note: For the test data, the \"scene\" column has no meaning, and the rotation_matrix and translation_vector columns are random.\n    if row.dataset not in samples:\n        samples[row.dataset] = []\n    samples[row.dataset].append(\n        Prediction(\n            image_id=None if is_train else row.image_id,\n            dataset=row.dataset,\n            filename=row.image\n        )\n    )\n\nfor dataset in samples:\n    print(f'Dataset \"{dataset}\" -> num_images={len(samples[dataset])}')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-31T03:18:13.295752Z","iopub.execute_input":"2025-05-31T03:18:13.295998Z","iopub.status.idle":"2025-05-31T03:18:13.462805Z","shell.execute_reply.started":"2025-05-31T03:18:13.295966Z","shell.execute_reply":"2025-05-31T03:18:13.461999Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"gc.collect()\n\nmax_images = None \ndatasets_to_process = None \n\nif is_train:\n\n    \n        datasets_to_process = [\n        \t# New data.\n        \t'amy_gardens',\n        \t'ETs',\n        \t'fbk_vineyard',\n        \t'stairs',\n        ]\n\ntimings = {\n    \"shortlisting\":[],\n    \"feature_detection\": [],\n    \"feature_matching\":[],\n    \"RANSAC\": [],\n    \"Reconstruction\": [],\n}\nmapping_result_strs = []\n\nfrom torchvision.io import read_image as T_read_image\nfrom torchvision.io import ImageReadMode\nfrom torchvision import transforms as T\nfrom check_orientation.pre_trained_models import create_model\n\nfrom torch.utils.data import Dataset, DataLoader\nimport torchvision.transforms as T\nimport torch\nimport os\n\n\n# 这段代码定义了一个函数 convert_rot_k，用于将一个索引值转换为旋转的 k 值。\ndef convert_rot_k(index: int) -> int:\n    \"\"\"\n    将索引值转换为旋转 k 值。\n\n    参数:\n        index (int): 输入索引值。\n\n    返回:\n        int: 对应的旋转 k 值。\n    \"\"\"\n    # 定义索引到旋转 k 值的映射规则\n    mapping = [0, 3, 2, 1]\n\n    # 如果索引在有效范围内（0 到 3），返回映射值；否则返回默认值 1\n    return mapping[index] if 0 <= index < 4 else 1\n    \n\n\n# 这段代码定义了一个 RotationImageDataset 类，继承自 PyTorch 的 Dataset，用于加载图像数据并支持可选的图像变换功能。\nclass RotationImageDataset(Dataset):\n    \"\"\"\n    自定义数据集类，用于加载图像并支持变换。\n\n    参数:\n        file_list (list of str): 图像文件路径列表。\n        transform (callable, optional): 图像变换函数（如数据增强、归一化等）。\n    \"\"\"\n    def __init__(self, file_list, transform=None):\n        \"\"\"\n        初始化数据集。\n\n        参数:\n            file_list (list of str): 包含图像文件路径的列表。\n            transform (callable, optional): 对加载的图像应用的变换函数。\n        \"\"\"\n        self.file_list = file_list  # 存储图像文件路径列表\n        self.transform = transform  # 图像变换函数\n\n    def __len__(self):\n        \"\"\"\n        返回数据集的大小（即图像文件的数量）。\n\n        返回:\n            int: 文件列表的长度。\n        \"\"\"\n        return len(self.file_list)  # 数据集的大小等于文件列表的长度\n\n    def __getitem__(self, idx):\n        \"\"\"\n        根据索引加载图像并应用变换。\n\n        参数:\n            idx (int): 图像在数据集中的索引。\n\n        返回:\n            Tensor: 图像数据（可能经过变换）。\n        \"\"\"\n        # 获取指定索引的图像路径\n        img_path = self.file_list[idx]\n\n        # 使用 torchvision.io.read_image 加载图像（以 RGB 模式读取）\n        img = T_read_image(img_path, mode=ImageReadMode.RGB)\n\n        # 如果有变换函数，将其应用到图像上；否则返回原始图像\n        return self.transform(img) if self.transform else img\n\n# 这段代码定义了 build_rotation_dataloader 函数，用于创建一个数据加载器 (DataLoader)，将图像数据以批处理的方式加载并应用预处理操作。\ndef build_rotation_dataloader(image_paths, batch_size=1):\n    \"\"\"\n    构建一个图像数据加载器。\n\n    参数:\n        image_paths (list of str): 图像文件路径列表。\n        batch_size (int): 每批次加载的图像数量，默认是 1。\n\n    返回:\n        DataLoader: PyTorch 数据加载器。\n    \"\"\"\n    # 定义图像预处理步骤\n    preprocess = T.Compose([\n        T.Resize((224, 224)),  # 将图像调整为 224x224 的固定大小\n        T.ConvertImageDtype(torch.float),  # 将图像像素值转换为浮点数（范围为 0.0 到 1.0）\n        T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))  # 对图像进行归一化\n    ])\n    \"\"\"\n    图像预处理说明:\n    - `T.Resize((224, 224))`:\n      确保所有输入图像的大小统一为 224x224 像素。\n    - `T.ConvertImageDtype(torch.float)`:\n      将图像的像素值从整数 (0-255) 转换为浮点数 (0.0-1.0)。\n    - `T.Normalize(mean, std)`:\n      对浮点图像进行标准化处理，使其符合常见的预训练模型（如 ImageNet 模型）输入要求:\n      - 均值: `(0.485, 0.456, 0.406)`，对应 RGB 通道的均值。\n      - 标准差: `(0.229, 0.224, 0.225)`，对应 RGB 通道的标准差。\n    \"\"\"\n\n    # 创建自定义数据集，传入图像路径列表和预处理管道\n    dataset = RotationImageDataset(image_paths, transform=preprocess)\n\n    # 使用 PyTorch 的 DataLoader 封装数据集\n    return DataLoader(\n        dataset,              # 自定义数据集对象\n        batch_size=batch_size,  # 每批加载的图像数量\n        shuffle=False,          # 是否打乱数据顺序（这里设置为 False）\n        pin_memory=True,        # 是否将数据加载到固定内存中，提升 GPU 数据传输速度\n        num_workers=2,          # 使用 2 个子线程加载数据\n        drop_last=False         # 如果最后一批数据不足 batch_size，是否丢弃（这里保留）\n    )\n\n\n\n# 该代码定义了一个函数，用于执行图像旋转检测任务，使用一个预训练模型预测图像的旋转角度，并返回每张图像的旋转参数。\ndef exec_rotation_detection(img_files, device):\n    \"\"\"\n    执行图像旋转检测任务。\n\n    参数:\n        img_files (list of str): 图像文件路径列表。\n        device (torch.device): 指定运行设备（如 CPU 或 GPU）。\n\n    返回:\n        list of int: 每张图像的旋转参数 rot_k（0, 1, 2, 3）。\n    \"\"\"\n    # 创建模型并加载预训练权重\n    model = create_model(\"swsl_resnext50_32x4d\")  # 使用预定义的模型架构 swsl_resnext50_32x4d\n    model.eval()  # 设置模型为评估模式，关闭 dropout 和 batchnorm 的训练行为\n    model.to(device)  # 将模型移动到指定设备（如 GPU 或 CPU）\n\n    # 使用自定义函数构建图像数据加载器\n    dataloader = build_rotation_dataloader(img_files)\n\n    # 存储每张图像预测的旋转参数 rot_k\n    rotations = []\n\n    # 遍历数据加载器中的每个批次图像\n    for i, img in enumerate(dataloader):\n        # 将图像数据转换为 float32 类型并移动到指定设备\n        img = img.to(torch.float32).to(device)\n\n        # 禁用梯度计算以节省内存并加速推理\n        with torch.no_grad():\n            # 使用模型对当前图像进行前向传播，获取预测结果\n            pred = model(img).cpu().numpy()  # 将预测结果移动到 CPU 并转换为 NumPy 数组\n\n            # 获取预测的旋转索引（pred[0].argmax() 返回预测概率最大的索引值）\n            rot_index = int(pred[0].argmax())\n\n            # 将旋转索引转换为旋转参数 rot_k（0, 1, 2, 3）\n            rot_k = convert_rot_k(rot_index)\n\n            # 将旋转参数添加到结果列表中\n            rotations.append(rot_k)\n\n            # 打印图像文件名和预测的旋转参数\n            print(f\"{os.path.basename(img_files[i])} > rot_k={rot_k}\")\n\n    # 返回所有图像的旋转参数\n    return rotations","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-31T03:18:13.46369Z","iopub.execute_input":"2025-05-31T03:18:13.463983Z","iopub.status.idle":"2025-05-31T03:18:16.578972Z","shell.execute_reply.started":"2025-05-31T03:18:13.463954Z","shell.execute_reply":"2025-05-31T03:18:16.578279Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 用于对指定的数据集执行完整的三维重建流程。该流程包括图像配对、特征提取、特征匹配、COLMAP 数据库创建、RANSAC 估计以及增量式三维重建等多个步骤。\ndef process_datasets(\n    sim_threshold=0.45,\n    min_pairs=40,\n    exhaustive_if_less=40,\n    aliked_max_keypoints=7000\n):\n    \"\"\"对所有数据集执行完整三维重建流程\n\n    Args:\n        sim_threshold (float): 相似度阈值，越高越严格\n        min_pairs (int): 每张图片至少选择的配对数量\n        exhaustive_if_less (int): 当图片数少于该值时用穷举配对\n        aliked_max_keypoints (int): aliked特征提取最大关键点数\n    \"\"\"\n    results = []\n\n    # 1. 遍历数据集\n    for dataset_name, preds in samples.items():\n        if datasets_to_process and dataset_name not in datasets_to_process:\n            print(f'Skipping \"{dataset_name}\"')\n            continue\n\n        img_dir = os.path.join(\n            data_dir, 'train' if is_train else 'test', dataset_name)\n        img_paths = [os.path.join(img_dir, p.filename) for p in preds]\n        if max_images is not None:\n            img_paths = img_paths[:max_images]\n\n        print(f'\\nProcessing dataset \"{dataset_name}\": {len(img_paths)} images')\n\n        fname_to_idx = {p.filename: idx for idx, p in enumerate(preds)}\n        feat_dir = os.path.join(workdir, 'featureout', dataset_name)\n        os.makedirs(feat_dir, exist_ok=True)\n\n        try:\n\n            # 2. 图像配对\n            t0 = time()\n            img_pairs = get_image_pairs_shortlist(\n                img_paths,\n                sim_th=sim_threshold,\n                min_pairs=min_pairs,\n                exhaustive_if_less=exhaustive_if_less,\n                device=device\n            )\n            rotations = exec_rotation_detection(img_paths, device)\n            timings['shortlisting'].append(time() - t0)\n            print(f'Shortlisting: {len(img_pairs)} pairs in {time() - t0:.4f} sec')\n            gc.collect()\n\n            # 3. 特征提取\n            t1 = time()\n            detect_aliked(img_paths, rotations, feat_dir, aliked_max_keypoints, device=device)\n            gc.collect()\n            timings['feature_detection'].append(time() - t1)\n            print(f'Features detected in {time() - t1:.4f} sec')\n\n            # 4. 特征匹配\n            t2 = time()\n            match_with_lightglue(\n                img_paths, img_pairs, feature_dir=feat_dir, device=device, verbose=False)\n            timings['feature_matching'].append(time() - t2)\n            print(f'Features matched in {time() - t2:.4f} sec')\n\n            # 5. COLMAP 数据库和 RANSAC\n            db_path = os.path.join(feat_dir, 'colmap.db')\n            if os.path.isfile(db_path):\n                os.remove(db_path)\n            gc.collect()\n            sleep(1)\n            import_into_colmap(img_dir, feature_dir=feat_dir, database_path=db_path)\n            out_path = f'{feat_dir}/colmap_rec_aliked'\n\n            t3 = time()\n            pycolmap.match_exhaustive(db_path)\n            timings['RANSAC'].append(time() - t3)\n            print(f'Ran RANSAC in {time() - t3:.4f} sec')\n\n\n            # 6. 增量三维重建\n            mapper_opts = pycolmap.IncrementalPipelineOptions()\n            mapper_opts.min_model_size = 8\n            mapper_opts.max_num_models = 35\n            os.makedirs(out_path, exist_ok=True)\n            t4 = time()\n            maps = pycolmap.incremental_mapping(\n                database_path=db_path,\n                image_path=img_dir,\n                output_path=out_path,\n                options=mapper_opts\n            )\n            sleep(1)\n            timings['Reconstruction'].append(time() - t4)\n            print(f'Reconstruction done in {time() - t4:.4f} sec')\n            print(maps)\n\n            clear_output(wait=False)\n\n            # 7. 结果存储\n            reg_count = 0\n            for map_idx, cur_map in maps.items():\n                for _, img in cur_map.images.items():\n                    pred_idx = fname_to_idx[img.name]\n                    preds[pred_idx].cluster_index = map_idx\n                    preds[pred_idx].rotation = deepcopy(img.cam_from_world.rotation.matrix())\n                    preds[pred_idx].translation = deepcopy(img.cam_from_world.translation)\n                    reg_count += 1\n            result_str = f'Dataset \"{dataset_name}\" -> Registered {reg_count} / {len(img_paths)} images in {len(maps)} clusters'\n            results.append(result_str)\n            print(result_str)\n            gc.collect()\n        except Exception as e:\n            print(e)\n            result_str = f'Dataset \"{dataset_name}\" -> Failed!'\n            results.append(result_str)\n            print(result_str)\n\n    print('\\nResults')\n    for s in results:\n        print(s)\n\n    print('\\nTimings')\n    for key, val in timings.items():\n        print(f'{key} -> total={sum(val):.02f} sec.')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-31T03:18:16.579709Z","iopub.execute_input":"2025-05-31T03:18:16.579979Z","iopub.status.idle":"2025-05-31T03:18:16.591885Z","shell.execute_reply.started":"2025-05-31T03:18:16.57995Z","shell.execute_reply":"2025-05-31T03:18:16.590912Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"process_datasets(aliked_max_keypoints=7000)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-31T03:18:16.593678Z","iopub.execute_input":"2025-05-31T03:18:16.593907Z","iopub.status.idle":"2025-05-31T03:20:54.63215Z","shell.execute_reply.started":"2025-05-31T03:18:16.593889Z","shell.execute_reply":"2025-05-31T03:20:54.631442Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 这段代码定义了一个流程，用于生成提交文件 submission.csv，并将三维重建结果写入文件中。代码的核心是处理预测结果 samples 中的每个数据集与图像输出，将其格式化为所需的 CSV 格式。\narray_to_str = lambda array: ';'.join([f\"{x:.09f}\" for x in array])\nnone_to_str = lambda n: ';'.join(['nan'] * n)\n\nsubmission_file = '/kaggle/working/submission.csv'\nwith open(submission_file, 'w') as f:\n    if is_train:\n        f.write('dataset,scene,image,rotation_matrix,translation_vector\\n')\n        for dataset in samples:\n            for prediction in samples[dataset]:\n                cluster_name = 'outliers' if prediction.cluster_index is None else f'cluster{prediction.cluster_index}'\n                rotation = none_to_str(9) if prediction.rotation is None else array_to_str(prediction.rotation.flatten())\n                translation = none_to_str(3) if prediction.translation is None else array_to_str(prediction.translation)\n                f.write(f'{prediction.dataset},{cluster_name},{prediction.filename},{rotation},{translation}\\n')\n    else:\n        f.write('image_id,dataset,scene,image,rotation_matrix,translation_vector\\n')\n        for dataset in samples:\n            for prediction in samples[dataset]:\n                cluster_name = 'outliers' if prediction.cluster_index is None else f'cluster{prediction.cluster_index}'\n                rotation = none_to_str(9) if prediction.rotation is None else array_to_str(prediction.rotation.flatten())\n                translation = none_to_str(3) if prediction.translation is None else array_to_str(prediction.translation)\n                f.write(f'{prediction.image_id},{prediction.dataset},{cluster_name},{prediction.filename},{rotation},{translation}\\n')\n\n!head {submission_file}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-31T03:20:54.633354Z","iopub.execute_input":"2025-05-31T03:20:54.633704Z","iopub.status.idle":"2025-05-31T03:20:54.803036Z","shell.execute_reply.started":"2025-05-31T03:20:54.633679Z","shell.execute_reply":"2025-05-31T03:20:54.802222Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}