{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.11.13"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":91498,"databundleVersionId":11655853,"sourceType":"competition"},{"sourceId":7884485,"sourceType":"datasetVersion","datasetId":4628051},{"sourceId":8026384,"sourceType":"datasetVersion","datasetId":4726252},{"sourceId":11575619,"sourceType":"datasetVersion","datasetId":7257638},{"sourceId":11754775,"sourceType":"datasetVersion","datasetId":7379527},{"sourceId":11924468,"sourceType":"datasetVersion","datasetId":6988459},{"sourceId":13961320,"sourceType":"datasetVersion","datasetId":8899689},{"sourceId":4534,"sourceType":"modelInstanceVersion","modelInstanceId":3326,"modelId":986},{"sourceId":17191,"sourceType":"modelInstanceVersion","modelInstanceId":14317,"modelId":21716},{"sourceId":17555,"sourceType":"modelInstanceVersion","modelInstanceId":14611,"modelId":22086},{"sourceId":370254,"sourceType":"modelInstanceVersion","modelInstanceId":128173,"modelId":2807}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true},"papermill":{"default_parameters":{},"duration":61.987086,"end_time":"2025-12-02T19:52:20.117521","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2025-12-02T19:51:18.130435","version":"2.6.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"ceb981ff","cell_type":"code","source":"!git clone --recursive https://github.com/naver/mast3r\n%cd mast3r\n!pip install -r requirements.txt\n!pip install -e .\n!pip install roma","metadata":{"papermill":{"duration":0.00494,"end_time":"2025-12-02T19:52:16.662017","exception":false,"start_time":"2025-12-02T19:52:16.657077","status":"completed"},"tags":[],"trusted":true,"execution":{"iopub.status.busy":"2025-12-08T19:21:26.491332Z","iopub.execute_input":"2025-12-08T19:21:26.491632Z","iopub.status.idle":"2025-12-08T19:21:40.182263Z","shell.execute_reply.started":"2025-12-08T19:21:26.491611Z","shell.execute_reply":"2025-12-08T19:21:40.181322Z"}},"outputs":[],"execution_count":null},{"id":"7cf37355-926b-4a1a-86b0-1fd2d56f4251","cell_type":"code","source":"import os\nimport gc\nimport torch\nimport numpy as np\nimport pandas as pd\nfrom mast3r.model import AsymmetricMASt3R\nfrom dust3r.inference import inference\nfrom dust3r.utils.image import load_images\nfrom dust3r.image_pairs import make_pairs\nfrom dust3r.cloud_opt import global_aligner, GlobalAlignerMode\n\nIMG_SIZE = 448\nMAX_IMAGES = 32\nWINDOW_SIZE = 3\nBATCH_SIZE = 1\nDEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\nN_RANSAC_RUNS = 7\n\nKAGGLE_INPUT_ROOT = '/kaggle/input/image-matching-challenge-2025/train'\nGT_CSV_PATH = '/kaggle/input/image-matching-challenge-2025/train_labels.csv'\n\ndef get_filenames_from_disk(scene_path):\n    exts = ('.png', '.jpg', '.jpeg', '.bmp', '.tif')\n    return sorted([f for f in os.listdir(scene_path) if f.lower().endswith(exts)])\n\ndef get_gt_poses(gt_df, scene_name):\n    mask = gt_df['dataset'].astype(str).str.lower() == scene_name.lower()\n    data = gt_df[mask]\n    poses = {}\n    \n    for _, row in data.iterrows():\n        name = str(row['image'])\n        basename = name.split('/')[-1]\n        \n        try:\n            R = np.array([float(x) for x in str(row['rotation_matrix']).split(';')]).reshape(3,3)\n            t = np.array([float(x) for x in str(row['translation_vector']).split(';')]).reshape(3,1)\n            pose = np.eye(4)\n            pose[:3, :3] = R.T\n            pose[:3, 3] = (-R.T @ t).flatten()\n            # Store with both full path and basename\n            poses[basename] = pose\n            poses[name] = pose\n        except:\n            continue\n    return poses\n\ndef smart_load_images(path, size, max_images):\n    fnames = get_filenames_from_disk(path)\n    if len(fnames) > max_images:\n        fnames = fnames[:max_images]\n    \n    loaded_imgs = load_images(path, size=size, verbose=False)\n    \n    if len(loaded_imgs) > len(fnames):\n        loaded_imgs = loaded_imgs[:len(fnames)]\n    \n    final_imgs = []\n    for img_obj in loaded_imgs:\n        img_tensor = img_obj['img']\n        _, _, h, w = img_tensor.shape\n        new_h = h - (h % 16)\n        new_w = w - (w % 16)\n        if new_h != h or new_w != w:\n            img_obj['img'] = img_tensor[:, :, :new_h, :new_w]\n        final_imgs.append(img_obj)\n    \n    return final_imgs, fnames\n\ndef multi_run_ransac_alignment(pred_poses, gt_poses, n_runs=7):\n    \"\"\"Multi-run RANSAC with better filename matching\"\"\"\n    common = sorted(list(set(pred_poses.keys()) & set(gt_poses.keys())))\n    \n    # Strategy 1: Direct match\n    if len(common) >= 5:\n        pass  # Use direct match\n    else:\n        # Strategy 2: Basename matching\n        gt_basenames = {k.split('/')[-1]: v for k, v in gt_poses.items()}\n        pred_basenames = {k.split('/')[-1]: v for k, v in pred_poses.items()}\n        common = sorted(list(set(pred_basenames.keys()) & set(gt_basenames.keys())))\n        \n        if len(common) >= 5:\n            pred_poses = pred_basenames\n            gt_poses = gt_basenames\n        else:\n            print(f\"    WARNING: Only {len(common)} common images found\")\n            print(f\"    Pred keys sample: {list(pred_poses.keys())[:3]}\")\n            print(f\"    GT keys sample: {list(gt_poses.keys())[:3]}\")\n            return 0.0, 0\n    \n    P_pred = np.array([pred_poses[k][:3, 3] for k in common]).T\n    P_gt = np.array([gt_poses[k][:3, 3] for k in common]).T\n    \n    if np.isnan(P_pred).any() or np.isnan(P_gt).any():\n        valid = ~np.isnan(P_pred).any(0) & ~np.isnan(P_gt).any(0)\n        P_pred, P_gt = P_pred[:, valid], P_gt[:, valid]\n        common = [common[i] for i in range(len(common)) if valid[i]]\n        if len(common) < 5:\n            return 0.0, 0\n    \n    scene_scale = np.linalg.norm(P_gt.max(1) - P_gt.min(1))\n    ransac_threshold = max(0.05 * scene_scale, 0.5)\n    \n    best_maa = 0.0\n    best_n = 0\n    \n    for run_seed in range(n_runs):\n        np.random.seed(run_seed)\n        \n        best_res = None\n        best_inliers = -1\n        \n        for _ in range(5000):\n            try:\n                idx = np.random.choice(P_pred.shape[1], 3, replace=False)\n                src, dst = P_pred[:, idx], P_gt[:, idx]\n                \n                mu_s, mu_d = src.mean(1, keepdims=True), dst.mean(1, keepdims=True)\n                src_c, dst_c = src - mu_s, dst - mu_d\n                \n                s_val = np.sum(src_c**2)\n                if s_val < 1e-7:\n                    continue\n                s = np.sqrt(np.sum(dst_c**2) / s_val)\n                \n                H = src_c @ dst_c.T\n                U, _, Vt = np.linalg.svd(H)\n                R = Vt.T @ U.T\n                \n                R_norm = R.copy()\n                if np.linalg.det(R_norm) < 0:\n                    Vt_tmp = Vt.copy()\n                    Vt_tmp[2] *= -1\n                    R_norm = Vt_tmp.T @ U.T\n                t_norm = mu_d - s * (R_norm @ mu_s)\n                err_norm = np.linalg.norm((s * (R_norm @ P_pred) + t_norm) - P_gt, axis=0)\n                in_norm = np.sum(err_norm < ransac_threshold)\n                \n                R_ref = R.copy()\n                if np.linalg.det(R_ref) > 0:\n                    Vt_tmp = Vt.copy()\n                    Vt_tmp[2] *= -1\n                    R_ref = Vt_tmp.T @ U.T\n                t_ref = mu_d - s * (R_ref @ mu_s)\n                err_ref = np.linalg.norm((s * (R_ref @ P_pred) + t_ref) - P_gt, axis=0)\n                in_ref = np.sum(err_ref < ransac_threshold)\n                \n                if in_norm >= in_ref:\n                    if in_norm > best_inliers:\n                        best_inliers = in_norm\n                        best_res = (s, R_norm, t_norm)\n                else:\n                    if in_ref > best_inliers:\n                        best_inliers = in_ref\n                        best_res = (s, R_ref, t_ref)\n            except:\n                continue\n        \n        if best_res is None:\n            continue\n        \n        s, R_a, t_a = best_res\n        \n        errors = []\n        scene_center = P_gt.mean(1)\n        P_aligned = s * (R_a @ P_pred) + t_a\n        \n        for i in range(len(common)):\n            k = common[i]\n            G = gt_poses[k]\n            P = pred_poses[k]\n            \n            R_pred_aligned = R_a @ P[:3, :3]\n            R_diff = G[:3, :3].T @ R_pred_aligned\n            \n            if np.linalg.det(R_diff) < 0:\n                R_diff = R_diff @ np.diag([1, 1, -1])\n            \n            tr = np.clip(np.trace(R_diff), -1, 3)\n            r_err = np.degrees(np.arccos(np.clip((tr - 1) / 2, -1, 1)))\n            if np.isnan(r_err):\n                r_err = 180.0\n            \n            v_gt = G[:3, 3] - scene_center\n            v_pred = P_aligned[:, i] - scene_center\n            n_gt, n_pred = np.linalg.norm(v_gt), np.linalg.norm(v_pred)\n            \n            if n_gt < 1e-6 or n_pred < 1e-6:\n                t_err = 0.0\n            else:\n                dot = np.dot(v_gt, v_pred) / (n_gt * n_pred)\n                t_err = np.degrees(np.arccos(np.clip(dot, -1.0, 1.0)))\n            \n            errors.append(max(r_err, t_err))\n        \n        errors = np.array(errors)\n        acc = [np.mean(errors < t) for t in [3, 5, 10]]\n        maa = np.mean(acc)\n        \n        if maa > best_maa:\n            best_maa = maa\n            best_n = len(common)\n    \n    return best_maa, best_n\n\ndef run_pipeline():\n    \"\"\"Final optimized pipeline\"\"\"\n    print(f\"Loading MASt3R on {DEVICE}...\")\n    model = AsymmetricMASt3R.from_pretrained(\n        \"naver/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric\"\n    ).to(DEVICE)\n    \n    gt_df = pd.read_csv(GT_CSV_PATH)\n    scenes = sorted([\n        d for d in os.listdir(KAGGLE_INPUT_ROOT)\n        if os.path.isdir(os.path.join(KAGGLE_INPUT_ROOT, d))\n    ])\n    \n    results = []\n    print(f\"FINAL (Memory-Optimized): {IMG_SIZE}px | Max: {MAX_IMAGES} | Window: {WINDOW_SIZE} | RANSAC: {N_RANSAC_RUNS} ⭐\")\n    \n    for scene_id in scenes:\n        print(f\"\\n--- {scene_id} ---\")\n        try:\n            gc.collect()\n            torch.cuda.empty_cache()\n            \n            path = os.path.join(KAGGLE_INPUT_ROOT, scene_id)\n            imgs, fnames = smart_load_images(path, IMG_SIZE, MAX_IMAGES)\n            \n            if len(imgs) < 2:\n                print(f\"  Skipping: only {len(imgs)} images\")\n                continue\n            \n            pairs = make_pairs(imgs, scene_graph=f'swin-{WINDOW_SIZE}', symmetrize=True)\n            \n            with torch.amp.autocast(device_type='cuda', dtype=torch.float16):\n                output = inference(pairs, model, device=DEVICE, batch_size=BATCH_SIZE, verbose=False)\n            \n            scene = global_aligner(output, device=DEVICE, mode=GlobalAlignerMode.PointCloudOptimizer)\n            # Reduced iterations to save memory\n            scene.compute_global_alignment(init=\"mst\", niter=120, schedule='cosine', lr=0.01)\n            \n            poses = scene.get_im_poses()\n            pred_poses = {}\n            for i in range(len(imgs)):\n                name = fnames[i]\n                pred_poses[name] = poses[i].detach().float().cpu().numpy()\n            \n            gt_data = get_gt_poses(gt_df, scene_id)\n            score, n = multi_run_ransac_alignment(pred_poses, gt_data, n_runs=N_RANSAC_RUNS)\n            \n            print(f\"  mAA: {score:.4f} | Matches: {n}\")\n            results.append({'scene': scene_id, 'mAA': score, 'status': 'OK'})\n            \n            del scene, output, pairs, imgs, poses, pred_poses\n            gc.collect()\n            torch.cuda.empty_cache()\n            \n        except Exception as e:\n            print(f\"  Error: {str(e)[:100]}\")\n            import traceback\n            traceback.print_exc()\n            results.append({'scene': scene_id, 'mAA': 0.0, 'status': 'Error'})\n            gc.collect()\n            torch.cuda.empty_cache()\n    \n    df = pd.DataFrame(results)\n    print(\"\\n\" + \"=\"*60)\n    print(\"FINAL RESULTS:\")\n    print(\"=\"*60)\n    print(df)\n    if not df.empty:\n        print(f\"\\nCombined mAA: {df['mAA'].mean():.4f}\")\n    print(\"=\"*60)\n\n# if __name__ == \"__main__\":\n#     run_pipeline()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-08T19:21:40.183916Z","iopub.execute_input":"2025-12-08T19:21:40.184417Z","iopub.status.idle":"2025-12-08T19:21:52.087685Z","shell.execute_reply.started":"2025-12-08T19:21:40.184390Z","shell.execute_reply":"2025-12-08T19:21:52.086592Z"}},"outputs":[],"execution_count":null},{"id":"6abfe6e3-35b6-4077-8d17-14ed89138e8b","cell_type":"code","source":"run_pipeline()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-08T19:21:52.088591Z","iopub.execute_input":"2025-12-08T19:21:52.089049Z","iopub.status.idle":"2025-12-08T20:00:31.296007Z","shell.execute_reply.started":"2025-12-08T19:21:52.089027Z","shell.execute_reply":"2025-12-08T20:00:31.295357Z"}},"outputs":[],"execution_count":null}]}