{"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":11924468,"sourceType":"datasetVersion","datasetId":6988459},{"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":false,"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},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"","metadata":{}},{"cell_type":"code","source":"# =========================\n# 1. SETUP & CONFIGURATION\n# =========================\nimport os\nimport sys          # <-- missing import, now added\nimport gc\nimport math\nimport time\nfrom pathlib import Path\nfrom collections import defaultdict\n\nimport numpy as np\nimport pandas as pd\nfrom tqdm.auto import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nimport kornia as K\nimport kornia.augmentation as Kaug\n\nfrom PIL import Image\nimport pycolmap\n\n# IMC25 utilities (COLMAPDatabase, add_keypoints, add_matches)\nsys.path.append('/kaggle/input/imc25-utils')\nfrom database import *      # gives COLMAPDatabase\nfrom h5_to_db import *      # gives add_keypoints, add_matches\n\n# HuggingFace DINOv2\nfrom transformers import AutoImageProcessor, AutoModel\n\n# --------- basic config ----------\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(\"Using device:\", DEVICE)\n\nSEED = 42\ntorch.manual_seed(SEED)\nnp.random.seed(SEED)\n\n# Paths – adapt to your Kaggle dataset structure\nCOMPETITION_DATA_DIR = Path(\"/kaggle/input/image-matching-challenge-2025\")\nTRAIN_LABELS_PATH = COMPETITION_DATA_DIR / \"train_labels.csv\"\nTRAIN_IMAGES_DIR = COMPETITION_DATA_DIR / \"train\"\nTEST_IMAGES_DIR = COMPETITION_DATA_DIR / \"test\"\n\n# Where we will store intermediate features / matches / colmap DBs\nWORK_DIR = Path(\"/kaggle/working\")\nFEATURES_DIR = WORK_DIR / \"features\"\nFEATURES_DIR.mkdir(parents=True, exist_ok=True)\n\n# DINOv2 base model location (adjust if your path differs)\nDINOV2_MODEL_DIR = Path(\"/kaggle/input/dinov2/pytorch/base/1\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================\n# 2. IMAGE I/O UTILITIES\n# =========================\n\ndef load_torch_image(path: Path, device: torch.device = DEVICE) -> torch.Tensor:\n    \"\"\"\n    Load an image from disk and convert to a 4D float tensor: (1, 3, H, W).\n    Keeps original aspect ratio and dynamic range in [0, 1].\n    \"\"\"\n    img = Image.open(path).convert(\"RGB\")\n    img = np.asarray(img).astype(np.float32) / 255.0  # (H, W, 3)\n    img = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0)  # (1, 3, H, W)\n    return img.to(device)\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================\n# 3. ENHANCED DINOV2 GLOBAL DESCRIPTORS\n# =========================\n\n_DINO_PROCESSOR = None\n_DINO_MODEL = None\n\ndef get_dinov2_model():\n    \"\"\"\n    Lazy-load DINOv2 processor and model once, reuse for all images.\n    \"\"\"\n    global _DINO_PROCESSOR, _DINO_MODEL\n    if _DINO_MODEL is None:\n        print(\"Loading DINOv2 base from:\", DINOV2_MODEL_DIR)\n        _DINO_PROCESSOR = AutoImageProcessor.from_pretrained(DINOV2_MODEL_DIR)\n        _DINO_MODEL = AutoModel.from_pretrained(DINOV2_MODEL_DIR)\n        _DINO_MODEL.eval().to(DEVICE)\n    return _DINO_PROCESSOR, _DINO_MODEL\n\ndef dino_descriptor_single_scale(img_4d: torch.Tensor,\n                                 p_gem: float = 3.0) -> torch.Tensor:\n    \"\"\"\n    img_4d: (1, 3, H, W) float tensor on DEVICE\n    Returns: (1, 2C) descriptor = concat(MAC, GeM), L2-normalized.\n    \"\"\"\n    processor, model = get_dinov2_model()\n\n    with torch.inference_mode():\n        # Use processor to apply DINOv2's standard transforms\n        # Note: do_rescale=False because we already scaled to [0, 1].\n        inputs = processor(images=img_4d, return_tensors=\"pt\", do_rescale=False)\n        # Processor may move tensors to CPU; push to DEVICE\n        inputs = {k: v.to(DEVICE) for k, v in inputs.items()}\n\n        outputs = model(**inputs)  # last_hidden_state: [B, N, C]\n        tokens = outputs.last_hidden_state[:, 1:, :]  # drop CLS token\n\n        # MAC pooling\n        mac = tokens.max(dim=1)[0]  # [B, C]\n\n        # GeM pooling\n        gem = (tokens.clamp(min=0).pow(p_gem).mean(dim=1)).pow(1.0 / p_gem)  # [B, C]\n\n        desc = torch.cat([mac, gem], dim=-1)  # [B, 2C]\n        desc = F.normalize(desc, p=2, dim=-1)\n    return desc  # [1, 2C]\n\ndef dino_descriptor_multiscale(img_path: Path,\n                               num_scales: int = 2,\n                               base_resize: int | None = None) -> torch.Tensor:\n    \"\"\"\n    Compute multi-scale DINOv2 descriptor for a single image path.\n    - num_scales: 1 or 2 (we use 2: original + downscaled by 1/sqrt(2))\n    - base_resize: if not None, first resize the *long side* to this value.\n    Returns: (2C,) tensor on CPU.\n    \"\"\"\n    img = load_torch_image(img_path, device=DEVICE)  # (1, 3, H, W)\n    _, _, H, W = img.shape\n\n    # Optionally normalize long side to base_resize to control runtime\n    if base_resize is not None:\n        long_side = max(H, W)\n        if long_side != base_resize:\n            scale = base_resize / float(long_side)\n            new_H, new_W = int(round(H * scale)), int(round(W * scale))\n            img = K.geometry.resize(img, (new_H, new_W))\n\n    scales = [1.0]\n    if num_scales >= 2:\n        scales.append(1.0 / math.sqrt(2.0))\n\n    descs_per_scale = []\n    for s in scales:\n        if abs(s - 1.0) > 1e-3:\n            _, _, Hs, Ws = img.shape\n            new_H, new_W = int(round(Hs * s)), int(round(Ws * s))\n            img_scaled = K.geometry.resize(img, (new_H, new_W))\n        else:\n            img_scaled = img\n\n        desc_s = dino_descriptor_single_scale(img_scaled)\n        descs_per_scale.append(desc_s)\n\n    desc = torch.stack(descs_per_scale, dim=0).mean(dim=0)  # [1, 2C]\n    desc = F.normalize(desc, p=2, dim=-1)\n    return desc.squeeze(0).detach().cpu()  # (2C,)\n\ndef compute_global_descriptors(image_paths: list[Path],\n                               num_scales: int = 2,\n                               base_resize: int | None = 1024) -> torch.Tensor:\n    \"\"\"\n    Compute multi-scale DINOv2 descriptors for all given image paths.\n    Returns: (N, D) tensor on CPU.\n    \"\"\"\n    global_descs = []\n    for path in tqdm(image_paths, desc=\"DINOv2 global descriptors\"):\n        desc = dino_descriptor_multiscale(path, num_scales=num_scales, base_resize=base_resize)\n        global_descs.append(desc)\n    global_descs = torch.stack(global_descs, dim=0)  # [N, D]\n    return global_descs\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================\n# 4. PAIRWISE SIMILARITY HEAD (for domain adaptation)\n# =========================\n\nclass DinoPairHead(nn.Module):\n    \"\"\"\n    Small MLP that predicts whether two DINO descriptors belong to the same scene.\n    Input: concat(d1, d2) of size 2D.\n    Output: logits (unbounded), we can apply sigmoid for probability.\n    \"\"\"\n    def __init__(self, in_dim: int, hidden_dim: int = 512):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(2 * in_dim, hidden_dim),\n            nn.ReLU(inplace=True),\n            nn.Linear(hidden_dim, hidden_dim // 2),\n            nn.ReLU(inplace=True),\n            nn.Linear(hidden_dim // 2, 1)\n        )\n\n    def forward(self, d1: torch.Tensor, d2: torch.Tensor) -> torch.Tensor:\n        # d1, d2: [B, D]\n        x = torch.cat([d1, d2], dim=-1)\n        logits = self.net(x).squeeze(-1)  # [B]\n        return logits\n\ndef shortlist_pairs_distance(global_descs: torch.Tensor,\n                             max_pairs_per_image: int = 30,\n                             max_global_pairs: int | None = None) -> list[tuple[int, int]]:\n    \"\"\"\n    global_descs: (N, D) float tensor on CPU or GPU.\n    Returns a list of (i, j) indices with i < j.\n    Strategy:\n      - Compute pairwise distances using torch.cdist (on DEVICE).\n      - For each image, keep top-k nearest neighbors (excluding itself).\n      - Optionally cap total pair count.\n    \"\"\"\n    N, D = global_descs.shape\n    descs = global_descs.to(DEVICE)\n\n    with torch.no_grad():\n        dists = torch.cdist(descs, descs, p=2.0)  # [N, N]\n        # Large value on diagonal so it is never selected\n        dists.fill_diagonal_(1e9)\n\n        # For each i, find its k smallest distances\n        k = min(max_pairs_per_image, N - 1)\n        knn_dists, knn_idx = torch.topk(dists, k=k, dim=1, largest=False)\n\n    pairs_set = set()\n    for i in range(N):\n        for j in knn_idx[i].tolist():\n            if i == j:\n                continue\n            a, b = (i, j) if i < j else (j, i)\n            pairs_set.add((a, b))\n\n    pairs = list(pairs_set)\n    # Optionally cap total pairs to avoid exploding runtime\n    if (max_global_pairs is not None) and (len(pairs) > max_global_pairs):\n        pairs = pairs[:max_global_pairs]\n\n    print(f\"Shortlisted {len(pairs)} pairs (distance-based).\")\n    return pairs\n\ndef shortlist_pairs_with_head(global_descs: torch.Tensor,\n                              pair_head: DinoPairHead,\n                              max_pairs_per_image: int = 30,\n                              max_global_pairs: int | None = None,\n                              threshold: float | None = None) -> list[tuple[int, int]]:\n    \"\"\"\n    Use a learned pair head to score pairs and shortlist neighbors.\n    - If threshold is given: keep all pairs with sigmoid(score) >= threshold.\n    - Else: for each image keep top-k neighbors by score.\n\n    Note: This is O(N^2). For large N you may need batching or approximate neighbors.\n    \"\"\"\n    N, D = global_descs.shape\n    descs = global_descs.to(DEVICE)\n    pair_head = pair_head.to(DEVICE).eval()\n\n    with torch.no_grad():\n        # Compute scores row by row to avoid memory explosion\n        scores = torch.empty((N, N), dtype=torch.float32, device=DEVICE)\n        for i in tqdm(range(N), desc=\"Scoring pairs with pair head\"):\n            d1 = descs[i].unsqueeze(0).repeat(N, 1)  # [N, D]\n            d2 = descs  # [N, D]\n            logits = pair_head(d1, d2)  # [N]\n            scores[i] = logits\n\n        # Symmetrize\n        scores = (scores + scores.T) / 2.0\n        scores.fill_diagonal_(-1e9)\n\n        if threshold is not None:\n            probs = scores.sigmoid()\n            mask = probs >= threshold\n            idx_i, idx_j = torch.where(mask)\n            pairs_set = {(int(i), int(j)) if i < j else (int(j), int(i))\n                         for i, j in zip(idx_i.tolist(), idx_j.tolist()) if i != j}\n            pairs = list(pairs_set)\n        else:\n            k = min(max_pairs_per_image, N - 1)\n            top_scores, top_idx = torch.topk(scores, k=k, dim=1, largest=True)\n            pairs_set = set()\n            for i in range(N):\n                for j in top_idx[i].tolist():\n                    if i == j:\n                        continue\n                    a, b = (i, j) if i < j else (j, i)\n                    pairs_set.add((a, b))\n            pairs = list(pairs_set)\n\n    if (max_global_pairs is not None) and (len(pairs) > max_global_pairs):\n        pairs = pairs[:max_global_pairs]\n\n    print(f\"Shortlisted {len(pairs)} pairs (pair-head-based).\")\n    return pairs\n\ndef get_shortlist_pairs(global_descs: torch.Tensor,\n                        use_pair_head: bool = False,\n                        pair_head: DinoPairHead | None = None,\n                        max_pairs_per_image: int = 30,\n                        max_global_pairs: int | None = None,\n                        pair_head_threshold: float | None = None) -> list[tuple[int, int]]:\n    \"\"\"\n    Unified entry point:\n      - If use_pair_head is False or pair_head is None: distance-based shortlist.\n      - Else: learned-head shortlist.\n    \"\"\"\n    if use_pair_head and (pair_head is not None):\n        return shortlist_pairs_with_head(\n            global_descs,\n            pair_head,\n            max_pairs_per_image=max_pairs_per_image,\n            max_global_pairs=max_global_pairs,\n            threshold=pair_head_threshold,\n        )\n    else:\n        return shortlist_pairs_distance(\n            global_descs,\n            max_pairs_per_image=max_pairs_per_image,\n            max_global_pairs=max_global_pairs,\n        )\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================\n# 5. ALIKED FEATURE EXTRACTION (CORRECTED)\n# =========================\nimport h5py\nimport kornia.feature as KF\nfrom lightglue import ALIKED  # <-- use ALIKED from lightglue\n\n\ndef extract_aliked_features(\n    image_paths: list[Path],\n    image_names: list[str],\n    feature_dir: Path,\n    num_features: int = 4096,\n    resize_to: int = 1024,\n    detection_threshold: float = 0.05,\n    device: torch.device = DEVICE,\n):\n    \"\"\"\n    Extract local features using ALIKED (from the LightGlue package) for all images.\n    Saves keypoints.h5 and descriptors.h5 in feature_dir.\n\n    - image_paths: list of Path objects to images\n    - image_names: list of unique names (strings) used as keys in HDF5\n    \"\"\"\n    feature_dir.mkdir(parents=True, exist_ok=True)\n    kpt_path = feature_dir / \"keypoints.h5\"\n    desc_path = feature_dir / \"descriptors.h5\"\n\n    # Remove old files if any\n    if kpt_path.exists():\n        kpt_path.unlink()\n    if desc_path.exists():\n        desc_path.unlink()\n\n    # Initialize ALIKED from LightGlue\n    # LightGlue docs: extractor = ALIKED(max_num_keypoints=..., resize=..., detection_threshold=...)\n    extractor = ALIKED(\n        max_num_keypoints=num_features,\n        detection_threshold=detection_threshold,\n        resize=resize_to,\n    ).eval().to(device)\n\n    with h5py.File(kpt_path, \"w\") as f_kpt, h5py.File(desc_path, \"w\") as f_desc:\n        for img_path, name in tqdm(\n            list(zip(image_paths, image_names)),\n            desc=\"ALIKED features\",\n            total=len(image_paths),\n        ):\n            # load_torch_image returns (1,3,H,W); ALIKED expects (3,H,W)\n            img_4d = load_torch_image(img_path, device=device)  # (1,3,H,W)\n            img = img_4d.squeeze(0)  # (3,H,W)\n\n            with torch.inference_mode():\n                feats = extractor.extract(img)  # dict with 'keypoints', 'descriptors', ...\n\n            # Flatten any batch dimension if present\n            kpts = feats[\"keypoints\"]\n            descs = feats[\"descriptors\"]\n\n            # Ensure shapes are (N,2) and (N,D)\n            kpts = kpts.reshape(-1, 2).detach().cpu().numpy()\n            descs = descs.reshape(len(kpts), -1).detach().cpu().numpy()\n\n            # Store per-image\n            f_kpt.create_dataset(name, data=kpts, compression=\"gzip\")\n            f_desc.create_dataset(name, data=descs, compression=\"gzip\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================\n# 6. LIGHTGLUE MATCHING\n# =========================\n\ndef match_with_lightglue(\n    image_names,\n    index_pairs,\n    feature_dir,\n    device=DEVICE,\n    min_matches=25,\n    verbose=False,\n):\n    \"\"\"\n    Use Kornia's LightGlueMatcher to match ALIKED features for shortlisted image pairs.\n    Saves matches.h5 in feature_dir.\n    \"\"\"\n    kpt_path = feature_dir / \"keypoints.h5\"\n    desc_path = feature_dir / \"descriptors.h5\"\n    match_path = feature_dir / \"matches.h5\"\n\n    if match_path.exists():\n        match_path.unlink()\n\n    f_kpt = h5py.File(kpt_path, \"r\")\n    f_desc = h5py.File(desc_path, \"r\")\n\n    lg_matcher = KF.LightGlueMatcher(\n        \"aliked\",\n        {\n            \"width_confidence\": -1,\n            \"depth_confidence\": -1,\n            \"mp\": True if \"cuda\" in str(device) else False,\n        },\n    ).eval().to(device)\n\n    with h5py.File(match_path, \"w\") as f_match:\n        for i, j in tqdm(index_pairs, desc=\"LightGlue matching\"):\n            name_i = image_names[i]\n            name_j = image_names[j]\n\n            kpts_i = torch.from_numpy(f_kpt[name_i][:]).float().to(device)\n            kpts_j = torch.from_numpy(f_kpt[name_j][:]).float().to(device)\n            desc_i = torch.from_numpy(f_desc[name_i][:]).float().to(device)\n            desc_j = torch.from_numpy(f_desc[name_j][:]).float().to(device)\n\n            with torch.inference_mode():\n                laf_i = KF.laf_from_center_scale_ori(kpts_i[None])\n                laf_j = KF.laf_from_center_scale_ori(kpts_j[None])\n                dists, idxs = lg_matcher(desc_i, desc_j, laf_i, laf_j)\n\n            n_matches = idxs.shape[0]\n            if n_matches < min_matches:\n                continue\n\n            grp = f_match.require_group(name_i)\n            grp.create_dataset(\n                name_j,\n                data=idxs.detach().cpu().numpy().reshape(-1, 2),\n                compression=\"gzip\",\n            )\n\n    f_kpt.close()\n    f_desc.close()\n\n    if verbose:\n        print(f\"Saved matches to: {match_path}\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================\n# 7. COLMAP RECONSTRUCTION HELPERS (FIXED)\n# =========================\n\ndef import_into_colmap(\n    images_dir: str | Path,\n    feature_dir: str | Path,\n    database_path: str | Path,\n):\n    \"\"\"\n    Thin wrapper around the IMC25 utils (COLMAPDatabase + h5_to_db).\n    Assumes:\n      - feature_dir contains keypoints.h5, descriptors.h5, matches.h5\n      - images_dir is the folder with the actual image files\n    \"\"\"\n    images_dir = str(images_dir)\n    feature_dir = str(feature_dir)\n    database_path = str(database_path)\n\n    # Create / connect DB and add features + matches.\n    db = COLMAPDatabase.connect(database_path)\n    db.create_tables()\n\n    # single_camera=False -> one camera per image folder (baseline behaviour)\n    single_camera = False\n    fname_to_id = add_keypoints(\n        db,\n        feature_dir,\n        images_dir,\n        '',                 # image list file (not used here)\n        'simple-pinhole',   # camera model\n        single_camera,\n    )\n\n    add_matches(\n        db,\n        feature_dir,\n        fname_to_id,\n    )\n\n    db.commit()\n    db.close()\n\n\ndef run_colmap_reconstruction(\n    images_dir: Path,\n    image_names: list[str],   # kept for compatibility, not used internally\n    feature_dir: Path,\n    output_dir: Path,\n    database_path: Path,\n    mapper_options: \"pycolmap.IncrementalPipelineOptions | None\" = None,\n):\n    \"\"\"\n    Build a COLMAP database from our ALIKED + LightGlue HDF5 files, then run\n    geometric verification and incremental mapping via pycolmap.\n\n    This uses the IMC25 utilities (COLMAPDatabase + add_keypoints + add_matches),\n    NOT pycolmap.Database directly, so it is consistent with the Kaggle baseline.\n    \"\"\"\n    output_dir.mkdir(parents=True, exist_ok=True)\n\n    # Fresh DB\n    if database_path.exists():\n        database_path.unlink()\n\n    # 1) Import our custom features & matches into COLMAP DB\n    import_into_colmap(images_dir, feature_dir, database_path)\n\n    # 2) Geometric verification (RANSAC) on the imported matches\n    #    (baseline uses match_exhaustive for this step)\n    pycolmap.match_exhaustive(str(database_path))\n\n    # 3) Incremental mapping\n    if mapper_options is None:\n        mapper_options = pycolmap.IncrementalPipelineOptions()\n    maps = pycolmap.incremental_mapping(\n        database_path=str(database_path),\n        image_path=str(images_dir),\n        output_path=str(output_dir),\n        options=mapper_options,\n    )\n\n    return maps\n\n\n# =========================\n# 7b. ADAPTIVE RECONSTRUCTION (uses fixed run_colmap_reconstruction)\n# =========================\n\ndef reconstruct_with_adaptive_fallback(\n    dataset_name: str,\n    images_dir: Path,\n    image_paths: list[Path],\n    image_names: list[str],\n    feature_dir: Path,\n    # base shortlist:\n    global_descs_base: torch.Tensor,\n    index_pairs_base: list[tuple[int, int]],\n    # ALIKED + LG base settings:\n    base_num_features: int = 4096,\n    base_resize: int = 1024,\n    base_min_matches: int = 25,\n    # fallback settings:\n    fallback_num_features: int = 8000,\n    fallback_resize: int = 1536,\n    fallback_min_matches: int = 15,\n    registration_ratio_threshold: float = 0.6,\n) -> dict:\n    \"\"\"\n    For a given dataset:\n      - Base run uses precomputed DINO descriptors + index_pairs for ALIKED+LG.\n      - If registration ratio is too low, reruns with more aggressive\n        ALIKED / LightGlue / DINO settings at higher resolution.\n    Returns: maps dict from pycolmap.\n    \"\"\"\n    print(f\"\\n=== Reconstruction for dataset: {dataset_name} ===\")\n    dataset_feature_dir = feature_dir\n    dataset_feature_dir.mkdir(parents=True, exist_ok=True)\n\n    # ---------- BASE RUN ----------\n    print(\"Base ALIKED/LightGlue settings:\",\n          f\"num_features={base_num_features}, resize={base_resize}, min_matches={base_min_matches}\")\n\n    # 1) ALIKED features (base settings)\n    extract_aliked_features(\n        image_paths=image_paths,\n        image_names=image_names,\n        feature_dir=dataset_feature_dir,\n        num_features=base_num_features,\n        resize_to=base_resize,\n        detection_threshold=0.05,\n        device=DEVICE,\n    )\n\n    # 2) LightGlue matching using precomputed shortlist\n    match_with_lightglue(\n        image_names=image_names,\n        index_pairs=index_pairs_base,\n        feature_dir=dataset_feature_dir,\n        device=DEVICE,\n        min_matches=base_min_matches,\n        verbose=True,\n    )\n\n    # 3) COLMAP reconstruction (base)\n    base_db_path = dataset_feature_dir / \"colmap_base.db\"\n    base_out_dir = dataset_feature_dir / \"colmap_rec_base\"\n\n    base_maps = run_colmap_reconstruction(\n        images_dir=images_dir,\n        image_names=image_names,\n        feature_dir=dataset_feature_dir,\n        output_dir=base_out_dir,\n        database_path=base_db_path,\n        mapper_options=None,\n    )\n\n    def count_registered(maps: dict) -> int:\n        reg = 0\n        for _, m in maps.items():\n            reg += len(m.images)\n        return reg\n\n    base_registered = count_registered(base_maps)\n    total_images = len(image_names)\n    base_ratio = base_registered / max(1, total_images)\n    print(f\"[BASE] Registered {base_registered}/{total_images} images ({base_ratio:.2%})\")\n\n    # If base ratio is adequate or nothing was reconstructed, use base\n    if base_ratio >= registration_ratio_threshold or base_registered == 0:\n        print(\"[ADAPTIVE] Using base reconstruction.\")\n        return base_maps\n\n    # ---------- FALLBACK RUN ----------\n    print(\"[ADAPTIVE] Registration ratio low, running fallback...\")\n    print(\"Fallback ALIKED/LightGlue settings:\",\n          f\"num_features={fallback_num_features}, resize={fallback_resize}, min_matches={fallback_min_matches}\")\n\n    # 1) ALIKED with more features / larger resize\n    extract_aliked_features(\n        image_paths=image_paths,\n        image_names=image_names,\n        feature_dir=dataset_feature_dir,  # reuse dir, overwrite features\n        num_features=fallback_num_features,\n        resize_to=fallback_resize,\n        detection_threshold=0.03,  # more sensitive\n        device=DEVICE,\n    )\n\n    # 2) Recompute DINO descriptors + distance-based shortlist for fallback\n    fallback_global_descs = compute_global_descriptors(\n        image_paths, num_scales=2, base_resize=fallback_resize\n    )\n    fallback_index_pairs = shortlist_pairs_distance(\n        fallback_global_descs, max_pairs_per_image=40\n    )\n\n    match_with_lightglue(\n        image_names=image_names,\n        index_pairs=fallback_index_pairs,\n        feature_dir=dataset_feature_dir,\n        device=DEVICE,\n        min_matches=fallback_min_matches,\n        verbose=True,\n    )\n\n    fallback_db_path = dataset_feature_dir / \"colmap_fallback.db\"\n    fallback_out_dir = dataset_feature_dir / \"colmap_rec_fallback\"\n\n    mapper_options = pycolmap.IncrementalPipelineOptions()\n    mapper_options.min_model_size = 3  # allow smaller initial models\n\n    fallback_maps = run_colmap_reconstruction(\n        images_dir=images_dir,\n        image_names=image_names,\n        feature_dir=dataset_feature_dir,\n        output_dir=fallback_out_dir,\n        database_path=fallback_db_path,\n        mapper_options=mapper_options,\n    )\n\n    fallback_registered = count_registered(fallback_maps)\n    fallback_ratio = fallback_registered / max(1, total_images)\n    print(f\"[FALLBACK] Registered {fallback_registered}/{total_images} images ({fallback_ratio:.2%})\")\n\n    # Select better run\n    if fallback_registered > base_registered:\n        print(\"[ADAPTIVE] Using fallback reconstruction results.\")\n        return fallback_maps\n    else:\n        print(\"[ADAPTIVE] Fallback did not improve registration; using base.\")\n        return base_maps\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================\n# 8. DOMAIN ADAPTATION: PAIR-HEAD TRAINING\n# =========================\n# This section only defines the utilities to TRAIN the pair head.\n# It does NOT use `global_descs` or `get_shortlist_pairs` anywhere.\n\ndef load_train_labels() -> pd.DataFrame:\n    \"\"\"\n    Load train_labels.csv for the competition.\n    Expected columns (at least): dataset, scene, image, rotation_matrix, translation_vector, ...\n    \"\"\"\n    df = pd.read_csv(TRAIN_LABELS_PATH)\n    return df\n\n\ndef build_train_image_paths(df: pd.DataFrame) -> dict:\n    \"\"\"\n    Returns a dict mapping each dataset name to a list of image Paths.\n\n    { dataset_name: [Path(...), Path(...), ...], ... }\n    \"\"\"\n    datasets = df[\"dataset\"].unique().tolist()\n    image_paths_by_dataset = {}\n    for d in datasets:\n        df_d = df[df[\"dataset\"] == d]\n        paths = [\n            TRAIN_IMAGES_DIR / d / img_name\n            for img_name in df_d[\"image\"].tolist()\n        ]\n        image_paths_by_dataset[d] = paths\n    return image_paths_by_dataset\n\n\ndef compute_train_descriptors(df: pd.DataFrame) -> dict:\n    \"\"\"\n    Compute DINOv2 descriptors for all train images using the enhanced global descriptor function.\n    Returns a dict:\n        key   = \"dataset/image_name\"\n        value = descriptor tensor of shape (D,) on CPU.\n    \"\"\"\n    descs = {}\n    for dataset, df_d in df.groupby(\"dataset\"):\n        print(f\"[TRAIN DESCS] Dataset: {dataset}\")\n        image_names = df_d[\"image\"].tolist()\n        image_paths = [TRAIN_IMAGES_DIR / dataset / n for n in image_names]\n\n        # Reuse the enhanced DINOv2 function defined earlier\n        global_descs = compute_global_descriptors(\n            image_paths,\n            num_scales=2,\n            base_resize=1024,\n        )  # (N, D)\n\n        for name, desc in zip(image_names, global_descs):\n            full_key = f\"{dataset}/{name}\"\n            descs[full_key] = desc  # (D,)\n    return descs\n\n\nclass PairDataset(torch.utils.data.Dataset):\n    \"\"\"\n    Supervised pair dataset for training the pair head:\n    - Positive pairs: same dataset + same scene.\n    - Negative pairs: same dataset + different scenes.\n    \"\"\"\n    def __init__(\n        self,\n        df: pd.DataFrame,\n        desc_dict: dict,\n        num_pos_per_scene: int = 3,\n        num_neg_pairs_per_scene_pair: int = 3,\n    ):\n        self.samples = []\n        self.desc_dict = desc_dict\n\n        # -------- Positive pairs: same dataset & same scene --------\n        for (dataset, scene), df_grp in df.groupby([\"dataset\", \"scene\"]):\n            imgs = df_grp[\"image\"].tolist()\n            key_imgs = [f\"{dataset}/{img}\" for img in imgs]\n            n = len(key_imgs)\n            for i in range(n):\n                for j in range(i + 1, min(i + 1 + num_pos_per_scene, n)):\n                    self.samples.append((key_imgs[i], key_imgs[j], 1.0))\n\n        # -------- Negative pairs: same dataset, different scenes --------\n        for dataset, df_d in df.groupby(\"dataset\"):\n            scenes = df_d[\"scene\"].unique().tolist()\n            for s1_idx in range(len(scenes)):\n                for s2_idx in range(s1_idx + 1, len(scenes)):\n                    s1 = scenes[s1_idx]\n                    s2 = scenes[s2_idx]\n                    imgs1 = df_d[df_d[\"scene\"] == s1][\"image\"].tolist()\n                    imgs2 = df_d[df_d[\"scene\"] == s2][\"image\"].tolist()\n\n                    key1 = [f\"{dataset}/{img}\" for img in imgs1]\n                    key2 = [f\"{dataset}/{img}\" for img in imgs2]\n\n                    for i in range(min(num_neg_pairs_per_scene_pair, len(key1))):\n                        for j in range(min(num_neg_pairs_per_scene_pair, len(key2))):\n                            self.samples.append((key1[i], key2[j], 0.0))\n\n        print(f\"[PAIR DATASET] Total pairs: {len(self.samples)}\")\n\n    def __len__(self):\n        return len(self.samples)\n\n    def __getitem__(self, idx: int):\n        k1, k2, label = self.samples[idx]\n        d1 = self.desc_dict[k1]   # (D,)\n        d2 = self.desc_dict[k2]   # (D,)\n        return (\n            d1.clone(),            # (D,)\n            d2.clone(),            # (D,)\n            torch.tensor(label, dtype=torch.float32),\n        )\n\n\ndef train_pair_head(\n    df: pd.DataFrame,\n    desc_dict: dict,\n    hidden_dim: int = 512,\n    batch_size: int = 256,\n    num_epochs: int = 3,\n    lr: float = 1e-3,\n    save_path: Path = WORK_DIR / \"dino_pair_head.pth\",\n):\n    \"\"\"\n    Train the DinoPairHead model using precomputed DINO descriptors and train labels.\n    Produces a file dino_pair_head.pth with the trained weights.\n    \"\"\"\n    # Build the supervised pair dataset\n    dataset = PairDataset(df, desc_dict)\n    loader = torch.utils.data.DataLoader(\n        dataset,\n        batch_size=batch_size,\n        shuffle=True,\n        drop_last=True,\n    )\n\n    # Determine descriptor dimension from any entry in desc_dict\n    any_key = next(iter(desc_dict.keys()))\n    in_dim = desc_dict[any_key].shape[0]\n\n    # DinoPairHead was defined earlier in the notebook\n    model = DinoPairHead(in_dim=in_dim, hidden_dim=hidden_dim).to(DEVICE)\n    optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n    criterion = nn.BCEWithLogitsLoss()\n\n    for epoch in range(num_epochs):\n        model.train()\n        running_loss = 0.0\n\n        for d1, d2, labels in tqdm(loader, desc=f\"Epoch {epoch+1}/{num_epochs}\"):\n            d1 = d1.to(DEVICE)\n            d2 = d2.to(DEVICE)\n            labels = labels.to(DEVICE)\n\n            logits = model(d1, d2)\n            loss = criterion(logits, labels)\n\n            optimizer.zero_grad()\n            loss.backward()\n            optimizer.step()\n\n            running_loss += loss.item() * d1.size(0)\n\n        epoch_loss = running_loss / len(dataset)\n        print(f\"[PAIR HEAD] Epoch {epoch+1}, loss = {epoch_loss:.4f}\")\n\n    torch.save(model.state_dict(), save_path)\n    print(f\"[PAIR HEAD] Saved weights to: {save_path}\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================\n# 9. INFERENCE PIPELINE: LIST DATASETS & IMAGES\n# =========================\n\ndef list_test_datasets_and_images(test_images_dir: Path) -> dict[str, dict]:\n    \"\"\"\n    Returns a dict:\n      {\n        dataset_name: {\n           \"images_dir\": Path,\n           \"image_names\": [str],\n           \"image_paths\": [Path]\n        },\n        ...\n      }\n\n    Ignores non-image files (e.g., LICENSE.txt).\n    \"\"\"\n    # Allowed image extensions\n    IMG_EXTS = {\".jpg\", \".jpeg\", \".png\", \".bmp\", \".tif\", \".tiff\"}\n\n    result = {}\n    for dataset_dir in sorted(test_images_dir.iterdir()):\n        if not dataset_dir.is_dir():\n            continue\n\n        dataset_name = dataset_dir.name\n\n        # Only keep files with allowed extensions\n        image_paths = sorted(\n            [\n                p for p in dataset_dir.iterdir()\n                if p.is_file() and p.suffix.lower() in IMG_EXTS\n            ]\n        )\n\n        if len(image_paths) == 0:\n            # Safety: skip folders that do not contain any images\n            print(f\"[WARN] Dataset '{dataset_name}' has no image files. Skipping.\")\n            continue\n\n        image_names = [p.name for p in image_paths]\n        result[dataset_name] = {\n            \"images_dir\": dataset_dir,\n            \"image_names\": image_names,\n            \"image_paths\": image_paths,\n        }\n\n    return result\n\n\n# Optionally load trained pair head\nPAIR_HEAD_WEIGHTS = WORK_DIR / \"dino_pair_head.pth\"\n\nuse_pair_head = PAIR_HEAD_WEIGHTS.exists()\npair_head = None\n\nif use_pair_head:\n    # We will actually instantiate the model after computing the first global_descs,\n    # so we know the descriptor dimension.\n    print(\"[PAIR HEAD] Weights found, will load after first descriptor batch.\")\nelse:\n    print(\"[PAIR HEAD] No weights found, using distance-based shortlist only.\")\n\n\n# =========================\n# 9. INFERENCE PIPELINE: MAIN LOOP\n# =========================\n\ntest_info = list_test_datasets_and_images(TEST_IMAGES_DIR)\nall_rows = []  # rows for submission.csv\n\nfor dataset_name, info in test_info.items():\n    images_dir = info[\"images_dir\"]\n    image_names = info[\"image_names\"]\n    image_paths = info[\"image_paths\"]\n    print(f\"\\n******** Dataset: {dataset_name}, {len(image_names)} images ********\")\n\n    # 1) Compute enhanced DINOv2 descriptors for this dataset\n    global_descs = compute_global_descriptors(\n        image_paths=image_paths,\n        num_scales=2,\n        base_resize=1024,   # you can adjust this if you like\n    )  # (N, D)\n\n    # 2) If we have pair-head weights and have not yet constructed the model, do it now\n    if use_pair_head and (pair_head is None):\n        in_dim = global_descs.shape[1]\n        pair_head = DinoPairHead(in_dim=in_dim)\n        pair_head.load_state_dict(torch.load(PAIR_HEAD_WEIGHTS, map_location=DEVICE))\n        print(\"[PAIR HEAD] Loaded with in_dim =\", in_dim)\n\n    # 3) Shortlist pairs (distance-based or pair-head-based)\n    index_pairs = get_shortlist_pairs(\n        global_descs=global_descs,\n        use_pair_head=use_pair_head and (pair_head is not None),\n        pair_head=pair_head,\n        max_pairs_per_image=30,\n        max_global_pairs=None,\n        pair_head_threshold=None,  # you can experiment with a threshold if desired\n    )\n\n    # 4) Run ALIKED + LightGlue + COLMAP with adaptive fallback\n    dataset_feature_dir = FEATURES_DIR / dataset_name\n    maps = reconstruct_with_adaptive_fallback(\n        dataset_name=dataset_name,\n        images_dir=images_dir,\n        image_paths=image_paths,\n        image_names=image_names,\n        feature_dir=dataset_feature_dir,\n        global_descs_base=global_descs,\n        index_pairs_base=index_pairs,\n        base_num_features=4096,\n        base_resize=1024,\n        base_min_matches=25,\n        fallback_num_features=8000,\n        fallback_resize=1536,\n        fallback_min_matches=15,\n        registration_ratio_threshold=0.6,\n    )\n\n    # 5) Convert maps into per-image predictions and append to all_rows\n    #    We choose: cluster_index = map_index (int).\n    #    Unregistered images: cluster_index = -1, poses = nan.\n    #    Rotation/translation flattened row-major, semicolon-separated.\n\n    # Initialize prediction containers\n    pred_cluster = {name: -1 for name in image_names}\n    pred_R = {}\n    pred_t = {}\n\n    for map_index, cur_map in maps.items():\n        for img_id, img in cur_map.images.items():\n            name = img.name  # may be full path or just filename\n            if name not in pred_cluster:\n                # Sometimes pycolmap stores relative paths; get just the basename\n                name = os.path.basename(name)\n                if name not in pred_cluster:\n                    continue\n\n            pred_cluster[name] = map_index\n            R = img.cam_from_world.rotation.matrix()  # (3, 3)\n            t = img.cam_from_world.translation        # (3,)\n            pred_R[name] = R\n            pred_t[name] = t\n\n    # Prepare NaN strings for unregistered images\n    nan_rot = \";\".join([\"nan\"] * 9)\n    nan_trans = \";\".join([\"nan\"] * 3)\n\n    # Build rows for this dataset\n    for name in image_names:\n        cluster_index = pred_cluster[name]\n\n        if name in pred_R:\n            R_flat = pred_R[name].reshape(-1).tolist()\n            t_flat = pred_t[name].reshape(-1).tolist()\n            rot_str = \";\".join([f\"{x:.8f}\" for x in R_flat])\n            trans_str = \";\".join([f\"{x:.8f}\" for x in t_flat])\n        else:\n            rot_str = nan_rot\n            trans_str = nan_trans\n\n        # For the competition, \"scene\" is just a cluster label.\n        # We use the COLMAP map index (cluster_index).\n        all_rows.append(\n            {\n                \"dataset\": dataset_name,\n                \"scene\": cluster_index,\n                \"image\": name,\n                \"rotation_matrix\": rot_str,\n                \"translation_vector\": trans_str,\n            }\n        )\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================\n# 10. BUILD SUBMISSION CSV (TEST MODE ONLY)\n# =========================\n\n# all_rows was filled in the main loop:\n# {\n#   \"dataset\": dataset_name,\n#   \"scene\": cluster_index,           # int or -1\n#   \"image\": name,\n#   \"rotation_matrix\": rot_str,\n#   \"translation_vector\": trans_str,\n# }\n\n# Load sample_submission to get the exact image_id/dataset/image ordering\nsample_sub = pd.read_csv(COMPETITION_DATA_DIR / \"sample_submission.csv\")\n\n# Build a lookup from (dataset, image) -> our prediction row\npred_map = {(row[\"dataset\"], row[\"image\"]): row for row in all_rows}\n\noutput_rows = []\nfor _, row in sample_sub.iterrows():\n    image_id = row[\"image_id\"]\n    dataset = row[\"dataset\"]\n    image = row[\"image\"]\n\n    key = (dataset, image)\n    pred = pred_map.get(key, None)\n\n    if pred is None:\n        # No prediction found: mark as outlier with nan poses\n        scene_idx = -1\n        rot_str = \";\".join([\"nan\"] * 9)\n        trans_str = \";\".join([\"nan\"] * 3)\n    else:\n        scene_idx = pred[\"scene\"]\n        rot_str = pred[\"rotation_matrix\"]\n        trans_str = pred[\"translation_vector\"]\n\n    # Convert scene index -> string label\n    if scene_idx is None or scene_idx < 0:\n        scene_str = \"outliers\"\n    else:\n        scene_str = f\"cluster{scene_idx}\"\n\n    output_rows.append(\n        {\n            \"image_id\": image_id,\n            \"dataset\": dataset,\n            \"scene\": scene_str,\n            \"image\": image,\n            \"rotation_matrix\": rot_str,\n            \"translation_vector\": trans_str,\n        }\n    )\n\nsubmission_df = pd.DataFrame(output_rows)\nsubmission_df = submission_df[\n    [\"image_id\", \"dataset\", \"scene\", \"image\",\n     \"rotation_matrix\", \"translation_vector\"]\n]\n\n# Write the final file Kaggle will read\nsubmission_df.to_csv(\"submission.csv\", index=False)\nprint(\"Saved submission.csv with shape:\", submission_df.shape)\nprint(submission_df.head())\n","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}