{"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":"none","dataSources":[{"sourceId":91498,"databundleVersionId":11655853,"isSourceIdPinned":false,"sourceType":"competition"}],"dockerImageVersionId":31192,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# --- CELL 1: INSTALLATION ---\n# 1. Install Libraries (Quietly)\n!pip install \"numpy<2\" kornia pycolmap transformers h5py --upgrade --quiet\n!pip install git+https://github.com/cvg/LightGlue.git --quiet\n\n# 2. Clone Repositories\n!git clone --recursive https://github.com/cvg/Hierarchical-Localization/ --quiet\n!git clone https://github.com/magicleap/SuperGluePretrainedNetwork.git --quiet\n\nprint(\"✅ Environment Ready.\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile main.py\nimport sys\nimport os\nimport torch\nimport numpy as np\nimport pandas as pd\nfrom pathlib import Path\nfrom PIL import Image\nimport shutil\nimport warnings\nimport torch.multiprocessing as mp\nimport time\nimport torchvision.transforms as T\nimport torch.nn.functional as F\n\n# Suppress warnings\nwarnings.filterwarnings(\"ignore\")\n\n# --- SETUP PATHS ---\ncurrent_dir = Path.cwd()\nsys.path.append(str(current_dir / 'Hierarchical-Localization'))\nsys.path.append(str(current_dir / 'SuperGluePretrainedNetwork'))\n\ntry:\n    import lightglue\nexcept ImportError:\n    import site\n    sys.path.append(site.getsitepackages()[0])\n\nimport pycolmap\nfrom hloc import extract_features, match_features, reconstruction\nfrom hloc.utils.read_write_model import read_model\n\n# --- CONFIG ---\nKAGGLE_INPUT_DIR = Path('/kaggle/input/image-matching-challenge-2025')\nTRAIN_TEST_DIR = KAGGLE_INPUT_DIR / 'train' \nOUTPUT_DIR = Path('/kaggle/working/outputs')\n\nif OUTPUT_DIR.exists(): shutil.rmtree(OUTPUT_DIR)\nOUTPUT_DIR.mkdir(exist_ok=True, parents=True)\n\n# --- WORKER FUNCTION (PARALLEL GPU) ---\ndef gpu_worker(gpu_id, datasets):\n    # 1. MASK GPU\n    os.environ[\"CUDA_VISIBLE_DEVICES\"] = str(gpu_id)\n    device = 'cuda'\n    print(f\"🎮 Worker {gpu_id}: Started processing {len(datasets)} datasets.\")\n\n    # 2. LOAD MODELS\n    try:\n        dino = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitl14').to(device).eval()\n    except:\n        print(f\"⚠️ Worker {gpu_id}: Failed to load DINOv2. Proceeding without rotation check.\")\n        dino = None\n\n    transform = T.Compose([\n        T.Resize((224, 224)), T.ToTensor(),\n        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n    ])\n\n    # HLOC Configs\n    conf_ex = extract_features.confs['aliked-n16']\n    conf_ex['model']['max_num_keypoints'] = 4096\n    conf_ex['preprocessing']['resize_max'] = 2048\n    conf_ma = match_features.confs['aliked+lightglue']\n\n    def correct_rotations(img_paths):\n        if not dino or len(img_paths) < 5: return\n        # Calculate Consensus\n        ref_feats = []\n        for p in img_paths[:5]:\n            try:\n                t = transform(Image.open(p).convert('RGB')).unsqueeze(0).to(device)\n                with torch.no_grad(): ref_feats.append(F.normalize(dino(t), p=2, dim=1))\n            except: pass\n        if not ref_feats: return\n        center = torch.mean(torch.cat(ref_feats), dim=0, keepdim=True)\n        \n        count = 0\n        for p in img_paths:\n            try:\n                img_org = Image.open(p).convert('RGB')\n                best_s, best_r, best_i = -1, 0, img_org\n                for r in [0, 90, 180, 270]:\n                    im = img_org if r==0 else img_org.rotate(r, expand=True)\n                    t = transform(im).unsqueeze(0).to(device)\n                    with torch.no_grad(): f = F.normalize(dino(t), p=2, dim=1)\n                    s = torch.mm(f, center.t()).item()\n                    if s > best_s: best_s, best_r, best_i = s, r, im\n                if best_r != 0:\n                    best_i.save(p)\n                    count += 1\n            except: pass\n        print(f\"   Worker {gpu_id}: Rotated {count} images.\")\n\n    # Dataset Loop\n    for ds_name in datasets:\n        src = Path(TRAIN_TEST_DIR) / ds_name\n        if not src.exists(): continue\n        \n        print(f\"🚀 Worker {gpu_id}: Starting {ds_name}...\")\n        try:\n            temp = Path(f'/kaggle/working/temp_{ds_name}_{gpu_id}')\n            out = OUTPUT_DIR / ds_name\n            out.mkdir(parents=True, exist_ok=True)\n            \n            if temp.exists(): shutil.rmtree(temp)\n            shutil.copytree(src, temp)\n            \n            images = sorted([p for p in temp.glob('**/*') if p.suffix.lower() in {'.jpg', '.png'}])\n            if len(images) < 3: continue\n            \n            correct_rotations(images)\n            \n            # PAIR GENERATION\n            pairs = []\n            rels = [p.relative_to(temp).as_posix() for p in images]\n            \n            if len(images) < 300: \n                print(f\"   Worker {gpu_id} [{ds_name}]: EXHAUSTIVE matching ({len(images)} images)\")\n                for i in range(len(rels)):\n                    for j in range(i+1, len(rels)):\n                        pairs.append((rels[i], rels[j]))\n            elif dino: \n                print(f\"   Worker {gpu_id} [{ds_name}]: DINOv2 matching\")\n                feats = []\n                for i in range(0, len(images), 16):\n                    batch = []\n                    for p in images[i:i+16]:\n                        batch.append(transform(Image.open(p).convert('RGB')))\n                    batch = torch.stack(batch).to(device)\n                    with torch.no_grad(): feats.append(F.normalize(dino(batch), p=2, dim=1))\n                all_f = torch.cat(feats)\n                sim = torch.mm(all_f, all_f.t())\n                for i in range(len(images)):\n                    sc = sim[i]\n                    valid = torch.where(sc > 0.15)[0]\n                    valid = valid[valid != i]\n                    if len(valid)>0:\n                        topk = valid[torch.topk(sc[valid], min(50, len(valid))).indices].tolist()\n                        for j in topk:\n                            if i<j: pairs.append((rels[i], rels[j]))\n                            elif j<i: pairs.append((rels[j], rels[i]))\n                pairs = sorted(list(set(pairs)))\n\n            p_path = out / 'pairs.txt'\n            with open(p_path, 'w') as f:\n                for p1, p2 in pairs: f.write(f\"{p1} {p2}\\n\")\n            \n            # HLOC PIPELINE\n            f_path = out / 'feats.h5'\n            m_path = out / 'matches.h5'\n            if f_path.exists(): f_path.unlink()\n            if m_path.exists(): m_path.unlink()\n            \n            img_list = [p.relative_to(temp).as_posix() for p in images]\n            extract_features.main(conf_ex, temp, feature_path=f_path, image_list=img_list)\n            match_features.main(conf_ma, pairs=p_path, features=f_path, matches=m_path)\n            reconstruction.main(out / 'colmap', temp, pairs=p_path, features=f_path, matches=m_path, camera_mode='AUTO', verbose=False)\n            \n            if temp.exists(): shutil.rmtree(temp)\n            print(f\"✅ Worker {gpu_id}: Finished {ds_name}\")\n            \n        except Exception as e:\n            print(f\"❌ Worker {gpu_id} Error on {ds_name}: {e}\")\n\n# --- MAIN CONTROLLER ---\nif __name__ == \"__main__\":\n    print(\"🚀 Main Pipeline Started...\")\n    \n    # 1. IDENTIFY DATASETS\n    all_datasets = sorted([d.name for d in TRAIN_TEST_DIR.iterdir() if d.is_dir()])\n    # Filter out stairs\n    all_datasets = [d for d in all_datasets if 'stairs' not in d]\n    \n    # 2. SPLIT WORKLOAD\n    mid = len(all_datasets) // 2 + 1\n    g0 = all_datasets[:mid]\n    g1 = all_datasets[mid:]\n    \n    print(f\"📋 Workload: GPU 0 has {len(g0)} datasets, GPU 1 has {len(g1)} datasets.\")\n    \n    # 3. LAUNCH WORKERS\n    mp.set_start_method('spawn', force=True)\n    p0 = mp.Process(target=gpu_worker, args=(0, g0))\n    p1 = mp.Process(target=gpu_worker, args=(1, g1))\n    \n    p0.start(); p1.start()\n    p0.join(); p1.join()\n    \n    print(\"\\n🛠️ ALL PROCESSING DONE. MERGING RESULTS...\")\n\n    # --- MERGER LOGIC (High Accuracy) ---\n    def qvec2rotmat_np(qvec):\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\n    submission_data = {}\n    solved_count = 0\n    \n    for dataset_dir in OUTPUT_DIR.iterdir():\n        if not dataset_dir.is_dir(): continue\n        ds_name = dataset_dir.name\n        \n        # FIND ALL MODELS (0, 1, 2...)\n        model_files = list(dataset_dir.glob('**/images.bin'))\n        \n        if model_files:\n            print(f\"   📂 {ds_name}: Merging {len(model_files)} fragments...\")\n            for m_file in model_files:\n                try:\n                    _, images, _ = read_model(m_file.parent, ext='.bin')\n                    for _, img in images.items():\n                        if (ds_name, img.name) not in submission_data:\n                            R = qvec2rotmat_np(img.qvec)\n                            submission_data[(ds_name, img.name)] = (\n                                \";\".join(map(str, R.flatten())),\n                                \";\".join(map(str, img.tvec.flatten()))\n                            )\n                            solved_count += 1\n                except: pass\n\n    # --- SAVE CSV ---\n    sample_path = KAGGLE_INPUT_DIR / 'sample_submission.csv'\n    if sample_path.exists():\n        df = pd.read_csv(sample_path)\n        out_r, out_t = [], []\n        for idx, row in df.iterrows():\n            key = (row['dataset'], row['image'])\n            if key in submission_data:\n                out_r.append(submission_data[key][0])\n                out_t.append(submission_data[key][1])\n            else:\n                out_r.append(\"1.0;0.0;0.0;0.0;1.0;0.0;0.0;0.0;1.0\")\n                out_t.append(\"0.0;0.0;0.0\")\n        df['rotation_matrix'] = out_r\n        df['translation_vector'] = out_t\n        df.to_csv('submission.csv', index=False)\n        print(f\"\\n🎉 SUCCESS! Final submission.csv saved. Solved {solved_count} images.\")\n    else:\n        # Fallback for dry run\n        print(\"No sample submission found. Saving raw data.\")\n        rows = [{'dataset': k[0], 'image': k[1], 'R': v[0], 'T': v[1]} for k, v in submission_data.items()]\n        pd.DataFrame(rows).to_csv('submission.csv', index=False)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- RUN THE PIPELINE ---\nprint(\"🚀 Launching main.py in a fresh process...\")\n!python main.py","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}