romav2 2.0.0__py3-none-any.whl

romav2/__init__.py

@@ -0,0 +1,8 @@
+import logging as _logging
+from .logging import configure_logger as configure_logger
+from .logging import logger as _logger
+
+if not any(not isinstance(h, _logging.NullHandler) for h in _logger.handlers):
+    configure_logger()
+
+from .romav2 import RoMaV2 as RoMaV2

romav2/benchmarks/__init__.py

@@ -0,0 +1,4 @@
+from .scannet1500 import ScanNet1500 as ScanNet1500
+from .mega1500 import Mega1500 as Mega1500
+from .wxbs import WxBSBenchmark as WxBSBenchmark
+from .satast import SatAst as SatAst

romav2/benchmarks/mega1500.py

@@ -0,0 +1,115 @@
+import numpy as np
+import logging
+from romav2.geometry import (
+    compute_pose_error,
+    pose_auc,
+    estimate_pose_cv2_ransac,
+    compute_relative_pose,
+)
+from PIL import Image
+from tqdm import tqdm
+
+logger = logging.getLogger(__name__)
+
+
+class Mega1500:
+    def __init__(self, data_root="data/megadepth") -> None:
+        self.scene_names = [
+            "0015_0.1_0.3.npz",
+            "0015_0.3_0.5.npz",
+            "0022_0.1_0.3.npz",
+            "0022_0.3_0.5.npz",
+            "0022_0.5_0.7.npz",
+        ]
+        self.scenes = [
+            np.load(f"{data_root}/{scene}", allow_pickle=True)
+            for scene in self.scene_names
+        ]
+        self.data_root = data_root
+
+    def benchmark(self, model, model_name=None):
+        data_root = self.data_root
+        tot_e_t, tot_e_R, tot_e_pose = [], [], []
+        thresholds = [5, 10, 20]
+        for scene_ind in range(len(self.scenes)):
+            scene = self.scenes[scene_ind]
+            pairs = scene["pair_infos"]
+            intrinsics = scene["intrinsics"]
+            poses = scene["poses"]
+            im_paths = scene["image_paths"]
+            pair_inds = range(len(pairs))
+            for pairind in (pbar := tqdm(pair_inds, desc="Mega1500 eval")):
+                idx1, idx2 = pairs[pairind][0]
+                K1 = intrinsics[idx1].copy()
+                T1 = poses[idx1].copy()
+                R1, t1 = T1[:3, :3], T1[:3, 3]
+                K2 = intrinsics[idx2].copy()
+                T2 = poses[idx2].copy()
+                R2, t2 = T2[:3, :3], T2[:3, 3]
+                R, t = compute_relative_pose(R1, t1, R2, t2)
+                im_A_path = f"{data_root}/{im_paths[idx1]}"
+                im_B_path = f"{data_root}/{im_paths[idx2]}"
+                preds = model.match(im_A_path, im_B_path)
+                im_A = Image.open(im_A_path)
+                w1, h1 = im_A.size
+                im_B = Image.open(im_B_path)
+                w2, h2 = im_B.size
+                if True:  # Note: we keep this true as it was used in DKM/RoMa papers. There is very little difference compared to setting to False.
+                    scale1 = 1200 / max(w1, h1)
+                    scale2 = 1200 / max(w2, h2)
+                    w1, h1 = scale1 * w1, scale1 * h1
+                    w2, h2 = scale2 * w2, scale2 * h2
+                    K1, K2 = K1.copy(), K2.copy()
+                    K1[:2] = K1[:2] * scale1
+                    K2[:2] = K2[:2] * scale2
+                for _ in range(5):
+                    matches, _, _, _ = model.sample(preds, 5_000)
+                    kpts1, kpts2 = model.to_pixel_coordinates(matches, h1, w1, h2, w2)
+                    kpts1, kpts2 = kpts1.cpu().numpy(), kpts2.cpu().numpy()
+                    shuffling = np.random.permutation(np.arange(len(kpts1)))
+                    kpts1 = kpts1[shuffling]
+                    kpts2 = kpts2[shuffling]
+                    try:
+                        threshold = 0.5
+                        norm_threshold = threshold / (
+                            np.mean(np.abs(K1[:2, :2])) + np.mean(np.abs(K2[:2, :2]))
+                        )
+                        R_est, t_est, _ = estimate_pose_cv2_ransac(
+                            kpts1,
+                            kpts2,
+                            K1,
+                            K2,
+                            norm_threshold,
+                            conf=0.99999,
+                        )
+                        e_t, e_R = compute_pose_error(R_est, t_est[:, 0], R, t)
+                        e_pose = max(e_t, e_R)
+                    except Exception as e:
+                        logger.debug(f"Pose estimation error: {e}")
+                        e_t, e_R = 90, 90
+                        e_pose = max(e_t, e_R)
+                    tot_e_t.append(e_t)
+                    tot_e_R.append(e_R)
+                    tot_e_pose.append(e_pose)
+                pbar.set_postfix(
+                    auc=f"{[f'{a.item():.3f}' for a in pose_auc(tot_e_pose, thresholds)]}"
+                )
+
+        tot_e_pose = np.array(tot_e_pose)
+        auc = pose_auc(tot_e_pose, thresholds)
+        acc_5 = (tot_e_pose < 5).mean()
+        acc_10 = (tot_e_pose < 10).mean()
+        acc_15 = (tot_e_pose < 15).mean()
+        acc_20 = (tot_e_pose < 20).mean()
+        map_5 = acc_5
+        map_10 = np.mean([acc_5, acc_10])
+        map_20 = np.mean([acc_5, acc_10, acc_15, acc_20])
+        logger.info("%s auc: %s", model_name, auc)
+        return {
+            "auc_5": auc[0],
+            "auc_10": auc[1],
+            "auc_20": auc[2],
+            "map_5": map_5,
+            "map_10": map_10,
+            "map_20": map_20,
+        }

