{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":91498,"databundleVersionId":11655853,"sourceType":"competition"}],"dockerImageVersionId":31154,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-03T04:57:57.632747Z","iopub.execute_input":"2025-11-03T04:57:57.633519Z","iopub.status.idle":"2025-11-03T04:58:03.691778Z","shell.execute_reply.started":"2025-11-03T04:57:57.633483Z","shell.execute_reply":"2025-11-03T04:58:03.691071Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport pandas as pd\nimport numpy as np\nfrom PIL import Image\nfrom torchvision import transforms\n# DBSCAN is no longer the primary method, but we keep it for context if needed\nfrom sklearn.cluster import DBSCAN \nfrom sklearn.metrics import adjusted_rand_score\nfrom tqdm.notebook import tqdm\nfrom pathlib import Path\nimport timm\nimport subprocess\nimport os\nimport shutil\n\n# --- NEW IMPORTS for the Coarse-to-Fine Strategy ---\nimport cv2\nimport networkx as nx\nfrom sklearn.neighbors import NearestNeighbors","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T04:58:03.693039Z","iopub.execute_input":"2025-11-03T04:58:03.693404Z","iopub.status.idle":"2025-11-03T04:58:18.493450Z","shell.execute_reply.started":"2025-11-03T04:58:03.693361Z","shell.execute_reply":"2025-11-03T04:58:18.492859Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- Configuration ---\nROOT_DIR = Path(\"/kaggle/input/image-matching-challenge-2025\")\nTRAIN_DIR = ROOT_DIR / \"train\"\nTEST_DIR = ROOT_DIR / \"test\"\nTRAIN_LABELS_PATH = ROOT_DIR / \"train_labels.csv\"\n\n\n# 将缓存和输出目录放在可写的路径下（当前工作目录）\nCACHE_DIR = Path(\"/kaggle/working/cache\")  # 修改为可写路径\nCOLMAP_OUTPUT_DIR = Path(\"/kaggle/working/colmap_output\")  # 修改为可写路径\nSUBMISSION_PATH = \"/kaggle/working/submission.csv\"\n\nCACHE_DIR.mkdir(parents=True, exist_ok=True)\nCOLMAP_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)\n\n# Model configuration\nDINO_MODEL_NAME = \"dinov2_vits14\"\n\n# --- NEW Hyperparameters for Coarse-to-Fine Strategy ---\n# 1. For Coarse Search: How many initial candidates to check for each image\nK_NEIGHBORS = 50 \n\n# 2. For Fine Matching: Minimum number of geometrically consistent matches (inliers)\n#    to consider two images as part of the same scene.\nMIN_INLIERS = 15","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T05:44:46.523221Z","iopub.execute_input":"2025-11-03T05:44:46.523568Z","iopub.status.idle":"2025-11-03T05:44:46.529480Z","shell.execute_reply.started":"2025-11-03T05:44:46.523544Z","shell.execute_reply":"2025-11-03T05:44:46.528778Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ### Part 2: Core Helper Functions ###\n# This cell defines all necessary helper functions for device handling,\n# model loading, and feature extraction.\n\ndef get_device():\n    if torch.cuda.is_available(): return torch.device(\"cuda\")\n    # if torch.backends.mps.is_available(): return torch.device(\"mps\")\n    return torch.device(\"cpu\")\n\nDEVICE = get_device()\nprint(f\"Using device: {DEVICE}\")\n\ndef load_models(device):\n    \"\"\"Loads the DINOv2 model. SigLIP is no longer needed for clustering.\"\"\"\n    print(\"Loading DINOv2 model...\")\n    dino_model = torch.hub.load('facebookresearch/dinov2', DINO_MODEL_NAME, verbose=False)\n    dino_model.to(device).eval()\n    \n    dino_transform = transforms.Compose([\n        transforms.Resize(224, interpolation=transforms.InterpolationMode.BICUBIC),\n        transforms.CenterCrop(224),\n        transforms.ToTensor(),\n        transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n    ])\n    \n    # We return placeholders for the siglip model to keep the function signature consistent\n    # if other parts of the code were to use it, but it's not loaded to save memory.\n    return dino_model, None, dino_transform, None\n\ndef extract_features(image_paths, model, transform, device, cache_path):\n    \"\"\"Extracts features for a list of images using a given model.\"\"\"\n    if cache_path.exists():\n        print(f\"Loading features from cache: {cache_path}\")\n        return np.load(cache_path)\n        \n    print(f\"Extracting features for {cache_path.stem}...\")\n    features = []\n    with torch.no_grad():\n        for img_path in tqdm(image_paths, desc=f\"Extracting {cache_path.stem}\"):\n            try:\n                img = Image.open(img_path).convert(\"RGB\")\n                img_tensor = transform(img).unsqueeze(0).to(device)\n                feature = model(img_tensor)\n                features.append(feature.cpu().numpy().flatten())\n            except Exception as e:\n                print(f\"Failed to process {img_path}: {e}\")\n                # DINOv2 small model has an embedding dimension of 384\n                features.append(np.zeros(384))\n                \n    features_np = np.array(features, dtype=np.float32)\n    features_np = features_np / np.linalg.norm(features_np, axis=1, keepdims=True)\n    \n    print(f\"Saving features to cache: {cache_path}\")\n    np.save(cache_path, features_np)\n    return features_np","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T04:58:18.500773Z","iopub.execute_input":"2025-11-03T04:58:18.501131Z","iopub.status.idle":"2025-11-03T04:58:18.598723Z","shell.execute_reply.started":"2025-11-03T04:58:18.501107Z","shell.execute_reply":"2025-11-03T04:58:18.597903Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ### Part 3: New Helper Function for Geometric Verification\n# This cell replaces the old tuning logic with the core function for our new strategy.\n\ndef geometric_verification(img_path1, img_path2, min_inliers=MIN_INLIERS):\n    \"\"\"\n    Performs geometric verification between two images using SIFT and RANSAC.\n    \n    Returns:\n        True if the number of inlier matches is above the threshold, False otherwise.\n    \"\"\"\n    try:\n        # 1. Read images and convert to grayscale\n        img1 = cv2.imread(str(img_path1), cv2.IMREAD_GRAYSCALE)\n        img2 = cv2.imread(str(img_path2), cv2.IMREAD_GRAYSCALE)\n        \n        if img1 is None or img2 is None: return False\n\n        # 2. Detect and compute SIFT features\n        sift = cv2.SIFT_create()\n        kp1, des1 = sift.detectAndCompute(img1, None)\n        kp2, des2 = sift.detectAndCompute(img2, None)\n\n        if des1 is None or des2 is None or len(kp1) < 2 or len(kp2) < 2:\n            return False\n\n        # 3. Match features using a Brute-Force Matcher with Lowe's Ratio Test\n        bf = cv2.BFMatcher()\n        matches = bf.knnMatch(des1, des2, k=2)\n        \n        good_matches = []\n        for m, n in matches:\n            if m.distance < 0.75 * n.distance:\n                good_matches.append(m)\n\n        # 4. Perform RANSAC to find homography and count inliers\n        if len(good_matches) > min_inliers:\n            src_pts = np.float32([kp1[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2)\n            dst_pts = np.float32([kp2[m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2)\n            \n            # findHomography is a robust way to check for geometric consistency\n            M, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)\n            \n            if mask is None: return False\n            \n            num_inliers = np.sum(mask)\n            return num_inliers >= min_inliers\n        else:\n            return False\n            \n    except Exception as e:\n        # print(f\"Error during geometric verification between {img_path1.name} and {img_path2.name}: {e}\")\n        return False","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T04:58:18.600552Z","iopub.execute_input":"2025-11-03T04:58:18.601161Z","iopub.status.idle":"2025-11-03T04:58:18.615356Z","shell.execute_reply.started":"2025-11-03T04:58:18.601129Z","shell.execute_reply":"2025-11-03T04:58:18.614570Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ### Part 4: New Clustering logic for the `test` set (Coarse-to-Fine)\n\nall_test_clusters = {}\n\n# We only need the DINO model for the coarse search phase\ndino_model, _, dino_transform, _ = load_models(DEVICE)\ndino_model.eval()\n\nfor dataset_path in sorted(TEST_DIR.iterdir()):\n    if not dataset_path.is_dir():\n        continue\n    dataset_name = dataset_path.name\n    print(f\"\\n--- Processing Test Dataset: {dataset_name} with Coarse-to-Fine Strategy ---\")\n    \n    test_image_paths = sorted(list(dataset_path.glob(\"*.png\")))\n    num_images = len(test_image_paths)\n    if num_images == 0:\n        print(f\"No images found in {dataset_name}, skipping.\")\n        continue\n\n    # --- 1. Coarse Search: Use DINO features to find candidate pairs ---\n    print(f\"Step 1/3: Coarse search using DINOv2 features for {num_images} images.\")\n    # Use only DINO features, as they are better for instance-level geometry\n    dino_features_test = extract_features(test_image_paths, dino_model, dino_transform, DEVICE, CACHE_DIR / f\"test_{dataset_name}_features_dino.npy\")\n    \n    # Use NearestNeighbors to find top K candidates for each image\n    nn_model = NearestNeighbors(n_neighbors=min(K_NEIGHBORS, num_images), metric='cosine', algorithm='brute')\n    nn_model.fit(dino_features_test)\n    distances, indices = nn_model.kneighbors(dino_features_test)\n    \n    # --- 2. Fine Matching: Build a graph using geometric verification ---\n    print(f\"Step 2/3: Fine matching on candidate pairs using geometric verification.\")\n    G = nx.Graph()\n    G.add_nodes_from(range(num_images)) # Each image is a node\n\n    # A set to keep track of pairs we've already checked\n    verified_pairs = set()\n\n    for i in tqdm(range(num_images), desc=\"Verifying pairs\"):\n        for j in indices[i, 1:]: # The first neighbor is always the image itself\n            # Ensure we check each pair only once (e.g., (1, 5) but not (5, 1))\n            if (min(i, j), max(i, j)) in verified_pairs:\n                continue\n            \n            if geometric_verification(test_image_paths[i], test_image_paths[j]):\n                G.add_edge(i, j)\n            \n            verified_pairs.add((min(i, j), max(i, j)))\n\n    # --- 3. Final Clustering: Find connected components in the graph ---\n    print(f\"Step 3/3: Finding clusters from the verified graph.\")\n    components = list(nx.connected_components(G))\n    \n    # Format the results into the required dictionary structure\n    dataset_clusters = {}\n    used_indices = set()\n    for cluster_id, component in enumerate(components):\n        # Only consider components with at least 2 images as a valid cluster\n        if len(component) >= 2:\n            scene_name = f\"cluster_{cluster_id}\"\n            dataset_clusters[scene_name] = [test_image_paths[i] for i in component]\n            used_indices.update(component)\n\n    # Images not in any valid cluster are outliers\n    dataset_clusters[\"outliers\"] = [test_image_paths[i] for i in range(num_images) if i not in used_indices]\n    \n    all_test_clusters[dataset_name] = dataset_clusters\n    \n    num_clusters = len(components) - (1 if 'outliers' in dataset_clusters and not dataset_clusters['outliers'] else 0)\n    num_outliers = len(dataset_clusters.get(\"outliers\", []))\n    print(f\"Clustering for {dataset_name} complete: Found {num_clusters} clusters and {num_outliers} outliers.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T04:58:18.616252Z","iopub.execute_input":"2025-11-03T04:58:18.616612Z","iopub.status.idle":"2025-11-03T05:17:20.450620Z","shell.execute_reply.started":"2025-11-03T04:58:18.616588Z","shell.execute_reply":"2025-11-03T05:17:20.449682Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# !pip install pycolmap\n!apt install colmap","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T05:32:21.238718Z","iopub.execute_input":"2025-11-03T05:32:21.239304Z","iopub.status.idle":"2025-11-03T05:32:40.408650Z","shell.execute_reply.started":"2025-11-03T05:32:21.239285Z","shell.execute_reply":"2025-11-03T05:32:40.407629Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import subprocess\nimport shutil\nfrom pathlib import Path\nimport numpy as np\n\ndef qvec_to_rotmat(qvec):\n    \"\"\"手动实现四元数到旋转矩阵的转换\"\"\"\n    w, x, y, z = qvec\n    return np.array([\n        [1 - 2*y*y - 2*z*z, 2*x*y - 2*z*w,     2*x*z + 2*y*w],\n        [2*x*y + 2*z*w,     1 - 2*x*x - 2*z*z, 2*y*z - 2*x*w],\n        [2*x*z - 2*y*w,     2*y*z + 2*x*w,     1 - 2*x*x - 2*y*y]\n    ])\n\ndef parse_colmap_images_file(images_file_path):\n    \"\"\"解析COLMAP的images.txt文件来获取姿态\"\"\"\n    poses = {}\n    try:\n        with open(images_file_path, 'r') as f:\n            lines = f.readlines()\n        \n        i = 0\n        while i < len(lines):\n            line = lines[i].strip()\n            if line and not line.startswith('#'):\n                parts = line.split()\n                if len(parts) >= 10:\n                    image_name = parts[-1]\n                    qw, qx, qy, qz = map(float, parts[1:5])\n                    tx, ty, tz = map(float, parts[5:8])\n                    \n                    qvec = [qw, qx, qy, qz]\n                    tvec = [tx, ty, tz]\n                    \n                    rotation_matrix = qvec_to_rotmat(qvec)\n                    \n                    rot_str = \";\".join([f\"{x:.8f}\" for x in rotation_matrix.flatten()])\n                    trans_str = \";\".join([f\"{x:.8f}\" for x in tvec])\n                    \n                    poses[image_name] = (rot_str, trans_str)\n                \n                i += 4\n            else:\n                i += 1\n                \n    except Exception as e:\n        print(f\"解析images.txt文件出错: {e}\")\n    \n    return poses\n\ndef convert_colmap_model_to_txt(model_path, output_path):\n    \"\"\"使用model_converter将二进制模型转换为文本格式\"\"\"\n    cmd_converter = [\n        \"colmap\", \"model_converter\",\n        \"--input_path\", str(model_path),\n        \"--output_path\", str(output_path),\n        \"--output_type\", \"TXT\"\n    ]\n    try:\n        result = subprocess.run(cmd_converter, check=True, capture_output=True, text=True)\n        print(\"Model conversion successful\")\n        return True\n    except subprocess.CalledProcessError as e:\n        print(f\"Model conversion failed: {e.stderr}\")\n        return False\n\ndef estimate_poses_with_colmap(image_paths_in_cluster, work_dir):\n    \"\"\"使用命令行COLMAP进行姿态估计\"\"\"\n    print(f\"Starting pose estimation for cluster: {work_dir.name}\")\n    \n    # 设置工作目录\n    image_dir = work_dir / \"images\"\n    database_path = work_dir / \"database.db\"\n    \n    if work_dir.exists():\n        shutil.rmtree(work_dir)\n    image_dir.mkdir(parents=True)\n    \n    # 拷贝图片\n    for img_path in image_paths_in_cluster:\n        shutil.copy2(img_path, image_dir / img_path.name)\n\n    # 1. 特征提取\n    print(\"Step 1/4: Feature Extraction\")\n    cmd_feature_extractor = [\n        \"colmap\", \"feature_extractor\",\n        \"--database_path\", str(database_path),\n        \"--image_path\", str(image_dir),\n        \"--ImageReader.camera_model\", \"SIMPLE_RADIAL\",\n        \"--SiftExtraction.use_gpu\", \"false\",\n    ]\n    try:\n        subprocess.run(cmd_feature_extractor, check=True, capture_output=True, text=True)\n        print(\"Feature extraction successful\")\n    except subprocess.CalledProcessError as e:\n        print(f\"Feature Extraction FAILED: {e.stderr}\")\n        return None\n    except FileNotFoundError:\n        print(\"COLMAP not found. Please ensure COLMAP is installed.\")\n        return None\n\n    # 2. 特征匹配\n    print(\"Step 2/4: Feature Matching\")\n    cmd_exhaustive_matcher = [\n        \"colmap\", \"exhaustive_matcher\",\n        \"--database_path\", str(database_path),\n        \"--SiftMatching.use_gpu\", \"false\",\n    ]\n    try:\n        subprocess.run(cmd_exhaustive_matcher, check=True, capture_output=True, text=True)\n        print(\"Feature matching successful\")\n    except subprocess.CalledProcessError as e:\n        print(f\"Feature Matching FAILED: {e.stderr}\")\n        return None\n\n    # 3. 场景重建\n    print(\"Step 3/4: Scene Reconstruction\")\n    sparse_model_dir = work_dir / \"sparse\"\n    sparse_model_dir.mkdir(parents=True, exist_ok=True)\n    \n    cmd_mapper = [\n        \"colmap\", \"mapper\",\n        \"--database_path\", str(database_path),\n        \"--image_path\", str(image_dir),\n        \"--output_path\", str(sparse_model_dir),\n    ]\n    try:\n        subprocess.run(cmd_mapper, check=True, capture_output=True, text=True)\n        print(\"Scene reconstruction successful\")\n    except subprocess.CalledProcessError as e:\n        print(f\"Mapper FAILED: {e.stderr}\")\n        return None\n\n    # 4. 查找重建目录并转换为文本格式\n    print(\"Step 4/4: Converting and parsing reconstruction\")\n    \n    # 查找重建目录\n    reconstruction_dirs = []\n    for item in sparse_model_dir.iterdir():\n        if item.is_dir() and item.name.isdigit():\n            reconstruction_dirs.append(item)\n    \n    if not reconstruction_dirs:\n        print(\"No reconstruction directories found\")\n        return None\n    \n    reconstruction_path = reconstruction_dirs[0]\n    print(f\"Found reconstruction in {reconstruction_path}\")\n    \n    # 转换为文本格式\n    txt_output_dir = work_dir / \"txt_output\"\n    txt_output_dir.mkdir(parents=True, exist_ok=True)\n    \n    if not convert_colmap_model_to_txt(reconstruction_path, txt_output_dir):\n        return None\n    \n    # 解析images.txt文件\n    images_file_path = txt_output_dir / \"images.txt\"\n    if not images_file_path.exists():\n        print(\"images.txt not found after conversion\")\n        return None\n    \n    poses = parse_colmap_images_file(images_file_path)\n    \n    if poses:\n        print(f\"Pose estimation successful. Registered {len(poses)} images.\")\n    else:\n        print(\"No poses extracted from images.txt\")\n    \n    return poses\n\n# --- 主循环 ---\nall_poses = {}\nnan_rotation = \";\".join([\"nan\"] * 9)\nnan_translation = \";\".join([\"nan\"] * 3)\n\nfor dataset_name, clusters in all_test_clusters.items():\n    dataset_poses = {}\n    print(f\"\\n--- Estimating Poses for Dataset: {dataset_name} ---\")\n    \n    for scene_name, image_paths in tqdm(clusters.items(), desc=f\"COLMAP for {dataset_name}\"):\n        if scene_name == \"outliers\" or len(image_paths) < 2:\n            for img_path in image_paths: \n                dataset_poses[img_path.name] = (nan_rotation, nan_translation)\n            continue\n            \n        work_dir = COLMAP_OUTPUT_DIR / dataset_name / scene_name\n        \n        poses = estimate_poses_with_colmap(image_paths, work_dir)\n        \n        # 记录结果\n        for img_path in image_paths:\n            if poses and img_path.name in poses:\n                rot, trans = poses[img_path.name]\n                dataset_poses[img_path.name] = (rot, trans)\n            else:\n                dataset_poses[img_path.name] = (nan_rotation, nan_translation)\n        \n        registered_count = sum(1 for v in dataset_poses.values() if \"nan\" not in v[0])\n        print(f\"{scene_name}: {registered_count}/{len(image_paths)} images registered\")\n                \n    all_poses[dataset_name] = dataset_poses","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T05:39:15.580481Z","iopub.execute_input":"2025-11-03T05:39:15.580743Z","iopub.status.idle":"2025-11-03T05:40:52.541328Z","shell.execute_reply.started":"2025-11-03T05:39:15.580726Z","shell.execute_reply":"2025-11-03T05:40:52.540290Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission_data = []\nfor dataset_name, clusters in all_test_clusters.items():\n    for scene_name, image_paths in clusters.items():\n        for img_path in image_paths:\n            img_name = img_path.name\n            rotation, translation = all_poses[dataset_name].get(img_name, (nan_rotation, nan_translation))\n            submission_data.append({\n                \"dataset\": dataset_name,\n                \"scene\": scene_name,\n                \"image\": img_name,\n                \"rotation_matrix\": rotation,\n                \"translation_vector\": translation,\n            })\n            \nsubmission_df = pd.DataFrame(submission_data)\nsubmission_df.to_csv(SUBMISSION_PATH, index=False)\n\nprint(f\"\\nSubmission file created at: {SUBMISSION_PATH}\")\nprint(\"--- Submission File Head ---\")\nprint(submission_df)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T05:44:52.591867Z","iopub.execute_input":"2025-11-03T05:44:52.592497Z","iopub.status.idle":"2025-11-03T05:44:52.604048Z","shell.execute_reply.started":"2025-11-03T05:44:52.592475Z","shell.execute_reply":"2025-11-03T05:44:52.603329Z"}},"outputs":[],"execution_count":null}]}