romav2/benchmarks/satast.py

@@ -0,0 +1,463 @@
+import os
+import logging
+import torch
+import tempfile
+import json
+import numpy as np
+import cv2
+from PIL import Image
+from tqdm import tqdm
+from romav2.geometry import pose_auc
+import matplotlib.pyplot as plt
+
+OFFSET = 0.5
+
+logger = logging.getLogger(__name__)
+
+
+class SatAst:
+    """
+    A benchmark class to evaluate image matchers using a folder of JSON files.
+
+    The final metric is the Area Under the Curve (AUC) of the cumulative
+    distribution of *all* individual errors, calculated at different
+    pixel thresholds.
+
+    Args:
+        json_folder_path (str): Path to the directory containing the .json files.
+        image_dataset_root (str): The root directory to prepend to the
+        image paths (query_path, pred_path)
+        found inside the JSON files.
+    """
+
+    def __init__(
+        self,
+        json_folder_path="data/satast/annotations-v1",
+        image_dataset_root="data/satast",
+    ) -> None:
+        self.json_folder_path = json_folder_path
+        self.image_dataset_root = image_dataset_root
+
+        try:
+            self.json_files = sorted(
+                [f for f in os.listdir(json_folder_path) if f.endswith(".json")]
+            )
+            if not self.json_files:
+                logger.warning("No .json files found in %s", json_folder_path)
+        except FileNotFoundError:
+            logger.error("Directory not found: %s", json_folder_path)
+            self.json_files = []
+
+        logger.info("Found %d JSON files.", len(self.json_files))
+
+    def _pixel_to_normalized(self, pts_pix, w, h, offset=OFFSET):
+        """
+        Converts pixel coordinates [0, n-1] to normalized [-1, 1].
+        """
+        pts_norm = np.zeros_like(pts_pix, dtype=np.float64)
+        pts_norm[..., 0] = (2.0 * (pts_pix[..., 0] + offset) / w) - 1.0
+        pts_norm[..., 1] = (2.0 * (pts_pix[..., 1] + offset) / h) - 1.0
+        return pts_norm
+
+    def _normalized_to_pixel(self, pts_norm, w, h, offset=OFFSET):
+        """
+        Converts normalized coordinates [-1, 1] to pixel [0, n-1].
+        """
+        pts_pix = np.zeros_like(pts_norm, dtype=np.float64)
+        pts_pix[..., 0] = (w * (pts_norm[..., 0] + 1.0) / 2.0) - offset
+        pts_pix[..., 1] = (h * (pts_norm[..., 1] + 1.0) / 2.0) - offset
+        return pts_pix
+
+    def benchmark(self, model, model_name=None, visualize=False, num_visualize=10):
+        """
+        Runs the full benchmark on the provided model.
+        This version benchmarks correspondences through an estimated homography.
+
+        Args:
+            model: The matcher model object. Must have a 'match(im_A_path, im_B_path)'
+            method and a 'sample(matches, certainty, num_samples)' method.
+            Assumed to return matches in normalized [-1, 1] coordinates.
+            model_name (str, optional): A name for the model.
+            visualize (bool, optional): If True, plots the first `num_visualize`
+            successful homography estimations.
+            num_visualize (int, optional): The maximum number of pairs to visualize.
+
+        Returns:
+            dict: A dictionary containing the AUC scores.
+        """
+        # This list will store *all* individual reprojection errors
+        all_reprojection_errors = []
+        viz_counter = 0
+        temp_file = tempfile.NamedTemporaryFile(suffix=".jpeg", delete=False)
+        im_B_path = temp_file.name
+
+        for json_name in tqdm(self.json_files, desc="Running Benchmark"):
+            json_path = os.path.join(self.json_folder_path, json_name)
+
+            # 1. Load JSON and Image Paths
+            with open(json_path, "r") as f:
+                json_data = json.load(f)
+
+            im_A_path_rel = json_data["query_path"].replace("\\", "/")
+            im_B_path_rel = json_data["pred_path"].replace("\\", "/")
+
+            im_A_path = os.path.join(self.image_dataset_root, im_A_path_rel)
+            im_B_path0 = os.path.join(self.image_dataset_root, im_B_path_rel)
+
+            im_A = Image.open(im_A_path)
+            w1, h1 = im_A.size
+            im_B0 = Image.open(im_B_path0)
+            # if "__rot0__" not in im_B_path0:
+            #     continue
+            w2, h2 = im_B0.size
+
+            for rot_idx in range(1):
+                im_B0 = Image.open(im_B_path0)
+
+                if rot_idx == 0:
+                    rotated_im_B = im_B0
+                elif rot_idx == 1:
+                    rotated_im_B = im_B0.transpose(method=Image.Transpose.ROTATE_90)
+                elif rot_idx == 2:
+                    rotated_im_B = im_B0.transpose(method=Image.Transpose.ROTATE_180)
+                elif rot_idx == 3:
+                    rotated_im_B = im_B0.transpose(method=Image.Transpose.ROTATE_270)
+                rotated_im_B.save(im_B_path, format="JPEG")
+
+                # 2. Run the matcher model
+                dense_preds = model.match(im_A_path, im_B_path)
+                if isinstance(dense_preds, dict):
+                    sparse_preds = model.sample(dense_preds, 10_000)
+                else:
+                    sparse_preds = model.sample(*dense_preds, 10_000)
+                good_matches = sparse_preds[0]
+                pos_a_norm = good_matches[:, :2]
+                pos_b_norm = good_matches[:, 2:]
+
+                # Rotate back
+                if rot_idx == 1:
+                    pos_b_norm = torch.stack(
+                        [-pos_b_norm[:, 1], pos_b_norm[:, 0]], dim=1
+                    )
+                elif rot_idx == 2:
+                    pos_b_norm = -pos_b_norm
+                elif rot_idx == 3:
+                    pos_b_norm = torch.stack(
+                        [pos_b_norm[:, 1], -pos_b_norm[:, 0]], dim=1
+                    )
+
+                # 3. Estimate Homography in Normalized Coordinates
+                H_pred_norm = None
+                if len(pos_a_norm) >= 4:
+                    try:
+                        norm_thresh = 0.001
+
+                        H_pred_norm, inliers = cv2.findHomography(
+                            pos_a_norm.cpu().numpy(),
+                            pos_b_norm.cpu().numpy(),
+                            method=cv2.USAC_DEFAULT,
+                            confidence=0.99999999,
+                            maxIters=100_000,
+                            ransacReprojThreshold=norm_thresh,
+                        )
+                    except cv2.error as e:
+                        print(e)
+                        H_pred_norm = None
+
+                # 4. Get Ground Truth Correspondences (Pixel)
+                last_iteration_corrs = json_data["correspondences"][-1]
+                pts_src_pix = np.array(last_iteration_corrs["pts_src"])
+                pts_dst_pix = np.array(last_iteration_corrs["pts_dst"])
+
+                if pts_src_pix.shape[0] == 0:
+                    continue  # No GT points, skip this pair
+
+                if H_pred_norm is None:
+                    # If RANSAC fails, append 'inf' for all GT points
+                    all_reprojection_errors.extend(
+                        [float("inf")] * pts_src_pix.shape[0]
+                    )
+                    continue
+
+                if visualize and viz_counter < num_visualize:
+                    logger.info("Visualizing pair: %s", json_name)
+
+                    # 1. Create pixel-to-norm (M1) and norm-to-pixel (M2_inv) matrices
+
+                    # M1: P1 -> N1
+                    M1 = np.array(
+                        [
+                            [2.0 / w1, 0, (2.0 * OFFSET / w1) - 1.0],
+                            [0, 2.0 / h1, (2.0 * OFFSET / h1) - 1.0],
+                            [0, 0, 1.0],
+                        ],
+                        dtype=np.float64,
+                    )
+
+                    # M2_inv: N2 -> P2
+                    M2_inv = np.array(
+                        [
+                            [w2 / 2.0, 0, w2 / 2.0 - OFFSET],
+                            [0, h2 / 2.0, h2 / 2.0 - OFFSET],
+                            [0, 0, 1.0],
+                        ],
+                        dtype=np.float64,
+                    )
+
+                    # 2. Convert H_pred_norm (N1 -> N2) to H_pix (P1 -> P2)
+                    H_pix = M2_inv @ H_pred_norm @ M1
+
+                    # 3. Invert H_pix to get (P2 -> P1) for warping Image 2 to 1
+                    try:
+                        H_pix_inv = np.linalg.inv(H_pix)
+                    except np.linalg.LinAlgError:
+                        logger.warning(
+                            "Could not invert homography for %s. Skipping visualization.",
+                            json_name,
+                        )
+                        continue  # Skip this viz
+
+                    # --- NEW: Extract and sample inliers for visualization ---
+                    plot_pts_a = None
+                    plot_pts_b = None
+                    plot_colors = None
+
+                    # Get inlier points from the RANSAC
+                    inliers_mask = inliers.flatten().astype(bool)
+                    inlier_pts_a_norm = pos_a_norm.cpu().numpy()[inliers_mask]
+                    inlier_pts_b_norm = pos_b_norm.cpu().numpy()[inliers_mask]
+                    num_inliers = inlier_pts_a_norm.shape[0]
+
+                    if num_inliers > 0:
+                        # Sample up to 10 inliers
+                        num_to_sample = min(num_inliers, 10)
+                        indices = np.random.choice(
+                            num_inliers, num_to_sample, replace=False
+                        )
+
+                        sampled_pts_a_norm = inlier_pts_a_norm[indices]
+                        sampled_pts_b_norm = inlier_pts_b_norm[indices]
+
+                        # Convert to pixel coordinates for plotting
+                        plot_pts_a = self._normalized_to_pixel(
+                            sampled_pts_a_norm, w1, h1
+                        )
+                        plot_pts_b = self._normalized_to_pixel(
+                            sampled_pts_b_norm, w2, h2
+                        )
+
+                        # Generate random colors for each pair
+                        plot_colors = np.random.rand(num_to_sample, 3)
+
+                    # 4. Load images with OpenCV for warping
+                    img1_cv = cv2.imread(im_A_path)  # Target
+                    img2_cv = cv2.imread(im_B_path0)  # Source
+
+                    # 5. Warp image 2 to image 1's frame
+                    warped_img2 = cv2.warpPerspective(img2_cv, H_pix_inv, (w1, h1))
+
+                    # 6. Plot
+                    img1_rgb = cv2.cvtColor(img1_cv, cv2.COLOR_BGR2RGB)
+                    img2_rgb = cv2.cvtColor(img2_cv, cv2.COLOR_BGR2RGB)
+                    warped_img2_rgb = cv2.cvtColor(warped_img2, cv2.COLOR_BGR2RGB)
+
+                    fig, axes = plt.subplots(1, 3, figsize=(21, 7))
+                    fig.suptitle(
+                        f"Homography Visualization: {os.path.basename(im_B_path0)} -> {os.path.basename(im_A_path)}",
+                        fontsize=16,
+                    )
+
+                    axes[0].imshow(img1_rgb)
+                    axes[0].set_title("Image 1 (Target)")
+                    axes[0].axis("off")
+                    if plot_pts_a is not None:
+                        axes[0].scatter(
+                            plot_pts_a[:, 0],
+                            plot_pts_a[:, 1],
+                            s=40,
+                            c=plot_colors,
+                            marker="x",
+                        )
+
+                    axes[1].imshow(img2_rgb)
+                    axes[1].set_title("Image 2 (Source)")
+                    axes[1].axis("off")
+                    if plot_pts_b is not None:
+                        axes[1].scatter(
+                            plot_pts_b[:, 0],
+                            plot_pts_b[:, 1],
+                            s=40,
+                            c=plot_colors,
+                            marker="x",
+                        )
+
+                    axes[2].imshow(warped_img2_rgb)
+                    axes[2].set_title("Image 2 Warped to Image 1")
+                    axes[2].axis("off")
+
+                    plt.tight_layout()
+                    plt.show()
+
+                    viz_counter += 1
+
+                # 5. Normalize GT Source Points
+                pts_src_norm = self._pixel_to_normalized(pts_src_pix, w1, h1)
+
+                # 6. Warp Normalized GT Points
+                pts_src_norm_h = np.hstack(
+                    (pts_src_norm, np.ones((pts_src_norm.shape[0], 1)))
+                )
+                warped_pts_norm_h = np.dot(pts_src_norm_h, H_pred_norm.T)
+
+                # De-homogenize
+                epsilon = 1e-8
+                warped_pts_norm = warped_pts_norm_h[:, :2] / (
+                    warped_pts_norm_h[:, 2, np.newaxis] + epsilon
+                )
+
+                # 7. Convert Warped Points back to Pixel Coords
+                warped_pts_pix = self._normalized_to_pixel(warped_pts_norm, w2, h2)
+
+                # 8. Calculate and Store All Errors
+                errors = np.linalg.norm(warped_pts_pix - pts_dst_pix, axis=1)
+                all_reprojection_errors.extend(errors)
+
+        temp_file.close()
+
+        # 9. Compute AUC over *all* errors
+        thresholds = np.arange(1, 31)
+        auc = pose_auc(np.array(all_reprojection_errors), thresholds)
+        logger.info("=== AUC OF ERRORS ===")
+        logger.info("%s", auc)
+
+        results = {
+            "reprojection_auc_5px": auc[4],
+            "reprojection_auc_10px": auc[9],
+            "reprojection_auc_20px": auc[19],
+            "reprojection_auc_30px": auc[29],
+        }
+        logger.info("=== MAIN RESULTS ===")
+        logger.info("%s", results)
+
+        return results
+
+    def benchmark_warp(self, model, model_name=None, visualize=False, num_visualize=10):
+        """
+        Runs the full benchmark on the provided model.
+        This version benchmarks correspondences through an estimated dense warp.
+
+        Args:
+            model: The matcher model object. Must have a 'match(im_A_path, im_B_path)'
+            method and a 'sample(matches, certainty, num_samples)' method.
+            Assumed to return matches in normalized [-1, 1] coordinates.
+            model_name (str, optional): A name for the model.
+            visualize (bool, optional): If True, plots the first `num_visualize`
+            successful homography estimations.
+            num_visualize (int, optional): The maximum number of pairs to visualize.
+
+        Returns:
+            dict: A dictionary containing the AUC scores.
+        """
+        # This list will store *all* individual reprojection errors
+        all_reprojection_errors = []
+        temp_file = tempfile.NamedTemporaryFile(suffix=".jpeg", delete=False)
+        im_B_path = temp_file.name
+
+        for json_name in tqdm(self.json_files, desc="Running Benchmark"):
+            json_path = os.path.join(self.json_folder_path, json_name)
+
+            # 1. Load JSON and Image Paths
+            with open(json_path, "r") as f:
+                json_data = json.load(f)
+
+            im_A_path_rel = json_data["query_path"].replace("\\", "/")
+            im_B_path_rel = json_data["pred_path"].replace("\\", "/")
+
+            im_A_path = os.path.join(self.image_dataset_root, im_A_path_rel)
+            im_B_path0 = os.path.join(self.image_dataset_root, im_B_path_rel)
+
+            im_A = Image.open(im_A_path)
+            w1, h1 = im_A.size
+            im_B0 = Image.open(im_B_path0)
+            w2, h2 = im_B0.size
+
+            for rot_idx in range(4):
+                im_B0 = Image.open(im_B_path0)
+
+                if rot_idx == 0:
+                    rotated_im_B = im_B0
+                elif rot_idx == 1:
+                    rotated_im_B = im_B0.transpose(method=Image.Transpose.ROTATE_90)
+                elif rot_idx == 2:
+                    rotated_im_B = im_B0.transpose(method=Image.Transpose.ROTATE_180)
+                elif rot_idx == 3:
+                    rotated_im_B = im_B0.transpose(method=Image.Transpose.ROTATE_270)
+                rotated_im_B.save(im_B_path, format="JPEG")
+
+                # 2. Run the matcher model
+                preds = model.match(im_A_path, im_B_path)
+                warp = preds["warp_AB"]
+
+                # 3. Get Ground Truth Correspondences (Pixel)
+                last_iteration_corrs = json_data["correspondences"][-1]
+                pts_src_pix = np.array(last_iteration_corrs["pts_src"])
+                pts_dst_pix = np.array(last_iteration_corrs["pts_dst"])
+
+                if pts_src_pix.shape[0] == 0:
+                    continue  # No GT points, skip this pair
+
+                # 5. Normalize GT Source Points
+                pts_src_norm = torch.tensor(
+                    self._pixel_to_normalized(pts_src_pix, w1, h1),
+                    device=warp.device,
+                    dtype=warp.dtype,
+                )
+
+                # 6. Warp Normalized GT Points
+                warp_src_to_pred = warp[0, :, : warp.shape[2] // 2, -2:]
+                warped_pts_norm = torch.nn.functional.grid_sample(
+                    warp_src_to_pred.permute(2, 0, 1)[None],
+                    pts_src_norm[None, None],
+                    align_corners=False,
+                    mode="bilinear",
+                )[0, :, 0].mT
+
+                # Rotate back
+                if rot_idx == 1:
+                    warped_pts_norm = torch.stack(
+                        [-warped_pts_norm[:, 1], warped_pts_norm[:, 0]], dim=1
+                    )
+                elif rot_idx == 2:
+                    warped_pts_norm = -warped_pts_norm
+                elif rot_idx == 3:
+                    warped_pts_norm = torch.stack(
+                        [warped_pts_norm[:, 1], -warped_pts_norm[:, 0]], dim=1
+                    )
+
+                # 7. Convert Warped Points back to Pixel Coords
+                warped_pts_pix = self._normalized_to_pixel(
+                    warped_pts_norm.cpu().numpy(), w2, h2
+                )
+
+                # 8. Calculate and Store All Errors
+                errors = np.linalg.norm(warped_pts_pix - pts_dst_pix, axis=1)
+                all_reprojection_errors.extend(errors)
+
+        temp_file.close()
+
+        # 9. Compute AUC over *all* errors
+        thresholds = np.arange(1, 31)
+        auc = pose_auc(np.array(all_reprojection_errors), thresholds)
+        logger.info("=== AUC OF ERRORS ===")
+        logger.info("%s", auc)
+
+        results = {
+            "reprojection_auc_5px": auc[4],
+            "reprojection_auc_10px": auc[9],
+            "reprojection_auc_20px": auc[19],
+            "reprojection_auc_30px": auc[29],
+        }
+        logger.info("=== MAIN RESULTS ===")
+        logger.info("%s", results)
+
+        return results