initial commit

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,31 @@
+Software Copyright License for Academic Use of RIPE, Version 2.0
+
+© Copyright (2025) Fraunhofer-Gesellschaft zur Förderung der angewandten Forschung e.V.
+
+1. INTRODUCTION
+
+RIPE which means any source code, object code or binary files provided by Fraunhofer excluding third party software and materials, is made available under this Software Copyright License.
+
+2. COPYRIGHT LICENSE
+
+Internal use of RIPE, in source and binary forms, with or without modification, is permitted without payment of copyright license fees for non-commercial purposes of evaluation, testing and academic research.
+
+No right or license, express or implied, is granted to any part of RIPE except and solely to the extent as expressly set forth herein. Any commercial use or exploitation of RIPE and/or any modifications thereto under this license are prohibited.
+
+For any other use of RIPE than permitted by this software copyright license You need another license from Fraunhofer. In such case please contact Fraunhofer under the CONTACT INFORMATION below.
+
+3. LIMITED PATENT LICENSE
+
+If Fraunhofer patents are implemented by RIPE their use is permitted for internal non-commercial purposes of evaluation, testing and academic research. No patent grant is provided for any other use, including but not limited to commercial use or exploitation.
+
+Fraunhofer provides no warranty of patent non-infringement with respect to RIPE.
+
+4. DISCLAIMER
+
+RIPE is provided by Fraunhofer "AS IS" and WITHOUT ANY EXPRESS OR IMPLIED WARRANTIES, including but not limited to the implied warranties of fitness for a particular purpose. IN NO EVENT SHALL FRAUNHOFER BE LIABLE for any direct, indirect, incidental, special, exemplary, or consequential damages, including but not limited to procurement of substitute goods or services; loss of use, data, or profits, or business interruption, however caused and on any theory of liability, whether in contract, strict liability, or tort (including negligence), arising in any way out of the use of the Fraunhofer Software, even if advised of the possibility of such damage.
+
+5. CONTACT INFORMATION
+
+Fraunhofer-Institut für Nachrichtentechnik, Heinrich-Hertz-Institut, HHI
+Einsteinufer 37, 10587 Berlin, Germany
+info@hhi.fraunhofer.de

app.py

@@ -0,0 +1,268 @@
+# This is a small gradio interface to access our RIPE keypoint extractor.
+# You can either upload two images or use one of the example image pairs.
+
+import os
+
+import gradio as gr
+from PIL import Image
+
+from ripe import vgg_hyper
+
+SEED = 32000
+os.environ["PYTHONHASHSEED"] = str(SEED)
+
+import random
+from pathlib import Path
+
+import numpy as np
+import torch
+
+torch.manual_seed(SEED)
+np.random.seed(SEED)
+random.seed(SEED)
+import cv2
+import kornia.feature as KF
+import kornia.geometry as KG
+
+from ripe.utils.utils import cv2_matches_from_kornia, to_cv_kpts
+
+MIN_SIZE = 512
+MAX_SIZE = 768
+
+description_text = """
+<p align='center'>
+  <h1 align='center'>🌊🌺 ICCV 2025 🌺🌊</h1>
+  <p align='center'>
+    <a href='https://scholar.google.com/citations?user=ybMR38kAAAAJ'>Johannes Künzel</a> · 
+    <a href='https://scholar.google.com/citations?user=5yTuyGIAAAAJ'>Anna Hilsmann</a> · 
+    <a href='https://scholar.google.com/citations?user=BCElyCkAAAAJ'>Peter Eisert</a>
+  </p>
+  <h2 align='center'>
+    <a href='https://arxiv.org/abs/2507.04839'>Arxiv</a> | 
+    <a href='https://fraunhoferhhi.github.io/RIPE/'>Project Page</a> | 
+    <a href='https://github.com/fraunhoferhhi/RIPE'>Code</a>
+  </h2>
+</p>
+
+<br/>
+<div align='center'>
+
+### This demo showcases our new keypoint extractor model, RIPE (Reinforcement Learning on Unlabeled Image Pairs for Robust Keypoint Extraction).
+
+### RIPE is trained without requiring pose or depth supervision or artificial augmentations. By leveraging reinforcement learning, it learns to extract keypoints solely based on whether an image pair depicts the same scene or not.
+
+### For more detailed information, please refer to our [paper](link to be added).
+
+The demo code extracts the 2048-top keypoints from the two input images. It uses the mutual nearest neighbor (MNN) descriptor matcher from kornia to find matches between the two images.
+If the number of matches is greater than 8, it applies RANSAC to filter out outliers based on the inlier threshold provided by the user.
+Images are resized to fit within a maximum size of 2048x2048 pixels with maintained aspect ratio.
+
+</div>
+"""
+
+model = vgg_hyper("./weights_ripe.pth")
+
+
+def get_new_image_size(image, min_size=1600, max_size=2048):
+    """
+    Get a new size for the image that is scaled to fit between min_size and max_size while maintaining the aspect ratio.
+
+    Args:
+        image (PIL.Image): Input image.
+        min_size (int): Minimum allowed size for width and height.
+        max_size (int): Maximum allowed size for width and height.
+
+    Returns:
+        tuple: New size (width, height) for the image.
+    """
+    width, height = image.size
+
+    aspect_ratio = width / height
+    if width > height:
+        new_width = max(min_size, min(max_size, width))
+        new_height = int(new_width / aspect_ratio)
+    else:
+        new_height = max(min_size, min(max_size, height))
+        new_width = int(new_height * aspect_ratio)
+
+    new_size = (new_width, new_height)
+
+    return new_size
+
+
+def extract_keypoints(image1, image2, inl_th):
+    """
+    Extract keypoints from two input images using the RIPE model.
+
+    Args:
+        image1 (PIL.Image): First input image.
+        image2 (PIL.Image): Second input image.
+        inl_th (float): RANSAC inlier threshold.
+
+    Returns:
+        dict: A dictionary containing keypoints and matches.
+    """
+    log_text = "Extracting keypoints and matches with RIPE\n"
+
+    log_text += f"Image 1 size: {image1.size}\n"
+    log_text += f"Image 2 size: {image2.size}\n"
+
+    # check not larger than 2048x2048
+    new_size = get_new_image_size(image1, min_size=MIN_SIZE, max_size=MAX_SIZE)
+    image1 = image1.resize(new_size)
+
+    new_size = get_new_image_size(image2, min_size=MIN_SIZE, max_size=MAX_SIZE)
+    image2 = image2.resize(new_size)
+
+    log_text += f"Resized Image 1 size: {image1.size}\n"
+    log_text += f"Resized Image 2 size: {image2.size}\n"
+
+    dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(dev)
+
+    image1 = image1.convert("RGB")
+    image2 = image2.convert("RGB")
+
+    image1_original = image1.copy()
+    image2_original = image2.copy()
+
+    # convert PIL images to numpy arrays
+    image1_original = np.array(image1_original)
+    image2_original = np.array(image2_original)
+
+    # convert PIL images to tensors
+    image1 = torch.tensor(np.array(image1)).permute(2, 0, 1).float() / 255.0
+    image2 = torch.tensor(np.array(image2)).permute(2, 0, 1).float() / 255.0
+
+    image1 = image1.to(dev).unsqueeze(0)  # Add batch dimension
+    image2 = image2.to(dev).unsqueeze(0)  # Add batch dimension
+
+    kpts_1, desc_1, score_1 = model.detectAndCompute(image1, threshold=0.5, top_k=2048)
+    kpts_2, desc_2, score_2 = model.detectAndCompute(image2, threshold=0.5, top_k=2048)
+
+    log_text += f"Number of keypoints in image 1: {kpts_1.shape[0]}\n"
+    log_text += f"Number of keypoints in image 2: {kpts_2.shape[0]}\n"
+
+    matcher = KF.DescriptorMatcher("mnn")  # threshold is not used with mnn
+    match_dists, match_idxs = matcher(desc_1, desc_2)
+
+    log_text += f"Number of MNN matches: {match_idxs.shape[0]}\n"
+
+    cv2_matches = cv2_matches_from_kornia(match_dists, match_idxs)
+
+    do_ransac = match_idxs.shape[0] > 8
+
+    if do_ransac:
+        matched_pts_1 = kpts_1[match_idxs[:, 0]]
+        matched_pts_2 = kpts_2[match_idxs[:, 1]]
+
+        H, mask = KG.ransac.RANSAC(model_type="fundamental", inl_th=inl_th)(matched_pts_1, matched_pts_2)
+        matchesMask = mask.int().ravel().tolist()
+
+        log_text += f"RANSAC found {mask.sum().item()} inliers out of {mask.shape[0]} matches with an inlier threshold of {inl_th}.\n"
+    else:
+        log_text += "Not enough matches for RANSAC, skipping RANSAC step.\n"
+
+    kpts_1 = to_cv_kpts(kpts_1, score_1)
+    kpts_2 = to_cv_kpts(kpts_2, score_2)
+
+    keypoints_raw_1 = cv2.drawKeypoints(image1_original, kpts_1, image1_original, color=(0, 255, 0))
+    keypoints_raw_2 = cv2.drawKeypoints(image2_original, kpts_2, image2_original, color=(0, 255, 0))
+
+    # pad height smaller image to match the height of the larger image
+    if keypoints_raw_1.shape[0] < keypoints_raw_2.shape[0]:
+        pad_height = keypoints_raw_2.shape[0] - keypoints_raw_1.shape[0]
+        keypoints_raw_1 = np.pad(
+            keypoints_raw_1, ((0, pad_height), (0, 0), (0, 0)), mode="constant", constant_values=255
+        )
+    elif keypoints_raw_1.shape[0] > keypoints_raw_2.shape[0]:
+        pad_height = keypoints_raw_1.shape[0] - keypoints_raw_2.shape[0]
+        keypoints_raw_2 = np.pad(
+            keypoints_raw_2, ((0, pad_height), (0, 0), (0, 0)), mode="constant", constant_values=255
+        )
+
+    # concatenate keypoints images horizontally
+    keypoints_raw = np.concatenate((keypoints_raw_1, keypoints_raw_2), axis=1)
+    keypoints_raw_pil = Image.fromarray(keypoints_raw)
+
+    result_raw = cv2.drawMatches(
+        image1_original,
+        kpts_1,
+        image2_original,
+        kpts_2,
+        cv2_matches,
+        None,
+        matchColor=(0, 255, 0),
+        matchesMask=None,
+        # matchesMask=None,
+        flags=cv2.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS,
+    )
+
+    if not do_ransac:
+        result_ransac = None
+    else:
+        result_ransac = cv2.drawMatches(
+            image1_original,
+            kpts_1,
+            image2_original,
+            kpts_2,
+            cv2_matches,
+            None,
+            matchColor=(0, 255, 0),
+            matchesMask=matchesMask,
+            singlePointColor=(0, 0, 255),
+            flags=cv2.DrawMatchesFlags_DEFAULT,
+        )
+
+    # result = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)  # Convert BGR to RGB for display
+
+    # convert to PIL Image
+    result_raw_pil = Image.fromarray(result_raw)
+    if result_ransac is not None:
+        result_ransac_pil = Image.fromarray(result_ransac)
+    else:
+        result_ransac_pil = None
+
+    return log_text, result_ransac_pil, result_raw_pil, keypoints_raw_pil
+
+
+demo = gr.Interface(
+    fn=extract_keypoints,
+    inputs=[
+        gr.Image(type="pil", label="Image 1"),
+        gr.Image(type="pil", label="Image 2"),
+        gr.Slider(
+            minimum=0.1,
+            maximum=3.0,
+            step=0.1,
+            value=0.5,
+            label="RANSAC inlier threshold",
+            info="Threshold for RANSAC inlier detection. Lower values may yield fewer inliers but more robust matches.",
+        ),
+    ],
+    outputs=[
+        gr.Textbox(type="text", label="Log"),
+        gr.Image(type="pil", label="Keypoints and Matches (RANSAC)"),
+        gr.Image(type="pil", label="Keypoints and Matches"),
+        gr.Image(type="pil", label="Keypoint Detection Results"),
+    ],
+    title="RIPE: Reinforcement Learning on Unlabeled Image Pairs for Robust Keypoint Extraction",
+    description=description_text,
+    examples=[
+        [
+            "assets_gradio/all_souls_000013.jpg",
+            "assets_gradio/all_souls_000055.jpg",
+        ],
+        [
+            "assets_gradio/167170681_0e5c42fd21_o.jpg",
+            "assets_gradio/170804731_6bf4fbecd4_o.jpg",
+        ],
+        [
+            "assets_gradio/4171014767_0fe879b783_o.jpg",
+            "assets_gradio/4174108353_20422632d6_o.jpg",
+        ],
+    ],
+    flagging_mode="never",
+    theme="default",
+)
+demo.launch()

packages.txt

@@ -0,0 +1,4 @@
+libeigen3-dev
+cmake
+build-essentials
+python3-opencv

requirements.txt

@@ -0,0 +1,7 @@
+setuptools
+poselib @ git+https://github.com/PoseLib/PoseLib.git@56d158f744d3561b0b70174e6d8ca9a7fc9bd9c1
+opencv-python
+kornia
+numpy
+torch
+torchvision

ripe/__init__.py

@@ -0,0 +1 @@
+from .model_zoo import vgg_hyper  # noqa: F401

ripe/benchmarks/imw_2020.py

@@ -0,0 +1,320 @@
+import os
+from pathlib import Path
+
+import cv2
+import kornia.feature as KF
+import matplotlib.pyplot as plt
+import numpy as np
+import poselib
+import torch
+from tqdm import tqdm
+
+from ripe import utils
+from ripe.data.data_transforms import Compose, Normalize, Resize
+from ripe.data.datasets.disk_imw import DISK_IMW
+from ripe.utils.pose_error import AUCMetric, relative_pose_error
+from ripe.utils.utils import (
+    cv2_matches_from_kornia,
+    cv_resize_and_pad_to_shape,
+    to_cv_kpts,
+)
+
+log = utils.get_pylogger(__name__)
+
+
+class IMW_2020_Benchmark:
+    def __init__(
+        self,
+        use_predefined_subset: bool = True,
+        conf_inference=None,
+        edge_input_divisible_by=None,
+    ):
+        data_dir = os.getenv("DATA_DIR")
+        if data_dir is None:
+            raise ValueError("Environment variable DATA_DIR is not set.")
+        root_path = Path(data_dir) / "disk-data"
+
+        self.data = DISK_IMW(
+            str(
+                root_path
+            ),  # Resize only to ensure that the input size is divisible the value of edge_input_divisible_by
+            transforms=Compose(
+                [
+                    Resize(None, edge_input_divisible_by),
+                    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+                ]
+            ),
+        )
+        self.ids_subset = None
+        self.results = []
+        self.conf_inference = conf_inference
+
+        # fmt: off
+        if use_predefined_subset:
+            self.ids_subset = [4921, 3561, 3143, 6040, 802, 6828, 5338, 9275, 10764, 10085, 5124, 11355, 7, 10027, 2161, 4433, 6887, 3311, 10766,
+                               11451, 11433, 8539, 2581, 10300, 10562, 1723, 8803, 6275, 10140, 11487, 6238, 638, 8092, 9979, 201, 10394, 3414,
+                               9002, 7456, 2431, 632, 6589, 9265, 9889, 3139, 7890, 10619, 4899, 675, 176, 4309, 4814, 3833, 3519, 148, 4560, 10705,
+                               3744, 1441, 4049, 1791, 5106, 575, 1540, 1105, 6791, 1383, 9344, 501, 2504, 4335, 8992, 10970, 10786, 10405, 9317,
+                               5279, 1396, 5044, 9408, 11125, 10417, 7627, 7480, 1358, 7738, 5461, 10178, 9226, 8106, 2766, 6216, 4032, 7298, 259,
+                               3021, 2645, 8756, 7513, 3163, 2510, 6701, 6684, 3159, 9689, 7425, 6066, 1904, 6382, 3052, 777, 6277, 7409, 5997, 2987,
+                               11316, 2894, 4528, 1927, 10366, 8605, 2726, 1886, 2416, 2164, 3352, 2997, 6636, 6765, 5609, 3679, 76, 10956, 3612, 6699,
+                               1741, 8811, 3755, 1285, 9520, 2476, 3977, 370, 9823, 1834, 7551, 6227, 7303, 6399, 4758, 10713, 5050, 380, 11056, 7620,
+                               4826, 6090, 9011, 7523, 7355, 8021, 9801, 1801, 6522, 7138, 10017, 8732, 6402, 3116, 4031, 6088, 3975, 9841, 9082, 9412,
+                               5406, 217, 2385, 8791, 8361, 494, 4319, 5275, 3274, 335, 6731, 207, 10095, 3068, 5996, 3951, 2808, 5877, 6134, 7772, 10042,
+                               8574, 5501, 10885, 7871]
+            # self.ids_subset = self.ids_subset[:10]
+        # fmt: on
+
+    def evaluate_sample(self, model, sample, dev):
+        img_1 = sample["src_image"].unsqueeze(0).to(dev)
+        img_2 = sample["trg_image"].unsqueeze(0).to(dev)
+
+        scale_h_1, scale_w_1 = (
+            sample["orig_size_src"][0] / img_1.shape[2],
+            sample["orig_size_src"][1] / img_1.shape[3],
+        )
+        scale_h_2, scale_w_2 = (
+            sample["orig_size_trg"][0] / img_2.shape[2],
+            sample["orig_size_trg"][1] / img_2.shape[3],
+        )
+
+        M = None
+        info = {}
+        kpts_1, desc_1, score_1 = None, None, None
+        kpts_2, desc_2, score_2 = None, None, None
+        match_dists, match_idxs = None, None
+
+        try:
+            kpts_1, desc_1, score_1 = model.detectAndCompute(img_1, **self.conf_inference)
+            kpts_2, desc_2, score_2 = model.detectAndCompute(img_2, **self.conf_inference)
+
+            if kpts_1.dim() == 3:
+                assert kpts_1.shape[0] == 1 and kpts_2.shape[0] == 1, "Batch size must be 1"
+
+                kpts_1, desc_1, score_1 = (
+                    kpts_1.squeeze(0),
+                    desc_1[0].squeeze(0),
+                    score_1[0].squeeze(0),
+                )
+                kpts_2, desc_2, score_2 = (
+                    kpts_2.squeeze(0),
+                    desc_2[0].squeeze(0),
+                    score_2[0].squeeze(0),
+                )
+
+            scale_1 = torch.tensor([scale_w_1, scale_h_1], dtype=torch.float).to(dev)
+            scale_2 = torch.tensor([scale_w_2, scale_h_2], dtype=torch.float).to(dev)
+
+            kpts_1 = kpts_1 * scale_1
+            kpts_2 = kpts_2 * scale_2
+
+            matcher = KF.DescriptorMatcher("mnn")  # threshold is not used with mnn
+            match_dists, match_idxs = matcher(desc_1, desc_2)
+
+            matched_pts_1 = kpts_1[match_idxs[:, 0]]
+            matched_pts_2 = kpts_2[match_idxs[:, 1]]
+
+            camera_1 = sample["src_camera"]
+            camera_2 = sample["trg_camera"]
+
+            M, info = poselib.estimate_relative_pose(
+                matched_pts_1.cpu().numpy(),
+                matched_pts_2.cpu().numpy(),
+                camera_1.to_cameradict(),
+                camera_2.to_cameradict(),
+                {
+                    "max_epipolar_error": 0.5,
+                },
+                {},
+            )
+        except RuntimeError as e:
+            if "No keypoints detected" in str(e):
+                pass
+            else:
+                raise e
+
+        success = M is not None
+        if success:
+            M = {
+                "R": torch.tensor(M.R, dtype=torch.float),
+                "t": torch.tensor(M.t, dtype=torch.float),
+            }
+            inl = info["inliers"]
+        else:
+            M = {
+                "R": torch.eye(3, dtype=torch.float),
+                "t": torch.zeros((3), dtype=torch.float),
+            }
+            inl = np.zeros((0,)).astype(bool)
+
+        t_err, r_err = relative_pose_error(sample["s2t_R"].cpu(), sample["s2t_T"].cpu(), M["R"], M["t"])
+
+        rel_pose_error = max(t_err.item(), r_err.item()) if success else np.inf
+        ransac_inl = np.sum(inl)
+        ransac_inl_ratio = np.mean(inl)
+
+        if success:
+            assert match_dists is not None and match_idxs is not None, "Matches must be computed"
+            cv_keypoints_src = to_cv_kpts(kpts_1, score_1)
+            cv_keypoints_trg = to_cv_kpts(kpts_2, score_2)
+            cv_matches = cv2_matches_from_kornia(match_dists, match_idxs)
+            cv_mask = [int(m) for m in inl]
+        else:
+            cv_keypoints_src, cv_keypoints_trg = [], []
+            cv_matches, cv_mask = [], []
+
+        estimation = {
+            "success": success,
+            "M_0to1": M,
+            "inliers": torch.tensor(inl).to(img_1),
+            "rel_pose_error": rel_pose_error,
+            "ransac_inl": ransac_inl,
+            "ransac_inl_ratio": ransac_inl_ratio,
+            "path_src_image": sample["src_path"],
+            "path_trg_image": sample["trg_path"],
+            "cv_keypoints_src": cv_keypoints_src,
+            "cv_keypoints_trg": cv_keypoints_trg,
+            "cv_matches": cv_matches,
+            "cv_mask": cv_mask,
+        }
+
+        return estimation
+
+    def evaluate(self, model, dev, progress_bar=False):
+        model.eval()
+
+        # reset results
+        self.results = []
+
+        for idx in tqdm(
+            self.ids_subset if self.ids_subset is not None else range(len(self.data)),
+            disable=not progress_bar,
+        ):
+            sample = self.data[idx]
+            self.results.append(self.evaluate_sample(model, sample, dev))
+
+    def get_auc(self, threshold=5, downsampled=False):
+        if len(self.results) == 0:
+            raise ValueError("No results to log. Run evaluate first.")
+
+        summary_results = self.calc_auc(downsampled=downsampled)
+
+        return summary_results[f"rel_pose_error@{threshold}°{'__original' if not downsampled else '__downsampled'}"]
+
+    def plot_results(self, num_samples=10, logger=None, step=None, downsampled=False):
+        if len(self.results) == 0:
+            raise ValueError("No results to plot. Run evaluate first.")
+
+        plot_data = []
+
+        for result in self.results[:num_samples]:
+            img1 = cv2.imread(result["path_src_image"])
+            img2 = cv2.imread(result["path_trg_image"])
+
+            # from BGR to RGB
+            img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)
+            img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2RGB)
+
+            plt_matches = cv2.drawMatches(
+                img1,
+                result["cv_keypoints_src"],
+                img2,
+                result["cv_keypoints_trg"],
+                result["cv_matches"],
+                None,
+                matchColor=None,
+                matchesMask=result["cv_mask"],
+                flags=cv2.DrawMatchesFlags_DEFAULT,
+            )
+            file_name = (
+                Path(result["path_src_image"]).parent.parent.name
+                + "_"
+                + Path(result["path_src_image"]).stem
+                + Path(result["path_trg_image"]).stem
+                + ("_downsampled" if downsampled else "")
+                + ".png"
+            )
+            # print rel_pose_error on image
+            plt_matches = cv2.putText(
+                plt_matches,
+                f"rel_pose_error: {result['rel_pose_error']:.2f} num_inliers: {result['ransac_inl']} inl_ratio: {result['ransac_inl_ratio']:.2f} num_matches: {len(result['cv_matches'])} num_keypoints: {len(result['cv_keypoints_src'])}/{len(result['cv_keypoints_trg'])}",
+                (10, 30),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                1,
+                (0, 0, 0),
+                2,
+                cv2.LINE_8,
+            )
+
+            plot_data.append({"file_name": file_name, "image": plt_matches})
+
+        if logger is None:
+            log.info("No logger provided. Using plt to plot results.")
+            for image in plot_data:
+                plt.imsave(
+                    image["file_name"],
+                    cv_resize_and_pad_to_shape(image["image"], (1024, 2048)),
+                )
+                plt.close()
+        else:
+            import wandb
+
+            log.info(f"Logging images to wandb with step={step}")
+            if not downsampled:
+                logger.log(
+                    {
+                        "examples": [
+                            wandb.Image(cv_resize_and_pad_to_shape(image["image"], (1024, 2048))) for image in plot_data
+                        ]
+                    },
+                    step=step,
+                )
+            else:
+                logger.log(
+                    {
+                        "examples_downsampled": [
+                            wandb.Image(cv_resize_and_pad_to_shape(image["image"], (1024, 2048))) for image in plot_data
+                        ]
+                    },
+                    step=step,
+                )
+
+    def log_results(self, logger=None, step=None, downsampled=False):
+        if len(self.results) == 0:
+            raise ValueError("No results to log. Run evaluate first.")
+
+        summary_results = self.calc_auc(downsampled=downsampled)
+
+        if logger is not None:
+            logger.log(summary_results, step=step)
+        else:
+            log.warning("No logger provided. Printing results instead.")
+            print(self.calc_auc())
+
+    def print_results(self):
+        if len(self.results) == 0:
+            raise ValueError("No results to print. Run evaluate first.")
+
+        print(self.calc_auc())
+
+    def calc_auc(self, auc_thresholds=None, downsampled=False):
+        if auc_thresholds is None:
+            auc_thresholds = [5, 10, 20]
+        if not isinstance(auc_thresholds, list):
+            auc_thresholds = [auc_thresholds]
+
+        if len(self.results) == 0:
+            raise ValueError("No results to calculate auc. Run evaluate first.")
+
+        rel_pose_errors = [r["rel_pose_error"] for r in self.results]
+
+        pose_aucs = AUCMetric(auc_thresholds, rel_pose_errors).compute()
+        assert isinstance(pose_aucs, list) and len(pose_aucs) == len(auc_thresholds)
+
+        ext = "_downsampled" if downsampled else "_original"
+
+        summary = {}
+        for i, ath in enumerate(auc_thresholds):
+            summary[f"rel_pose_error@{ath}°_{ext}"] = pose_aucs[i]
+
+        return summary

ripe/data/data_transforms.py

@@ -0,0 +1,204 @@
+import collections
+import collections.abc
+
+import kornia.geometry as KG
+import numpy as np
+import torch
+from torchvision.transforms import functional as TF
+
+
+class Compose:
+    """Composes several transforms together. The transforms are applied in the order they are passed in.
+    Args:        transforms (list): A list of transforms to be applied.
+    """
+
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, src, trg, src_mask, trg_mask, h):
+        for t in self.transforms:
+            src, trg, src_mask, trg_mask, h = t(src, trg, src_mask, trg_mask, h)
+
+        return src, trg, src_mask, trg_mask, h
+
+
+class Transform:
+    """Base class for all transforms. It provides a method to apply a transformation function to the input images and masks.
+    Args:
+        src (torch.Tensor): The source image tensor.
+        trg (torch.Tensor): The target image tensor.
+        src_mask (torch.Tensor): The source image mask tensor.
+        trg_mask (torch.Tensor): The target image mask tensor.
+        h (torch.Tensor): The homography matrix tensor.
+    Returns:
+        tuple: A tuple containing the transformed source image, the transformed target image, the transformed source mask,
+        the transformed target mask and the updated homography matrix.
+    """
+
+    def __init__(self):
+        pass
+
+    def apply_transform(self, src, trg, src_mask, trg_mask, h, transfrom_function):
+        src, trg, src_mask, trg_mask, h = transfrom_function(src, trg, src_mask, trg_mask, h)
+        return src, trg, src_mask, trg_mask, h
+
+
+class Normalize(Transform):
+    def __init__(self, mean, std):
+        self.mean = mean
+        self.std = std
+
+    def __call__(self, src, trg, src_mask, trg_mask, h):
+        return self.apply_transform(src, trg, src_mask, trg_mask, h, self.transform_function)
+
+    def transform_function(self, src, trg, src_mask, trg_mask, h):
+        src = TF.normalize(src, mean=self.mean, std=self.std)
+        trg = TF.normalize(trg, mean=self.mean, std=self.std)
+        return src, trg, src_mask, trg_mask, h
+
+
+class ResizeAndPadWithHomography(Transform):
+    def __init__(self, target_size_longer_side=768):
+        self.target_size = target_size_longer_side
+
+    def __call__(self, src, trg, src_mask, trg_mask, h):
+        return self.apply_transform(src, trg, src_mask, trg_mask, h, self.transform_function)
+
+    def transform_function(self, src, trg, src_mask, trg_mask, h):
+        src_w, src_h = src.shape[-1], src.shape[-2]
+        trg_w, trg_h = trg.shape[-1], trg.shape[-2]
+
+        # Resizing logic for both images
+        scale_src, new_src_w, new_src_h = self.compute_resize(src_w, src_h)
+        scale_trg, new_trg_w, new_trg_h = self.compute_resize(trg_w, trg_h)
+
+        # Resize both images
+        src_resized = TF.resize(src, [new_src_h, new_src_w])
+        trg_resized = TF.resize(trg, [new_trg_h, new_trg_w])
+
+        src_mask_resized = TF.resize(src_mask, [new_src_h, new_src_w])
+        trg_mask_resized = TF.resize(trg_mask, [new_trg_h, new_trg_w])
+
+        # Pad the resized images to be square (768x768)
+        src_padded, src_padding = self.apply_padding(src_resized, new_src_w, new_src_h)
+        trg_padded, trg_padding = self.apply_padding(trg_resized, new_trg_w, new_trg_h)
+
+        src_mask_padded, _ = self.apply_padding(src_mask_resized, new_src_w, new_src_h)
+        trg_mask_padded, _ = self.apply_padding(trg_mask_resized, new_trg_w, new_trg_h)
+
+        # Update the homography matrix
+        h = self.update_homography(h, scale_src, src_padding, scale_trg, trg_padding)
+
+        return src_padded, trg_padded, src_mask_padded, trg_mask_padded, h
+
+    def compute_resize(self, w, h):
+        if w > h:
+            scale = self.target_size / w
+            new_w = self.target_size
+            new_h = int(h * scale)
+        else:
+            scale = self.target_size / h
+            new_h = self.target_size
+            new_w = int(w * scale)
+        return scale, new_w, new_h
+
+    def apply_padding(self, img, new_w, new_h):
+        pad_w = (self.target_size - new_w) // 2
+        pad_h = (self.target_size - new_h) // 2
+        padding = [
+            pad_w,
+            pad_h,
+            self.target_size - new_w - pad_w,
+            self.target_size - new_h - pad_h,
+        ]
+        img_padded = TF.pad(img, padding, fill=0)  # Zero-pad
+        return img_padded, padding
+
+    def update_homography(self, h, scale_src, padding_src, scale_trg, padding_trg):
+        # Create the scaling matrices
+        scale_matrix_src = np.array([[scale_src, 0, 0], [0, scale_src, 0], [0, 0, 1]])
+        scale_matrix_trg = np.array([[scale_trg, 0, 0], [0, scale_trg, 0], [0, 0, 1]])
+
+        # Create the padding translation matrices
+        pad_matrix_src = np.array([[1, 0, padding_src[0]], [0, 1, padding_src[1]], [0, 0, 1]])
+        pad_matrix_trg = np.array([[1, 0, -padding_trg[0]], [0, 1, -padding_trg[1]], [0, 0, 1]])
+
+        # Update the homography: apply scaling and translation
+        h_updated = (
+            pad_matrix_trg
+            @ scale_matrix_trg
+            @ h.numpy()
+            @ np.linalg.inv(scale_matrix_src)
+            @ np.linalg.inv(pad_matrix_src)
+        )
+
+        return torch.from_numpy(h_updated).float()
+
+
+class Resize(Transform):
+    def __init__(self, output_size, edge_divisible_by=None, side="long", antialias=True):
+        self.output_size = output_size
+        self.edge_divisible_by = edge_divisible_by
+        self.side = side
+        self.antialias = antialias
+
+    def __call__(self, src, trg, src_mask, trg_mask, h):
+        return self.apply_transform(src, trg, src_mask, trg_mask, h, self.transform_function)
+
+    def transform_function(self, src, trg, src_mask, trg_mask, h):
+        new_size_src = self.get_new_image_size(src)
+        new_size_trg = self.get_new_image_size(trg)
+
+        src, T_src = self.resize(src, new_size_src)
+        trg, T_trg = self.resize(trg, new_size_trg)
+
+        src_mask, _ = self.resize(src_mask, new_size_src)
+        trg_mask, _ = self.resize(trg_mask, new_size_trg)
+
+        h = torch.from_numpy(T_trg @ h.numpy() @ T_src).float()
+
+        return src, trg, src_mask, trg_mask, h
+
+    def resize(self, img, size):
+        h, w = img.shape[-2:]
+
+        img = KG.transform.resize(
+            img,
+            size,
+            side=self.side,
+            antialias=self.antialias,
+            align_corners=None,
+            interpolation="bilinear",
+        )
+
+        scale = torch.Tensor([img.shape[-1] / w, img.shape[-2] / h]).to(img)
+        T = np.diag([scale[0].item(), scale[1].item(), 1])
+
+        return img, T
+
+    def get_new_image_size(self, img):
+        h, w = img.shape[-2:]
+
+        if isinstance(self.output_size, collections.abc.Iterable):
+            assert len(self.output_size) == 2
+            return tuple(self.output_size)
+        if self.output_size is None:  # keep the original size, but possibly make it divisible by edge_divisible_by
+            size = (h, w)
+        else:
+            side_size = self.output_size
+            aspect_ratio = w / h
+            if self.side not in ("short", "long", "vert", "horz"):
+                raise ValueError(f"side can be one of 'short', 'long', 'vert', and 'horz'. Got '{self.side}'")
+            if self.side == "vert":
+                size = side_size, int(side_size * aspect_ratio)
+            elif self.side == "horz":
+                size = int(side_size / aspect_ratio), side_size
+            elif (self.side == "short") ^ (aspect_ratio < 1.0):
+                size = side_size, int(side_size * aspect_ratio)
+            else:
+                size = int(side_size / aspect_ratio), side_size
+
+        if self.edge_divisible_by is not None:
+            df = self.edge_divisible_by
+            size = list(map(lambda x: int(x // df * df), size))
+        return size

ripe/data/datasets/acdc.py

@@ -0,0 +1,154 @@
+from pathlib import Path
+from typing import Any, Callable, Dict, Optional
+
+import torch
+from torch.utils.data import Dataset
+from torchvision.io import read_image
+
+from ripe import utils
+from ripe.data.data_transforms import Compose
+from ripe.utils.utils import get_other_random_id
+
+log = utils.get_pylogger(__name__)
+
+
+class ACDC(Dataset):
+    def __init__(
+        self,
+        root: Path,
+        stage: str = "train",
+        condition: str = "rain",
+        transforms: Optional[Callable] = None,
+        positive_only: bool = False,
+    ) -> None:
+        self.root = root
+        self.stage = stage
+        self.condition = condition
+        self.transforms = transforms
+        self.positive_only = positive_only
+
+        if isinstance(self.root, str):
+            self.root = Path(self.root)
+
+        if not self.root.exists():
+            raise FileNotFoundError(f"Dataset not found at {self.root}")
+
+        if transforms is None:
+            self.transforms = Compose([])
+        else:
+            self.transforms = transforms
+
+        if self.stage not in ["train", "val", "test", "pred"]:
+            raise RuntimeError(
+                "Unknown option "
+                + self.stage
+                + " as training stage variable. Valid options: 'train', 'val', 'test' and 'pred'"
+            )
+
+        if self.stage == "pred":  # prediction uses the test set
+            self.stage = "test"
+
+        if self.stage in ["val", "test", "pred"]:
+            self.positive_only = True
+            log.info(f"{self.stage} stage: Using only positive pairs!")
+
+        weather_conditions = ["fog", "night", "rain", "snow"]
+
+        if self.condition not in weather_conditions + ["all"]:
+            raise RuntimeError(
+                "Unknown option "
+                + self.condition
+                + " as weather condition variable. Valid options: 'fog', 'night', 'rain', 'snow' and 'all'"
+            )
+
+        self.weather_condition_query = weather_conditions if self.condition == "all" else [self.condition]
+
+        self._read_sample_files()
+
+        if positive_only:
+            log.warning("Using only positive pairs!")
+        log.info(f"Found {len(self.src_images)} source images and {len(self.trg_images)} target images.")
+
+    def _read_sample_files(self):
+        file_name_pattern_ref = "_ref_anon.png"
+        file_name_pattern = "_rgb_anon.png"
+
+        self.trg_images = []
+        self.src_images = []
+
+        for weather_condition in self.weather_condition_query:
+            rgb_files = sorted(
+                list(self.root.glob("rgb_anon/" + weather_condition + "/" + self.stage + "/**/*" + file_name_pattern)),
+                key=lambda i: i.stem[:21],
+            )
+
+            src_images = sorted(
+                list(
+                    self.root.glob(
+                        "rgb_anon/" + weather_condition + "/" + self.stage + "_ref" + "/**/*" + file_name_pattern_ref
+                    )
+                ),
+                key=lambda i: i.stem[:21],
+            )
+
+            self.trg_images += rgb_files
+            self.src_images += src_images
+
+    def __len__(self) -> int:
+        if self.positive_only:
+            return len(self.trg_images)
+        return 2 * len(self.trg_images)
+
+    def __getitem__(self, idx: int) -> Dict[str, Any]:
+        sample: Any = {}
+
+        positive_sample = (idx % 2 == 0) or (self.positive_only)
+        if not self.positive_only:
+            idx = idx // 2
+
+        sample["label"] = positive_sample
+
+        if positive_sample:
+            sample["src_path"] = str(self.src_images[idx])
+            sample["trg_path"] = str(self.trg_images[idx])
+
+            assert self.src_images[idx].stem[:21] == self.trg_images[idx].stem[:21], (
+                f"Source and target image mismatch: {self.src_images[idx]} vs {self.trg_images[idx]}"
+            )
+
+            src_img = read_image(sample["src_path"])
+            trg_img = read_image(sample["trg_path"])
+
+            homography = torch.eye(3, dtype=torch.float32)
+        else:
+            sample["src_path"] = str(self.src_images[idx])
+            idx_other = get_other_random_id(idx, len(self) // 2)
+            sample["trg_path"] = str(self.trg_images[idx_other])
+
+            assert self.src_images[idx].stem[:21] != self.trg_images[idx_other].stem[:21], (
+                f"Source and target image match for negative sample: {self.src_images[idx]} vs {self.trg_images[idx_other]}"
+            )
+
+            src_img = read_image(sample["src_path"])
+            trg_img = read_image(sample["trg_path"])
+
+            homography = torch.zeros((3, 3), dtype=torch.float32)
+
+        src_img = src_img / 255.0
+        trg_img = trg_img / 255.0
+
+        _, H, W = src_img.shape
+
+        src_mask = torch.ones((1, H, W), dtype=torch.uint8)
+        trg_mask = torch.ones((1, H, W), dtype=torch.uint8)
+
+        if self.transforms:
+            src_img, trg_img, src_mask, trg_mask, _ = self.transforms(src_img, trg_img, src_mask, trg_mask, homography)
+
+        sample["src_image"] = src_img
+        sample["trg_image"] = trg_img
+        sample["src_mask"] = src_mask.to(torch.bool)
+        sample["trg_mask"] = trg_mask.to(torch.bool)
+        sample["homography"] = homography
+
+        return sample

ripe/data/datasets/dataset_combinator.py

@@ -0,0 +1,88 @@
+import torch
+
+from ripe import utils
+
+log = utils.get_pylogger(__name__)
+
+
+class DatasetCombinator:
+    """Combines multiple datasets into one. Length of the combined dataset is the length of the
+    longest dataset. Shorter datasets are looped over.
+
+    Args:
+        datasets: List of datasets to combine.
+        mode: How to sample from the datasets. Can be either "uniform" or "weighted".
+            In "uniform" mode, each dataset is sampled with equal probability.
+            In "weighted" mode, each dataset is sampled with probability proportional to its length.
+    """
+
+    def __init__(self, datasets, mode="uniform", weights=None):
+        self.datasets = datasets
+
+        names_datasets = [type(ds).__name__ for ds in self.datasets]
+        self.lengths = [len(ds) for ds in datasets]
+
+        if mode == "weighted":
+            self.probs_datasets = [length / sum(self.lengths) for length in self.lengths]
+        elif mode == "uniform":
+            self.probs_datasets = [1 / len(self.datasets) for _ in self.datasets]
+        elif mode == "custom":
+            assert weights is not None, "Weights must be provided in custom mode"
+            assert len(weights) == len(datasets), "Number of weights must match number of datasets"
+            assert sum(weights) == 1.0, "Weights must sum to 1"
+            self.probs_datasets = weights
+        else:
+            raise ValueError(f"Unknown mode {mode}")
+
+        log.info("Got the following datasets: ")
+
+        for name, length, prob in zip(names_datasets, self.lengths, self.probs_datasets):
+            log.info(f"{name} with {length} samples and probability {prob}")
+        log.info(f"Total number of samples: {sum(self.lengths)}")
+
+        self.num_samples = max(self.lengths)
+
+        self.dataset_dist = torch.distributions.Categorical(probs=torch.tensor(self.probs_datasets))
+
+    def __len__(self):
+        return self.num_samples
+
+    def __getitem__(self, idx: int):
+        positive_sample = idx % 2 == 0
+
+        if positive_sample:
+            dataset_idx = self.dataset_dist.sample().item()
+
+            idx = torch.randint(0, self.lengths[dataset_idx], (1,)).item()
+            while idx % 2 == 1:
+                idx = torch.randint(0, self.lengths[dataset_idx], (1,)).item()
+
+            return self.datasets[dataset_idx][idx]
+        else:
+            dataset_idx_1 = self.dataset_dist.sample().item()
+            dataset_idx_2 = self.dataset_dist.sample().item()
+
+            if dataset_idx_1 == dataset_idx_2:
+                idx = torch.randint(0, self.lengths[dataset_idx_1], (1,)).item()
+                while idx % 2 == 0:
+                    idx = torch.randint(0, self.lengths[dataset_idx_1], (1,)).item()
+                return self.datasets[dataset_idx_1][idx]
+
+            else:
+                idx_1 = torch.randint(0, self.lengths[dataset_idx_1], (1,)).item()
+                idx_2 = torch.randint(0, self.lengths[dataset_idx_2], (1,)).item()
+
+                sample_1 = self.datasets[dataset_idx_1][idx_1]
+                sample_2 = self.datasets[dataset_idx_2][idx_2]
+
+                sample = {
+                    "label": False,
+                    "src_path": sample_1["src_path"],
+                    "trg_path": sample_2["trg_path"],
+                    "src_image": sample_1["src_image"],
+                    "trg_image": sample_2["trg_image"],
+                    "src_mask": sample_1["src_mask"],
+                    "trg_mask": sample_2["trg_mask"],
+                    "homography": sample_2["homography"],
+                }
+                return sample

ripe/data/datasets/disk_imw.py

@@ -0,0 +1,160 @@
+import json
+import random
+from itertools import accumulate
+from pathlib import Path
+from typing import Any, Callable, Dict, Optional, Tuple
+
+import torch
+from torch.utils.data import Dataset
+from torchvision.io import read_image
+
+from ripe import utils
+from ripe.data.data_transforms import Compose
+from ripe.utils.image_utils import Camera, cameras2F
+
+log = utils.get_pylogger(__name__)
+
+
+class DISK_IMW(Dataset):
+    def __init__(
+        self,
+        root: str,
+        stage: str = "val",
+        # condition: str = "rain",
+        transforms: Optional[Callable] = None,
+    ) -> None:
+        self.root = root
+        self.stage = stage
+        self.transforms = transforms
+
+        if isinstance(self.root, str):
+            self.root = Path(self.root)
+
+        if not self.root.exists():
+            raise FileNotFoundError(f"Dataset not found at {self.root}")
+
+        if transforms is None:
+            self.transforms = Compose([])
+        else:
+            self.transforms = transforms
+
+        if self.stage not in ["val"]:
+            raise RuntimeError("Unknown option " + self.stage + " as training stage variable. Valid options: 'train'")
+
+        json_path = self.root / "imw2020-val" / "dataset.json"
+        with open(json_path) as json_file:
+            json_data = json.load(json_file)
+
+        self.scenes = []
+
+        for scene in json_data:
+            self.scenes.append(Scene(self.root / "imw2020-val", json_data[scene]))
+
+        self.tuples_per_scene = [len(scene) for scene in self.scenes]
+
+    def __len__(self) -> int:
+        return sum(self.tuples_per_scene)
+
+    def __getitem__(self, idx: int) -> Dict[str, Any]:
+        sample: Any = {}
+
+        i_scene, i_image = self._get_scene_and_image_id_from_idx(idx)
+
+        sample["src_path"], sample["trg_path"], path_calib_src, path_calib_trg = self.scenes[i_scene][i_image]
+
+        cam_src = Camera.from_calibration_file(path_calib_src)
+        cam_trg = Camera.from_calibration_file(path_calib_trg)
+
+        F = self.get_F(cam_src, cam_trg)
+        s2t_R, s2t_T = self.get_relative_pose(cam_src, cam_trg)
+
+        src_img = read_image(sample["src_path"]) / 255.0
+        trg_img = read_image(sample["trg_path"]) / 255.0
+
+        _, H_src, W_src = src_img.shape
+        _, H_trg, W_trg = trg_img.shape
+
+        src_mask = torch.ones((1, H_src, W_src), dtype=torch.uint8)
+        trg_mask = torch.ones((1, H_trg, W_trg), dtype=torch.uint8)
+
+        H = torch.eye(3)
+        if self.transforms:
+            src_img, trg_img, src_mask, trg_mask, _ = self.transforms(src_img, trg_img, src_mask, trg_mask, H)
+
+        # check if transformations in self.transforms. Only Normalize is allowed
+        for t in self.transforms.transforms:
+            if t.__class__.__name__ not in ["Normalize", "Resize"]:
+                raise ValueError(f"Transform {t.__class__.__name__} not allowed in DISK_IMW dataset")
+
+        sample["src_image"] = src_img
+        sample["trg_image"] = trg_img
+        sample["orig_size_src"] = (H_src, W_src)
+        sample["orig_size_trg"] = (H_trg, W_trg)
+        sample["src_mask"] = src_mask.to(torch.bool)
+        sample["trg_mask"] = trg_mask.to(torch.bool)
+        sample["F"] = F
+        sample["s2t_R"] = s2t_R
+        sample["s2t_T"] = s2t_T
+        sample["src_camera"] = cam_src
+        sample["trg_camera"] = cam_trg
+
+        return sample
+
+    def get_relative_pose(self, cam_src: Camera, cam_trg: Camera) -> Tuple[torch.Tensor, torch.Tensor]:
+        R = cam_trg.R @ cam_src.R.T
+        T = cam_trg.t - R @ cam_src.t
+
+        return R, T
+
+    def get_F(self, cam_src: Camera, cam_trg: Camera) -> torch.Tensor:
+        F = cameras2F(cam_src, cam_trg)
+
+        return F
+
+    def _get_scene_and_image_id_from_idx(self, idx: int) -> Tuple[int, int]:
+        accumulated_tuples = accumulate(self.tuples_per_scene)
+
+        if idx >= sum(self.tuples_per_scene):
+            raise IndexError(f"Index {idx} out of bounds")
+
+        idx_scene = None
+        for i, accumulated_tuple in enumerate(accumulated_tuples):
+            idx_scene = i
+            if idx < accumulated_tuple:
+                break
+
+        idx_image = idx - sum(self.tuples_per_scene[:idx_scene])
+
+        return idx_scene, idx_image
+
+    def _get_other_random_scene_and_image_id(self, scene_id_to_exclude: int) -> Tuple[int, int]:
+        possible_scene_ids = list(range(len(self.scenes)))
+        possible_scene_ids.remove(scene_id_to_exclude)
+
+        idx_scene = random.choice(possible_scene_ids)
+        idx_image = random.randint(0, len(self.scenes[idx_scene]) - 1)
+
+        return idx_scene, idx_image
+
+
+class Scene:
+    def __init__(self, root_path, scene_data: Dict[str, Any]) -> None:
+        self.root_path = root_path
+        self.image_path = Path(scene_data["image_path"])
+        self.calib_path = Path(scene_data["calib_path"])
+        self.image_names = scene_data["images"]
+        self.tuples = scene_data["tuples"]
+
+    def __len__(self) -> int:
+        return len(self.tuples)
+
+    def __getitem__(self, idx: int) -> Dict[str, Any]:
+        idx_1 = self.tuples[idx][0]
+        idx_2 = self.tuples[idx][1]
+
+        path_image_1 = str(self.root_path / self.image_path / self.image_names[idx_1]) + ".jpg"
+        path_image_2 = str(self.root_path / self.image_path / self.image_names[idx_2]) + ".jpg"
+        path_calib_1 = str(self.root_path / self.calib_path / ("calibration_" + self.image_names[idx_1])) + ".h5"
+        path_calib_2 = str(self.root_path / self.calib_path / ("calibration_" + self.image_names[idx_2])) + ".h5"
+
+        return path_image_1, path_image_2, path_calib_1, path_calib_2

ripe/data/datasets/disk_megadepth.py

@@ -0,0 +1,157 @@
+import json
+import random
+from itertools import accumulate
+from pathlib import Path
+from typing import Any, Callable, Dict, Optional, Tuple
+
+import torch
+from torch.utils.data import Dataset
+from torchvision.io import read_image
+
+from ripe import utils
+from ripe.data.data_transforms import Compose
+
+log = utils.get_pylogger(__name__)
+
+
+class DISK_Megadepth(Dataset):
+    def __init__(
+        self,
+        root: str,
+        max_scene_size: int,
+        stage: str = "train",
+        # condition: str = "rain",
+        transforms: Optional[Callable] = None,
+        positive_only: bool = False,
+    ) -> None:
+        self.root = root
+        self.stage = stage
+        self.transforms = transforms
+        self.positive_only = positive_only
+
+        if isinstance(self.root, str):
+            self.root = Path(self.root)
+
+        if not self.root.exists():
+            raise FileNotFoundError(f"Dataset not found at {self.root}")
+
+        if transforms is None:
+            self.transforms = Compose([])
+        else:
+            self.transforms = transforms
+
+        if self.stage not in ["train"]:
+            raise RuntimeError("Unknown option " + self.stage + " as training stage variable. Valid options: 'train'")
+
+        json_path = self.root / "megadepth" / "dataset.json"
+        with open(json_path) as json_file:
+            json_data = json.load(json_file)
+
+        self.scenes = []
+
+        for scene in json_data:
+            self.scenes.append(Scene(self.root / "megadepth", json_data[scene], max_scene_size))
+
+        self.tuples_per_scene = [len(scene) for scene in self.scenes]
+
+        if positive_only:
+            log.warning("Using only positive pairs!")
+
+    def __len__(self) -> int:
+        if self.positive_only:
+            return sum(self.tuples_per_scene)
+        return 2 * sum(self.tuples_per_scene)
+
+    def __getitem__(self, idx: int) -> Dict[str, Any]:
+        sample: Any = {}
+
+        positive_sample = idx % 2 == 0 or self.positive_only
+        if not self.positive_only:
+            idx = idx // 2
+
+        sample["label"] = positive_sample
+
+        i_scene, i_image = self._get_scene_and_image_id_from_idx(idx)
+
+        if positive_sample:
+            sample["src_path"], sample["trg_path"] = self.scenes[i_scene][i_image]
+
+            homography = torch.eye(3, dtype=torch.float32)
+        else:
+            sample["src_path"], _ = self.scenes[i_scene][i_image]
+
+            i_scene_other, i_image_other = self._get_other_random_scene_and_image_id(i_scene)
+
+            sample["trg_path"], _ = self.scenes[i_scene_other][i_image_other]
+
+            homography = torch.zeros((3, 3), dtype=torch.float32)
+
+        src_img = read_image(sample["src_path"]) / 255.0
+        trg_img = read_image(sample["trg_path"]) / 255.0
+
+        _, H_src, W_src = src_img.shape
+        _, H_trg, W_trg = trg_img.shape
+
+        src_mask = torch.ones((1, H_src, W_src), dtype=torch.uint8)
+        trg_mask = torch.ones((1, H_trg, W_trg), dtype=torch.uint8)
+
+        if self.transforms:
+            src_img, trg_img, src_mask, trg_mask, _ = self.transforms(src_img, trg_img, src_mask, trg_mask, homography)
+
+        sample["src_image"] = src_img
+        sample["trg_image"] = trg_img
+        sample["src_mask"] = src_mask.to(torch.bool)
+        sample["trg_mask"] = trg_mask.to(torch.bool)
+        sample["homography"] = homography
+
+        return sample
+
+    def _get_scene_and_image_id_from_idx(self, idx: int) -> Tuple[int, int]:
+        accumulated_tuples = accumulate(self.tuples_per_scene)
+
+        if idx >= sum(self.tuples_per_scene):
+            raise IndexError(f"Index {idx} out of bounds")
+
+        idx_scene = None
+        for i, accumulated_tuple in enumerate(accumulated_tuples):
+            idx_scene = i
+            if idx < accumulated_tuple:
+                break
+
+        idx_image = idx - sum(self.tuples_per_scene[:idx_scene])
+
+        return idx_scene, idx_image
+
+    def _get_other_random_scene_and_image_id(self, scene_id_to_exclude: int) -> Tuple[int, int]:
+        possible_scene_ids = list(range(len(self.scenes)))
+        possible_scene_ids.remove(scene_id_to_exclude)
+
+        idx_scene = random.choice(possible_scene_ids)
+        idx_image = random.randint(0, len(self.scenes[idx_scene]) - 1)
+
+        return idx_scene, idx_image
+
+
+class Scene:
+    def __init__(self, root_path, scene_data: Dict[str, Any], max_size_scene) -> None:
+        self.root_path = root_path
+        self.image_path = Path(scene_data["image_path"])
+        self.image_names = scene_data["images"]
+
+        # randomly sample tuples
+        if max_size_scene > 0:
+            self.tuples = random.sample(scene_data["tuples"], min(max_size_scene, len(scene_data["tuples"])))
+
+    def __len__(self) -> int:
+        return len(self.tuples)
+
+    def __getitem__(self, idx: int) -> Tuple[str, str]:
+        idx_1, idx_2 = random.sample([0, 1, 2], 2)
+
+        idx_1 = self.tuples[idx][idx_1]
+        idx_2 = self.tuples[idx][idx_2]
+
+        path_image_1 = str(self.root_path / self.image_path / self.image_names[idx_1])
+        path_image_2 = str(self.root_path / self.image_path / self.image_names[idx_2])
+
+        return path_image_1, path_image_2

ripe/data/datasets/tokyo247.py

@@ -0,0 +1,134 @@
+import os
+import random
+from glob import glob
+from typing import Any, Callable, Optional
+
+import torch
+from torch.utils.data import Dataset
+from torchvision.io import read_image
+
+from ripe import utils
+from ripe.data.data_transforms import Compose
+
+log = utils.get_pylogger(__name__)
+
+
+class Tokyo247(Dataset):
+    def __init__(
+        self,
+        root: str,
+        stage: str = "train",
+        transforms: Optional[Callable] = None,
+        positive_only: bool = False,
+    ):
+        if stage != "train":
+            raise ValueError("Tokyo247Dataset only supports the 'train' stage.")
+
+        # check if the root directory exists
+        if not os.path.isdir(root):
+            raise FileNotFoundError(f"Directory {root} does not exist.")
+
+        self.root_dir = root
+        self.transforms = transforms if transforms is not None else Compose([])
+        self.positive_only = positive_only
+
+        self.image_paths = []
+        self.positive_pairs = []
+
+        # Collect images grouped by location folder
+        self.locations = {}
+        for location_rough in sorted(os.listdir(self.root_dir)):
+            location_rough_path = os.path.join(self.root_dir, location_rough)
+
+            # check if the location_rough_path is a directory
+            if not os.path.isdir(location_rough_path):
+                continue
+
+            for location_fine in sorted(os.listdir(location_rough_path)):
+                location_fine_path = os.path.join(self.root_dir, location_rough, location_fine)
+
+                if os.path.isdir(location_fine_path):
+                    images = sorted(
+                        glob(os.path.join(location_fine_path, "*.png")),
+                        key=lambda i: int(i[-7:-4]),
+                    )
+                    if len(images) >= 12:
+                        self.locations[location_fine] = images
+                        self.image_paths.extend(images)
+
+        # Generate positive pairs
+        for _, images in self.locations.items():
+            for i in range(len(images) - 1):
+                self.positive_pairs.append((images[i], images[i + 1]))
+            self.positive_pairs.append((images[-1], images[0]))
+
+        if positive_only:
+            log.warning("Using only positive pairs!")
+
+        log.info(f"Found {len(self.positive_pairs)} image pairs.")
+
+    def __len__(self):
+        if self.positive_only:
+            return len(self.positive_pairs)
+        return 2 * len(self.positive_pairs)
+
+    def __getitem__(self, idx):
+        sample: Any = {}
+
+        positive_sample = (idx % 2 == 0) or (self.positive_only)
+        if not self.positive_only:
+            idx = idx // 2
+
+        sample["label"] = positive_sample
+
+        if positive_sample:  # Positive pair
+            img1_path, img2_path = self.positive_pairs[idx]
+
+            assert os.path.dirname(img1_path) == os.path.dirname(img2_path), (
+                f"Source and target image mismatch: {img1_path} vs {img2_path}"
+            )
+
+            homography = torch.eye(3, dtype=torch.float32)
+        else:  # Negative pair
+            img1_path = random.choice(self.image_paths)
+            img2_path = random.choice(self.image_paths)
+
+            # Ensure images are from different folders
+            esc = 0
+            while os.path.dirname(img1_path) == os.path.dirname(img2_path):
+                img2_path = random.choice(self.image_paths)
+
+                esc += 1
+                if esc > 100:
+                    raise RuntimeError("Could not find a negative pair.")
+
+            assert os.path.dirname(img1_path) != os.path.dirname(img2_path), (
+                f"Source and target image match for negative pair: {img1_path} vs {img2_path}"
+            )
+
+            homography = torch.zeros((3, 3), dtype=torch.float32)
+
+        sample["src_path"] = img1_path
+        sample["trg_path"] = img2_path
+
+        # Load images
+        src_img = read_image(sample["src_path"]) / 255.0
+        trg_img = read_image(sample["trg_path"]) / 255.0
+
+        _, H_src, W_src = src_img.shape
+        _, H_trg, W_trg = src_img.shape
+
+        src_mask = torch.ones((1, H_src, W_src), dtype=torch.uint8)
+        trg_mask = torch.ones((1, H_trg, W_trg), dtype=torch.uint8)
+
+        # Apply transformations
+        if self.transforms:
+            src_img, trg_img, src_mask, trg_mask, _ = self.transforms(src_img, trg_img, src_mask, trg_mask, homography)
+
+        sample["src_image"] = src_img
+        sample["trg_image"] = trg_img
+        sample["src_mask"] = src_mask.to(torch.bool)
+        sample["trg_mask"] = trg_mask.to(torch.bool)
+        sample["homography"] = homography
+
+        return sample

ripe/losses/contrastive_loss.py

@@ -0,0 +1,88 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def second_nearest_neighbor(desc1, desc2):
+    if desc2.shape[0] < 2:  # We cannot perform snn check, so output empty matches
+        raise ValueError("desc2 should have at least 2 descriptors")
+
+    dist = torch.cdist(desc1, desc2, p=2)
+
+    vals, idxs = torch.topk(dist, 2, dim=1, largest=False)
+    idxs_in_2 = idxs[:, 1]
+    idxs_in_1 = torch.arange(0, idxs_in_2.size(0), device=dist.device)
+
+    matches_idxs = torch.cat([idxs_in_1.view(-1, 1), idxs_in_2.view(-1, 1)], 1)
+
+    return vals[:, 1].view(-1, 1), matches_idxs
+
+
+def contrastive_loss(
+    desc1,
+    desc2,
+    matches,
+    inliers,
+    label,
+    logits_1,
+    logits_2,
+    pos_margin=1.0,
+    neg_margin=1.0,
+):
+    if inliers.sum() < 8:  # if there are too few inliers, calculate loss on all matches
+        inliers = torch.ones_like(inliers)
+
+    matched_inliers_descs1 = desc1[matches[:, 0][inliers]]
+    matched_inliers_descs2 = desc2[matches[:, 1][inliers]]
+
+    if logits_1 is not None and logits_2 is not None:
+        matched_inliers_logits1 = logits_1[matches[:, 0][inliers]]
+        matched_inliers_logits2 = logits_2[matches[:, 1][inliers]]
+        logits = torch.minimum(matched_inliers_logits1, matched_inliers_logits2)
+    else:
+        logits = torch.ones_like(matches[:, 0][inliers])
+
+    if label:
+        snn_match_dists_1, idx1 = second_nearest_neighbor(matched_inliers_descs1, desc2)
+        snn_match_dists_2, idx2 = second_nearest_neighbor(matched_inliers_descs2, desc1)
+
+        dists = torch.hstack((snn_match_dists_1, snn_match_dists_2))
+        min_dists_idx = torch.min(dists, dim=1).indices.unsqueeze(1)
+
+        dists_hard = torch.gather(dists, 1, min_dists_idx).squeeze(-1)
+        dists_pos = F.pairwise_distance(matched_inliers_descs1, matched_inliers_descs2)
+
+        contrastive_loss = torch.clamp(pos_margin + dists_pos - dists_hard, min=0.0)
+
+        contrastive_loss = contrastive_loss * logits
+
+        contrastive_loss = contrastive_loss.sum() / (logits.sum() + 1e-8)  # small epsilon to avoid division by zero
+    else:
+        dists = F.pairwise_distance(matched_inliers_descs1, matched_inliers_descs2)
+        contrastive_loss = torch.clamp(neg_margin - dists, min=0.0)
+
+        contrastive_loss = contrastive_loss * logits
+
+        contrastive_loss = contrastive_loss.sum() / (logits.sum() + 1e-8)  # small epsilon to avoid division by zero
+
+    return contrastive_loss
+
+
+class ContrastiveLoss(nn.Module):
+    def __init__(self, pos_margin=1.0, neg_margin=1.0):
+        super().__init__()
+        self.pos_margin = pos_margin
+        self.neg_margin = neg_margin
+
+    def forward(self, desc1, desc2, matches, inliers, label, logits_1=None, logits_2=None):
+        return contrastive_loss(
+            desc1,
+            desc2,
+            matches,
+            inliers,
+            label,
+            logits_1,
+            logits_2,
+            self.pos_margin,
+            self.neg_margin,
+        )

ripe/matcher/concurrent_matcher.py

@@ -0,0 +1,97 @@
+import concurrent.futures
+
+import torch
+
+
+class ConcurrentMatcher:
+    """A class that performs matching and geometric filtering in parallel using a thread pool executor.
+    It matches keypoints from two sets of descriptors and applies a robust estimator to filter the matches based on geometric constraints.
+
+    Args:
+        matcher (callable): A callable that takes two sets of descriptors and returns distances and indices of matches.
+        robust_estimator (callable): A callable that estimates a geometric transformation and returns inliers.
+        min_num_matches (int, optional): Minimum number of matches required to perform geometric filtering. Defaults to 8.
+        max_workers (int, optional): Maximum number of threads in the thread pool executor. Defaults to 12.
+    """
+
+    def __init__(self, matcher, robust_estimator, min_num_matches=8, max_workers=12):
+        self.matcher = matcher
+        self.robust_estimator = robust_estimator
+        self.min_num_matches = min_num_matches
+
+        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=max_workers)
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        kpts1,
+        kpts2,
+        pdesc1,
+        pdesc2,
+        selected_mask1,
+        selected_mask2,
+        inl_th,
+        label=None,
+    ):
+        dev = pdesc1.device
+        B = pdesc1.shape[0]
+
+        batch_rel_idx_matches = [None] * B
+        batch_idx_matches = [None] * B
+        future_results = [None] * B
+
+        for b in range(B):
+            if selected_mask1[b].sum() < 16 or selected_mask2[b].sum() < 16:
+                continue
+
+            dists, idx_matches = self.matcher(pdesc1[b][selected_mask1[b]], pdesc2[b][selected_mask2[b]])
+
+            batch_rel_idx_matches[b] = idx_matches.clone()
+
+            # calculate ABSOLUTE indexes
+            idx_matches[:, 0] = torch.nonzero(selected_mask1[b], as_tuple=False)[idx_matches[:, 0]].squeeze()
+            idx_matches[:, 1] = torch.nonzero(selected_mask2[b], as_tuple=False)[idx_matches[:, 1]].squeeze()
+
+            batch_idx_matches[b] = idx_matches
+
+            # if not enough matches
+            if idx_matches.shape[0] < self.min_num_matches:
+                ransac_inliers = torch.zeros((idx_matches.shape[0]), device=dev).bool()
+                future_results[b] = (None, ransac_inliers)
+                continue
+
+            # use label information to exclude negative pairs from geometric filtering process -> enforces more descriminative descriptors
+            if label is not None and label[b] == 0:
+                ransac_inliers = torch.ones((idx_matches.shape[0]), device=dev).bool()
+                future_results[b] = (None, ransac_inliers)
+                continue
+
+            mkpts1 = kpts1[b][idx_matches[:, 0]]
+            mkpts2 = kpts2[b][idx_matches[:, 1]]
+
+            future_results[b] = self.executor.submit(self.robust_estimator, mkpts1, mkpts2, inl_th)
+
+        batch_ransac_inliers = [None] * B
+        batch_Fm = [None] * B
+
+        for b in range(B):
+            future_result = future_results[b]
+            if future_result is None:
+                ransac_inliers = None
+                Fm = None
+            elif isinstance(future_result, tuple):
+                Fm, ransac_inliers = future_result
+            else:
+                Fm, ransac_inliers = future_result.result()
+
+                # if no inliers
+                if ransac_inliers.sum() == 0:
+                    ransac_inliers = ransac_inliers.squeeze(
+                        -1
+                    )  # kornia.geometry.ransac.RANSAC returns (N, 1) tensor if no inliers and (N,) tensor if inliers
+                    Fm = None
+
+            batch_ransac_inliers[b] = ransac_inliers
+            batch_Fm[b] = Fm
+
+        return batch_rel_idx_matches, batch_idx_matches, batch_ransac_inliers, batch_Fm

ripe/matcher/pose_estimator_poselib.py

@@ -0,0 +1,31 @@
+import poselib
+import torch
+
+
+class PoseLibRelativePoseEstimator:
+    """PoseLibRelativePoseEstimator estimates the fundamental matrix using poselib library.
+    It uses the poselib's estimate_fundamental function to compute the fundamental matrix and inliers based on the provided points.
+    Args:
+        None
+    """
+
+    def __init__(self):
+        pass
+
+    def __call__(self, pts0, pts1, inl_th):
+        F, info = poselib.estimate_fundamental(
+            pts0.cpu().numpy(),
+            pts1.cpu().numpy(),
+            {
+                "max_epipolar_error": inl_th,
+            },
+        )
+
+        success = F is not None
+        if success:
+            inliers = info.pop("inliers")
+            inliers = torch.tensor(inliers, dtype=torch.bool, device=pts0.device)
+        else:
+            inliers = torch.zeros(pts0.shape[0], dtype=torch.bool, device=pts0.device)
+
+        return F, inliers

ripe/model_zoo/__init__.py

@@ -0,0 +1 @@
+from .vgg_hyper import vgg_hyper  # noqa: F401

ripe/model_zoo/vgg_hyper.py

@@ -0,0 +1,39 @@
+from pathlib import Path
+
+import torch
+
+from ripe.models.backbones.vgg import VGG
+from ripe.models.ripe import RIPE
+from ripe.models.upsampler.hypercolumn_features import HyperColumnFeatures
+
+
+def vgg_hyper(model_path: Path = None, desc_shares=None):
+    if model_path is None:
+        # check if the weights file exists in the current directory
+        model_path = Path("/tmp/ripe_weights.pth")
+
+        if model_path.exists():
+            print(f"Using existing weights from {model_path}")
+        else:
+            print("Weights file not found. Downloading ...")
+            torch.hub.download_url_to_file(
+                "https://cvg.hhi.fraunhofer.de/RIPE/ripe_weights.pth",
+                "/tmp/ripe_weights.pth",
+            )
+    else:
+        if not model_path.exists():
+            print(f"Error: {model_path} does not exist.")
+            raise FileNotFoundError(f"Error: {model_path} does not exist.")
+
+    backbone = VGG(pretrained=False)
+    upsampler = HyperColumnFeatures()
+
+    extractor = RIPE(
+        net=backbone,
+        upsampler=upsampler,
+        desc_shares=desc_shares,
+    )
+
+    extractor.load_state_dict(torch.load(model_path, map_location="cpu"))
+
+    return extractor

ripe/models/backbones/backbone_base.py

@@ -0,0 +1,61 @@
+import torch
+import torch.nn as nn
+
+
+class BackboneBase(nn.Module):
+    """Base class for backbone networks. Provides a standard interface for preprocessing inputs and
+    defining encoder dimensions.
+
+    Args:
+        nchannels (int): Number of input channels.
+        use_instance_norm (bool): Whether to apply instance normalization.
+    """
+
+    def __init__(self, nchannels=3, use_instance_norm=False):
+        super().__init__()
+        assert nchannels > 0, "Number of channels must be positive."
+        self.nchannels = nchannels
+        self.use_instance_norm = use_instance_norm
+        self.norm = nn.InstanceNorm2d(nchannels) if use_instance_norm else None
+
+    def get_dim_layers_encoder(self):
+        """Get dimensions of encoder layers."""
+        raise NotImplementedError("Subclasses must implement this method.")
+
+    def _forward(self, x):
+        """Define the forward pass for the backbone."""
+        raise NotImplementedError("Subclasses must implement this method.")
+
+    def forward(self, x: torch.Tensor, preprocess=True):
+        """Forward pass with optional preprocessing.
+
+        Args:
+            x (Tensor): Input tensor.
+            preprocess (bool): Whether to apply channel reduction.
+        """
+        if preprocess:
+            if x.dim() != 4:
+                if x.dim() == 2 and x.shape[0] > 3 and x.shape[1] > 3:
+                    x = x.unsqueeze(0).unsqueeze(0)
+                elif x.dim() == 3:
+                    x = x.unsqueeze(0)
+                else:
+                    raise ValueError(f"Unexpected input shape: {x.shape}")
+
+            if self.nchannels == 1 and x.shape[1] != 1:
+                if len(x.shape) == 4:  # Assumes (batch, channel, height, width)
+                    x = torch.mean(x, axis=1, keepdim=True)
+                else:
+                    raise ValueError(f"Unexpected input shape: {x.shape}")
+
+            #
+            if self.nchannels == 3 and x.shape[1] == 1:
+                if len(x.shape) == 4:
+                    x = x.repeat(1, 3, 1, 1)
+                else:
+                    raise ValueError(f"Unexpected input shape: {x.shape}")
+
+        if self.use_instance_norm:
+            x = self.norm(x)
+
+        return self._forward(x)

ripe/models/backbones/vgg.py

@@ -0,0 +1,99 @@
+# adapted from: https://github.com/Parskatt/DeDoDe/blob/main/DeDoDe/encoder.py and https://github.com/Parskatt/DeDoDe/blob/main/DeDoDe/decoder.py
+
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .backbone_base import BackboneBase
+from .vgg_utils import VGG19, ConvRefiner, Decoder
+
+
+class VGG(BackboneBase):
+    def __init__(self, nchannels=3, pretrained=True, use_instance_norm=True, mode="dect"):
+        super().__init__(nchannels=nchannels, use_instance_norm=use_instance_norm)
+
+        self.nchannels = nchannels
+        self.mode = mode
+
+        if self.mode not in ["dect", "desc", "dect+desc"]:
+            raise ValueError("mode should be 'dect', 'desc' or 'dect+desc'")
+
+        NUM_OUTPUT_CHANNELS, hidden_blocks = self._get_mode_params(mode)
+        conv_refiner = self._create_conv_refiner(NUM_OUTPUT_CHANNELS, hidden_blocks)
+
+        self.encoder = VGG19(pretrained=pretrained, num_input_channels=nchannels)
+        self.decoder = Decoder(conv_refiner, num_prototypes=NUM_OUTPUT_CHANNELS)
+
+    def _get_mode_params(self, mode):
+        """Get the number of output channels and the number of hidden blocks for the ConvRefiner.
+
+        Depending on the mode, the ConvRefiner will have a different number of output channels.
+        """
+
+        if mode == "dect":
+            return 1, 8
+        elif mode == "desc":
+            return 256, 5
+        elif mode == "dect+desc":
+            return 256 + 1, 8
+
+    def _create_conv_refiner(self, num_output_channels, hidden_blocks):
+        return nn.ModuleDict(
+            {
+                "8": ConvRefiner(
+                    512,
+                    512,
+                    256 + num_output_channels,
+                    hidden_blocks=hidden_blocks,
+                    residual=True,
+                ),
+                "4": ConvRefiner(
+                    256 + 256,
+                    256,
+                    128 + num_output_channels,
+                    hidden_blocks=hidden_blocks,
+                    residual=True,
+                ),
+                "2": ConvRefiner(
+                    128 + 128,
+                    128,
+                    64 + num_output_channels,
+                    hidden_blocks=hidden_blocks,
+                    residual=True,
+                ),
+                "1": ConvRefiner(
+                    64 + 64,
+                    64,
+                    1 + num_output_channels,
+                    hidden_blocks=hidden_blocks,
+                    residual=True,
+                ),
+            }
+        )
+
+    def get_dim_layers_encoder(self):
+        return self.encoder.get_dim_layers()
+
+    def _forward(self, x):
+        features, sizes = self.encoder(x)
+        output = 0
+        context = None
+        scales = self.decoder.scales
+        for idx, (feature_map, scale) in enumerate(zip(reversed(features), scales)):
+            delta_descriptor, context = self.decoder(feature_map, scale=scale, context=context)
+            output = output + delta_descriptor
+            if idx < len(scales) - 1:
+                size = sizes[-(idx + 2)]
+                output = F.interpolate(output, size=size, mode="bilinear", align_corners=False)
+                context = F.interpolate(context, size=size, mode="bilinear", align_corners=False)
+
+        if self.mode == "dect":
+            return {"heatmap": output, "coarse_descs": features}
+        elif self.mode == "desc":
+            return {"fine_descs": output, "coarse_descs": features}
+        elif self.mode == "dect+desc":
+            logits = output[:, :1].contiguous()
+            descs = output[:, 1:].contiguous()
+
+            return {"heatmap": logits, "fine_descs": descs, "coarse_descs": features}
+        else:
+            raise ValueError("mode should be 'dect', 'desc' or 'dect+desc'")

ripe/models/backbones/vgg_utils.py

@@ -0,0 +1,143 @@
+# adapted from: https://github.com/Parskatt/DeDoDe/blob/main/DeDoDe/encoder.py and https://github.com/Parskatt/DeDoDe/blob/main/DeDoDe/decoder.py
+
+import torch
+import torch.nn as nn
+import torchvision.models as tvm
+
+from ripe import utils
+
+log = utils.get_pylogger(__name__)
+
+
+class Decoder(nn.Module):
+    def __init__(self, layers, *args, super_resolution=False, num_prototypes=1, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.layers = layers
+        self.scales = self.layers.keys()
+        self.super_resolution = super_resolution
+        self.num_prototypes = num_prototypes
+
+    def forward(self, features, context=None, scale=None):
+        if context is not None:
+            features = torch.cat((features, context), dim=1)
+        stuff = self.layers[scale](features)
+        logits, context = (
+            stuff[:, : self.num_prototypes],
+            stuff[:, self.num_prototypes :],
+        )
+        return logits, context
+
+
+class ConvRefiner(nn.Module):
+    def __init__(
+        self,
+        in_dim=6,
+        hidden_dim=16,
+        out_dim=2,
+        dw=True,
+        kernel_size=5,
+        hidden_blocks=5,
+        residual=False,
+    ):
+        super().__init__()
+        self.block1 = self.create_block(
+            in_dim,
+            hidden_dim,
+            dw=False,
+            kernel_size=1,
+        )
+        self.hidden_blocks = nn.Sequential(
+            *[
+                self.create_block(
+                    hidden_dim,
+                    hidden_dim,
+                    dw=dw,
+                    kernel_size=kernel_size,
+                )
+                for hb in range(hidden_blocks)
+            ]
+        )
+        self.hidden_blocks = self.hidden_blocks
+        self.out_conv = nn.Conv2d(hidden_dim, out_dim, 1, 1, 0)
+        self.residual = residual
+
+    def create_block(
+        self,
+        in_dim,
+        out_dim,
+        dw=True,
+        kernel_size=5,
+        bias=True,
+        norm_type=nn.BatchNorm2d,
+    ):
+        num_groups = 1 if not dw else in_dim
+        if dw:
+            assert out_dim % in_dim == 0, "outdim must be divisible by indim for depthwise"
+        conv1 = nn.Conv2d(
+            in_dim,
+            out_dim,
+            kernel_size=kernel_size,
+            stride=1,
+            padding=kernel_size // 2,
+            groups=num_groups,
+            bias=bias,
+        )
+        norm = norm_type(out_dim) if norm_type is nn.BatchNorm2d else norm_type(num_channels=out_dim)
+        relu = nn.ReLU(inplace=True)
+        conv2 = nn.Conv2d(out_dim, out_dim, 1, 1, 0)
+        return nn.Sequential(conv1, norm, relu, conv2)
+
+    def forward(self, feats):
+        b, c, hs, ws = feats.shape
+        x0 = self.block1(feats)
+        x = self.hidden_blocks(x0)
+        if self.residual:
+            x = (x + x0) / 1.4
+        x = self.out_conv(x)
+        return x
+
+
+class VGG19(nn.Module):
+    def __init__(self, pretrained=False, num_input_channels=3) -> None:
+        super().__init__()
+        self.layers = nn.ModuleList(tvm.vgg19_bn(pretrained=pretrained).features[:40])
+        # Maxpool layers: 6, 13, 26, 39
+
+        if num_input_channels != 3:
+            log.info(f"Changing input channels from 3 to {num_input_channels}")
+            self.layers[0] = nn.Conv2d(num_input_channels, 64, 3, 1, 1)
+
+    def get_dim_layers(self):
+        return [64, 128, 256, 512]
+
+    def forward(self, x, **kwargs):
+        feats = []
+        sizes = []
+        for layer in self.layers:
+            if isinstance(layer, nn.MaxPool2d):
+                feats.append(x)
+                sizes.append(x.shape[-2:])
+            x = layer(x)
+        return feats, sizes
+
+
+class VGG(nn.Module):
+    def __init__(self, size="19", pretrained=False) -> None:
+        super().__init__()
+        if size == "11":
+            self.layers = nn.ModuleList(tvm.vgg11_bn(pretrained=pretrained).features[:22])
+        elif size == "13":
+            self.layers = nn.ModuleList(tvm.vgg13_bn(pretrained=pretrained).features[:28])
+        elif size == "19":
+            self.layers = nn.ModuleList(tvm.vgg19_bn(pretrained=pretrained).features[:40])
+        # Maxpool layers: 6, 13, 26, 39
+
+    def forward(self, x, **kwargs):
+        feats = []
+        sizes = []
+        for layer in self.layers:
+            if isinstance(layer, nn.MaxPool2d):
+                feats.append(x)
+                sizes.append(x.shape[-2:])
+            x = layer(x)
+        return feats, sizes

ripe/models/ripe.py

@@ -0,0 +1,303 @@
+from typing import List, Optional
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ripe import utils
+from ripe.utils.utils import gridify
+
+log = utils.get_pylogger(__name__)
+
+
+class KeypointSampler(nn.Module):
+    """
+    Sample keypoints according to a Heatmap
+    Adapted from: https://github.com/verlab/DALF_CVPR_2023/blob/main/modules/models/DALF.py
+    """
+
+    def __init__(self, window_size=8):
+        super().__init__()
+        self.window_size = window_size
+        self.idx_cells = None  # Cache for meshgrid indices
+
+    def sample(self, grid):
+        """
+        Sample keypoints given a grid where each cell has logits stacked in last dimension
+        Input
+          grid: [B, C, H//w, W//w, w*w]
+
+        Returns
+          log_probs: [B, C, H//w, W//w ] - logprobs of selected samples
+          choices: [B, C, H//w, W//w] indices of choices
+          accept_mask: [B, C, H//w, W//w] mask of accepted keypoints
+
+        """
+        chooser = torch.distributions.Categorical(logits=grid)
+        choices = chooser.sample()
+        logits_selected = torch.gather(grid, -1, choices.unsqueeze(-1)).squeeze(-1)
+
+        flipper = torch.distributions.Bernoulli(logits=logits_selected)
+        accepted_choices = flipper.sample()
+
+        # Sum log-probabilities is equivalent to multiplying the probabilities
+        log_probs = chooser.log_prob(choices) + flipper.log_prob(accepted_choices)
+
+        accept_mask = accepted_choices.gt(0)
+
+        return (
+            log_probs.squeeze(1),
+            choices,
+            accept_mask.squeeze(1),
+            logits_selected.squeeze(1),
+        )
+
+    def precompute_idx_cells(self, H, W, device):
+        idx_cells = gridify(
+            torch.dstack(
+                torch.meshgrid(
+                    torch.arange(H, dtype=torch.float32, device=device),
+                    torch.arange(W, dtype=torch.float32, device=device),
+                )
+            )
+            .permute(2, 0, 1)
+            .unsqueeze(0)
+            .expand(1, -1, -1, -1),
+            window_size=self.window_size,
+        )
+
+        return idx_cells
+
+    def forward(self, x, mask_padding=None):
+        """
+        Sample keypoints from a heatmap
+        Input
+          x: [B, C, H, W] Heatmap
+          mask_padding: [B, 1, H, W] Mask for padding (optional)
+        Returns
+            keypoints: [B, H//w, W//w, 2] Keypoints in (x, y) format
+            log_probs: [B, H//w, W//w] Log probabilities of selected keypoints
+            mask: [B, H//w, W//w] Mask of accepted keypoints
+            mask_padding: [B, 1, H//w, W//w] Mask of padding (optional)
+            logits_selected: [B, H//w, W//w] Logits of selected keypoints
+        """
+
+        B, C, H, W = x.shape
+
+        keypoint_cells = gridify(x, self.window_size)
+
+        mask_padding = (
+            (torch.min(gridify(mask_padding, self.window_size), dim=4).values) if mask_padding is not None else None
+        )
+
+        if self.idx_cells is None or self.idx_cells.shape[2:4] != (
+            H // self.window_size,
+            W // self.window_size,
+        ):
+            self.idx_cells = self.precompute_idx_cells(H, W, x.device)
+
+        log_probs, idx, mask, logits_selected = self.sample(keypoint_cells)
+
+        keypoints = (
+            torch.gather(
+                self.idx_cells.expand(B, -1, -1, -1, -1),
+                -1,
+                idx.repeat(1, 2, 1, 1).unsqueeze(-1),
+            )
+            .squeeze(-1)
+            .permute(0, 2, 3, 1)
+        )
+
+        # flip keypoints to (x, y) format
+        return keypoints.flip(-1), log_probs, mask, mask_padding, logits_selected
+
+
+class RIPE(nn.Module):
+    """
+    Base class for extracting keypoints and descriptors
+    Input
+      x: [B, C, H, W] Images
+
+    Returns
+      kpts:
+        list of size [B] with detected keypoints
+      descs:
+        list of size [B] with descriptors
+    """
+
+    def __init__(
+        self,
+        net,
+        upsampler,
+        window_size: int = 8,
+        non_linearity_dect=None,
+        desc_shares: Optional[List[int]] = None,
+        descriptor_dim: int = 256,
+        device=None,
+    ):
+        super().__init__()
+        self.net = net
+
+        self.detector = KeypointSampler(window_size)
+        self.upsampler = upsampler
+        self.sampler = None
+        self.window_size = window_size
+        self.non_linearity_dect = non_linearity_dect if non_linearity_dect is not None else nn.Identity()
+
+        log.info(f"Training with window size {window_size}.")
+        log.info(f"Use {non_linearity_dect} as final non-linearity before the detection heatmap.")
+
+        dim_coarse_desc = self.get_dim_raw_desc()
+
+        if desc_shares is not None:
+            assert upsampler.name == "HyperColumnFeatures", (
+                "Individual descriptor convolutions are only supported with HyperColumnFeatures"
+            )
+            assert len(desc_shares) == 4, "desc_shares should have 4 elements"
+            assert sum(desc_shares) == descriptor_dim, f"sum of desc_shares should be {descriptor_dim}"
+
+            self.conv_dim_reduction_coarse_desc = nn.ModuleList()
+
+            for dim_in, dim_out in zip(dim_coarse_desc, desc_shares):
+                log.info(f"Training dim reduction descriptor with {dim_in} -> {dim_out} 1x1 conv")
+                self.conv_dim_reduction_coarse_desc.append(
+                    nn.Conv1d(dim_in, dim_out, kernel_size=1, stride=1, padding=0)
+                )
+        else:
+            if descriptor_dim is not None:
+                log.info(f"Training dim reduction descriptor with {sum(dim_coarse_desc)} -> {descriptor_dim} 1x1 conv")
+                self.conv_dim_reduction_coarse_desc = nn.Conv1d(
+                    sum(dim_coarse_desc),
+                    descriptor_dim,
+                    kernel_size=1,
+                    stride=1,
+                    padding=0,
+                )
+            else:
+                log.warning(
+                    f"No descriptor dimension specified, no 1x1 conv will be applied! Direct usage of {sum(dim_coarse_desc)}-dimensional raw descriptor"
+                )
+                self.conv_dim_reduction_coarse_desc = nn.Identity()
+
+    def get_dim_raw_desc(self):
+        layers_dims_encoder = self.net.get_dim_layers_encoder()
+
+        if self.upsampler.name == "InterpolateSparse2d":
+            return [layers_dims_encoder[-1]]
+        elif self.upsampler.name == "HyperColumnFeatures":
+            return layers_dims_encoder
+        else:
+            raise ValueError(f"Unknown interpolator {self.upsampler.name}")
+
+    @torch.inference_mode()
+    def detectAndCompute(self, img, threshold=0.5, top_k=2048, output_aux=False):
+        self.train(False)
+
+        if img.dim() == 3:
+            img = img.unsqueeze(0)
+
+        out = self(img, training=False)
+        B, K, H, W = out["heatmap"].shape
+
+        assert B == 1, "Batch size should be 1"
+
+        kpts = [{"xy": self.NMS(out["heatmap"][b], threshold)} for b in range(B)]
+
+        if top_k is not None:
+            for b in range(B):
+                scores = out["heatmap"][b].squeeze(0)[kpts[b]["xy"][:, 1].long(), kpts[b]["xy"][:, 0].long()]
+                sorted_idx = torch.argsort(-scores)
+                kpts[b]["xy"] = kpts[b]["xy"][sorted_idx[:top_k]]
+                if "logprobs" in kpts[b]:
+                    kpts[b]["logprobs"] = kpts[b]["xy"][sorted_idx[:top_k]]
+
+        if kpts[0]["xy"].shape[0] == 0:
+            raise RuntimeError("No keypoints detected")
+
+        # the following works for batch size 1 only
+
+        descs = self.get_descs(out["coarse_descs"], img, kpts[0]["xy"].unsqueeze(0), H, W)
+        descs = descs.squeeze(0)
+
+        score_map = out["heatmap"][0].squeeze(0)
+
+        kpts = kpts[0]["xy"]
+
+        scores = score_map[kpts[:, 1], kpts[:, 0]]
+        scores /= score_map.max()
+
+        sort_idx = torch.argsort(-scores)
+        kpts, descs, scores = kpts[sort_idx], descs[sort_idx], scores[sort_idx]
+
+        if output_aux:
+            return (
+                kpts.float(),
+                descs,
+                scores,
+                {
+                    "heatmap": out["heatmap"],
+                    "descs": out["coarse_descs"],
+                    "conv": self.conv_dim_reduction_coarse_desc,
+                },
+            )
+
+        return kpts.float(), descs, scores
+
+    def NMS(self, x, threshold=3.0, kernel_size=3):
+        pad = kernel_size // 2
+        local_max = nn.MaxPool2d(kernel_size=kernel_size, stride=1, padding=pad)(x)
+
+        pos = (x == local_max) & (x > threshold)
+        return pos.nonzero()[..., 1:].flip(-1)
+
+    def get_descs(self, feature_map, guidance, kpts, H, W):
+        descs = self.upsampler(feature_map, kpts, H, W)
+
+        if isinstance(self.conv_dim_reduction_coarse_desc, nn.ModuleList):
+            # individual descriptor convolutions for each layer
+            desc_conv = []
+            for desc, conv in zip(descs, self.conv_dim_reduction_coarse_desc):
+                desc_conv.append(conv(desc.permute(0, 2, 1)).permute(0, 2, 1))
+            desc = torch.cat(desc_conv, dim=-1)
+        else:
+            desc = torch.cat(descs, dim=-1)
+            desc = self.conv_dim_reduction_coarse_desc(desc.permute(0, 2, 1)).permute(0, 2, 1)
+
+        desc = F.normalize(desc, dim=2)
+
+        return desc
+
+    def forward(self, x, mask_padding=None, training=False):
+        B, C, H, W = x.shape
+        out = self.net(x)
+        out["heatmap"] = self.non_linearity_dect(out["heatmap"])
+        # print(out['map'].shape, out['descr'].shape)
+        if training:
+            kpts, log_probs, mask, mask_padding, logits_selected = self.detector(out["heatmap"], mask_padding)
+
+            filter_A = kpts[:, :, :, 0] >= 16
+            filter_B = kpts[:, :, :, 1] >= 16
+            filter_C = kpts[:, :, :, 0] < W - 16
+            filter_D = kpts[:, :, :, 1] < H - 16
+            filter_all = filter_A * filter_B * filter_C * filter_D
+
+            mask = mask * filter_all
+
+            return (
+                kpts.view(B, -1, 2),
+                log_probs.view(B, -1),
+                mask.view(B, -1),
+                mask_padding.view(B, -1),
+                logits_selected.view(B, -1),
+                out,
+            )
+        else:
+            return out
+
+
+def output_number_trainable_params(model):
+    model_parameters = filter(lambda p: p.requires_grad, model.parameters())
+    nb_params = sum([np.prod(p.size()) for p in model_parameters])
+
+    print(f"Number of trainable parameters: {nb_params:d}")

ripe/models/upsampler/hypercolumn_features.py

@@ -0,0 +1,54 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class HyperColumnFeatures(nn.Module):
+    """
+    Interpolate 3D tensor given N sparse 2D positions
+    Input
+      x: list([C, H, W]) list of feature tensors at different scales (e.g. from a U-Net) -> extract hypercolumn features
+      pos: [N, 2] tensor of positions
+      H: int, height of the OUTPUT map
+      W: int, width of the OUTPUT map
+
+    Returns
+      [N, C] sampled features at 2d positions
+    """
+
+    def __init__(self, mode="bilinear"):
+        super().__init__()
+        self.mode = mode
+        self.name = "HyperColumnFeatures"
+
+    def normgrid(self, x, H, W):
+        return 2.0 * (x / (torch.tensor([W - 1, H - 1], device=x.device, dtype=x.dtype))) - 1.0
+
+    def extract_values_at_poses(self, x, pos, H, W):
+        """Extract values from tensor x at the positions given by pos.
+
+        Args:
+        - x (Tensor): Tensor of size (C, H, W).
+        - pos (Tensor): Tensor of size (N, 2) containing the x, y positions.
+
+        Returns:
+        - values (Tensor): Tensor of size (N, C) with the values from f at the positions given by p.
+        """
+
+        # check if grid is float32
+        if x.dtype != torch.float32:
+            x = x.to(torch.float32)
+
+        grid = self.normgrid(pos, H, W).unsqueeze(-2)
+
+        x = F.grid_sample(x, grid, mode=self.mode, align_corners=True)
+        return x.permute(0, 2, 3, 1).squeeze(-2)
+
+    def forward(self, x, pos, H, W):
+        descs = []
+
+        for layer in x:
+            desc = self.extract_values_at_poses(layer, pos, H, W)
+            descs.append(desc)
+
+        return descs

ripe/models/upsampler/interpolate_sparse2d.py

@@ -0,0 +1,37 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class InterpolateSparse2d(nn.Module):
+    """
+    Interpolate 3D tensor given N sparse 2D positions
+    Input
+      x: list([C, H, W]) feature tensors at different scales (e.g. from a U-Net), ONLY the last one is used
+      pos: [N, 2] tensor of positions
+      H: int, height of the OUTPUT map
+      W: int, width of the OUTPUT map
+
+    Returns
+      [N, C] sampled features at 2d positions
+    """
+
+    def __init__(self, mode="bicubic"):
+        super().__init__()
+        self.mode = mode
+        self.name = "InterpolateSparse2d"
+
+    def normgrid(self, x, H, W):
+        return 2.0 * (x / (torch.tensor([W - 1, H - 1], device=x.device, dtype=x.dtype))) - 1.0
+
+    def forward(self, x, pos, H, W):
+        x = x[-1]  # only use the last layer
+
+        # check if grid is float32
+        if x.dtype != torch.float32:
+            x = x.to(torch.float32)
+
+        grid = self.normgrid(pos, H, W).unsqueeze(-2)
+
+        x = F.grid_sample(x, grid, mode=self.mode, align_corners=True)
+        return [x.permute(0, 2, 3, 1).squeeze(-2)]

ripe/scheduler/constant.py

@@ -0,0 +1,6 @@
+class ConstantScheduler:
+    def __init__(self, value):
+        self.value = value
+
+    def __call__(self, step):
+        return self.value

ripe/scheduler/expDecay.py

@@ -0,0 +1,26 @@
+import numpy as np
+
+from ripe import utils
+
+log = utils.get_pylogger(__name__)
+
+
+class ExpDecay:
+    """Exponential decay scheduler.
+    args:
+        a: float, a + c = initial value
+        b: decay rate
+        c: float, final value
+
+        f(x) = a * e^(-b * x) + c
+    """
+
+    def __init__(self, a, b, c):
+        self.a = a
+        self.b = b
+        self.c = c
+
+        log.info(f"ExpDecay: a={a}, b={b}, c={c}")
+
+    def __call__(self, step):
+        return self.a * np.exp(-self.b * step) + self.c

ripe/scheduler/linearLR.py

@@ -0,0 +1,37 @@
+class StepLinearLR:
+    """Decay the learning rate by a linearly changing factor at each STEP (not epoch).
+
+    Args:
+        optimizer (Optimizer): Wrapped optimizer.
+        num_steps (int): Total number of steps in the training process.
+        initial_lr (float): Initial learning rate.
+        final_lr (float): Final learning rate.
+    """
+
+    def __init__(self, optimizer, steps_init, num_steps, initial_lr, final_lr):
+        self.optimizer = optimizer
+        self.num_steps = num_steps
+        self.initial_lr = initial_lr
+        self.final_lr = final_lr
+        self.i_step = steps_init
+        self.decay_factor = (final_lr - initial_lr) / num_steps
+
+    def step(self):
+        """Decay the learning rate by decay_factor."""
+        self.i_step += 1
+
+        if self.i_step > self.num_steps:
+            return
+
+        lr = self.initial_lr + self.i_step * self.decay_factor
+        for param_group in self.optimizer.param_groups:
+            param_group["lr"] = lr
+
+    def get_lr(self):
+        return self.optimizer.param_groups[0]["lr"]
+
+    def get_last_lr(self):
+        return self.optimizer.param_groups[0]["lr"]
+
+    def get_step(self):
+        return self.i_step

ripe/scheduler/linear_with_plateaus.py

@@ -0,0 +1,44 @@
+from ripe import utils
+
+log = utils.get_pylogger(__name__)
+
+
+class LinearWithPlateaus:
+    """Linear scheduler with plateaus.
+
+    Linearly increases from `start_val` to `end_val`.
+    Stays at `start_val` for `plateau_start_steps` steps and at `end_val` for `plateau_end_steps` steps.
+    Linearly changes from `start_val` to `end_val` during the remaining steps.
+    """
+
+    def __init__(
+        self,
+        start_val,
+        end_val,
+        steps_total,
+        rel_length_start_plateau=0.0,
+        rel_length_end_plateu=0.0,
+    ):
+        self.start_val = start_val
+        self.end_val = end_val
+        self.steps_total = steps_total
+        self.plateau_start_steps = steps_total * rel_length_start_plateau
+        self.plateau_end_steps = steps_total * rel_length_end_plateu
+
+        assert self.plateau_start_steps >= 0
+        assert self.plateau_end_steps >= 0
+        assert self.plateau_start_steps + self.plateau_end_steps <= self.steps_total
+
+        self.slope = (end_val - start_val) / (steps_total - self.plateau_start_steps - self.plateau_end_steps)
+
+        log.info(
+            f"LinearWithPlateaus: start_val={start_val}, end_val={end_val}, steps_total={steps_total}, "
+            f"plateau_start_steps={self.plateau_start_steps}, plateau_end_steps={self.plateau_end_steps}"
+        )
+
+    def __call__(self, step):
+        if step < self.plateau_start_steps:
+            return self.start_val
+        if step < self.steps_total - self.plateau_end_steps:
+            return self.start_val + self.slope * (step - self.plateau_start_steps)
+        return self.end_val

ripe/train.py

@@ -0,0 +1,410 @@
+import pyrootutils
+
+root = pyrootutils.setup_root(
+    search_from=__file__,
+    indicator=[".git", "pyproject.toml"],
+    pythonpath=True,
+    dotenv=True,
+)
+
+SEED = 32000
+
+import collections
+import os
+
+import hydra
+from hydra.utils import instantiate
+from lightning.fabric import Fabric
+
+print(SEED)
+import random
+
+os.environ["PYTHONHASHSEED"] = str(SEED)
+
+import numpy as np
+import torch
+import tqdm
+import wandb
+from torch.optim.adamw import AdamW
+from torch.utils.data import DataLoader
+
+from ripe import utils
+from ripe.benchmarks.imw_2020 import IMW_2020_Benchmark
+from ripe.utils.utils import get_rewards
+from ripe.utils.wandb_utils import get_flattened_wandb_cfg
+
+log = utils.get_pylogger(__name__)
+from pathlib import Path
+
+torch.manual_seed(SEED)
+np.random.seed(SEED)
+random.seed(SEED)
+
+
+def unpack_batch(batch):
+    src_image = batch["src_image"]
+    trg_image = batch["trg_image"]
+    trg_mask = batch["trg_mask"]
+    src_mask = batch["src_mask"]
+    label = batch["label"]
+    H = batch["homography"]
+
+    return src_image, trg_image, src_mask, trg_mask, H, label
+
+
+@hydra.main(config_path="../conf/", config_name="config", version_base=None)
+def train(cfg):
+    """Main training function for the RIPE model."""
+    #  Prepare model, data and hyperparms
+
+    strategy = "ddp" if cfg.num_gpus > 1 else "auto"
+    fabric = Fabric(
+        accelerator="cuda",
+        devices=cfg.num_gpus,
+        precision=cfg.precision,
+        strategy=strategy,
+    )
+    fabric.launch()
+
+    output_dir = Path(cfg.output_dir)
+    experiment_name = output_dir.parent.parent.parent.name
+    run_id = output_dir.parent.parent.name
+    timestamp = output_dir.parent.name + "_" + output_dir.name
+
+    experiment_name = run_id + " " + timestamp + " " + experiment_name
+
+    # setup logger
+    wandb_logger = wandb.init(
+        project=cfg.project_name,
+        name=experiment_name,
+        config=get_flattened_wandb_cfg(cfg),
+        dir=cfg.output_dir,
+        mode=cfg.wandb_mode,
+    )
+
+    min_nums_matches = {"homography": 4, "fundamental": 8, "fundamental_7pt": 7}
+    min_num_matches = min_nums_matches[cfg.transformation_model]
+    print(f"Minimum number of matches for {cfg.transformation_model} is {min_num_matches}")
+
+    batch_size = cfg.batch_size
+    steps = cfg.num_steps
+    lr = cfg.lr
+
+    num_grad_accs = (
+        cfg.num_grad_accs
+    )  # this performs grad accumulation to simulate larger batch size, set to 1 to disable;
+
+    # instantiate dataset
+    ds = instantiate(cfg.data)
+
+    # prepare dataloader
+    dl = DataLoader(
+        ds,
+        batch_size=batch_size,
+        shuffle=True,
+        drop_last=True,
+        persistent_workers=False,
+        num_workers=cfg.num_workers,
+    )
+    dl = fabric.setup_dataloaders(dl)
+    i_dl = iter(dl)
+
+    # create matcher
+    matcher = instantiate(cfg.matcher)
+
+    if cfg.desc_loss_weight != 0.0:
+        descriptor_loss = instantiate(cfg.descriptor_loss)
+    else:
+        log.warning(
+            "Descriptor loss weight is 0.0, descriptor loss will not be used. 1x1 conv for descriptors will be deactivated!"
+        )
+        descriptor_loss = None
+
+    upsampler = instantiate(cfg.upsampler) if "upsampler" in cfg else None
+
+    # create network
+    net = instantiate(cfg.network)(
+        net=instantiate(cfg.backbones),
+        upsampler=upsampler,
+        descriptor_dim=cfg.descriptor_dim if descriptor_loss is not None else None,
+        device=fabric.device,
+    ).train()
+
+    # get num parameters
+    num_params = sum(p.numel() for p in net.parameters() if p.requires_grad)
+    log.info(f"Number of parameters: {num_params}")
+
+    fp_penalty = cfg.fp_penalty  # small penalty for not finding a match
+    kp_penalty = cfg.kp_penalty  # small penalty for low logprob keypoints
+
+    opt_pi = AdamW(filter(lambda x: x.requires_grad, net.parameters()), lr=lr, weight_decay=1e-5)
+    net, opt_pi = fabric.setup(net, opt_pi)
+
+    if cfg.lr_scheduler:
+        scheduler = instantiate(cfg.lr_scheduler)(optimizer=opt_pi, steps_init=0)
+    else:
+        scheduler = None
+
+    val_benchmark = IMW_2020_Benchmark(
+        use_predefined_subset=True,
+        conf_inference=cfg.conf_inference,
+        edge_input_divisible_by=None,
+    )
+
+    # mean average of skipped batches
+    # this is used to monitor how many batches were skipped due to not enough keypoints
+    # this is useful to detect if the model is not learning anything -> should be zero
+    ma_skipped_batches = collections.deque(maxlen=100)
+
+    opt_pi.zero_grad()
+
+    # initialize scheduler
+    alpha_scheduler = instantiate(cfg.alpha_scheduler)
+    beta_scheduler = instantiate(cfg.beta_scheduler)
+    inl_th_scheduler = instantiate(cfg.inl_th)
+
+    # ======  Training Loop  ======
+    # check if the model is in training mode
+    net.train()
+
+    with tqdm.tqdm(total=steps) as pbar:
+        for i_step in range(steps):
+            alpha = alpha_scheduler(i_step)
+            beta = beta_scheduler(i_step)
+            inl_th = inl_th_scheduler(i_step)
+
+            if scheduler:
+                scheduler.step()
+
+            # Initialize vars for current step
+            # We need to handle batching because the description can have arbitrary number of keypoints
+            sum_reward_batch = 0
+            sum_num_keypoints_1 = 0
+            sum_num_keypoints_2 = 0
+            loss = None
+            loss_policy_stack = None
+            loss_desc_stack = None
+            loss_kp_stack = None
+
+            try:
+                batch = next(i_dl)
+            except StopIteration:
+                i_dl = iter(dl)
+                batch = next(i_dl)
+
+            p1, p2, mask_padding_1, mask_padding_2, Hs, label = unpack_batch(batch)
+
+            (
+                kpts1,
+                logprobs1,
+                selected_mask1,
+                mask_padding_grid_1,
+                logits_selected_1,
+                out1,
+            ) = net(p1, mask_padding_1, training=True)
+            (
+                kpts2,
+                logprobs2,
+                selected_mask2,
+                mask_padding_grid_2,
+                logits_selected_2,
+                out2,
+            ) = net(p2, mask_padding_2, training=True)
+
+            # upsample coarse descriptors for all keypoints from the intermediate feature maps from the encoder
+            desc_1 = net.get_descs(out1["coarse_descs"], p1, kpts1, p1.shape[2], p1.shape[3])
+            desc_2 = net.get_descs(out2["coarse_descs"], p2, kpts2, p2.shape[2], p2.shape[3])
+
+            if cfg.padding_filter_mode == "ignore":  # remove keypoints that are in padding
+                batch_mask_selection_for_matching_1 = selected_mask1 & mask_padding_grid_1
+                batch_mask_selection_for_matching_2 = selected_mask2 & mask_padding_grid_2
+            elif cfg.padding_filter_mode == "punish":
+                batch_mask_selection_for_matching_1 = selected_mask1  # keep all keypoints
+                batch_mask_selection_for_matching_2 = selected_mask2  # punish the keypoints in the padding area
+            else:
+                raise ValueError(f"Unknown padding filter mode: {cfg.padding_filter_mode}")
+
+            (
+                batch_rel_idx_matches,
+                batch_abs_idx_matches,
+                batch_ransac_inliers,
+                batch_Fm,
+            ) = matcher(
+                kpts1,
+                kpts2,
+                desc_1,
+                desc_2,
+                batch_mask_selection_for_matching_1,
+                batch_mask_selection_for_matching_2,
+                inl_th,
+                label if cfg.no_filtering_negatives else None,
+            )
+
+            for b in range(batch_size):
+                # ignore if less than 16 keypoints have been detected
+                if batch_rel_idx_matches[b] is None:
+                    ma_skipped_batches.append(1)
+                    continue
+                else:
+                    ma_skipped_batches.append(0)
+
+                mask_selection_for_matching_1 = batch_mask_selection_for_matching_1[b]
+                mask_selection_for_matching_2 = batch_mask_selection_for_matching_2[b]
+
+                rel_idx_matches = batch_rel_idx_matches[b]
+                abs_idx_matches = batch_abs_idx_matches[b]
+                ransac_inliers = batch_ransac_inliers[b]
+
+                if cfg.selected_only:
+                    # every SELECTED keypoint with every other SELECTED keypoint
+                    dense_logprobs = logprobs1[b][mask_selection_for_matching_1].view(-1, 1) + logprobs2[b][
+                        mask_selection_for_matching_2
+                    ].view(1, -1)
+                else:
+                    if cfg.padding_filter_mode == "ignore":
+                        # every keypoint with every other keypoint, but WITHOUT keypoint in the padding area
+                        dense_logprobs = logprobs1[b][mask_padding_grid_1[b]].view(-1, 1) + logprobs2[b][
+                            mask_padding_grid_2[b]
+                        ].view(1, -1)
+                    elif cfg.padding_filter_mode == "punish":
+                        # every keypoint with every other keypoint, also WITH keypoints in the padding areas -> will be punished by the reward
+                        dense_logprobs = logprobs1[b].view(-1, 1) + logprobs2[b].view(1, -1)
+                    else:
+                        raise ValueError(f"Unknown padding filter mode: {cfg.padding_filter_mode}")
+
+                reward = None
+
+                if cfg.reward_type == "inlier":
+                    reward = (
+                        0.5 if cfg.no_filtering_negatives and not label[b] else 1.0
+                    )  # reward is 1.0 if the pair is positive, 0.5 if negative and no filtering is applied
+                elif cfg.reward_type == "inlier_ratio":
+                    ratio_inlier = ransac_inliers.sum() / len(abs_idx_matches)
+                    reward = ratio_inlier  # reward is the ratio of inliers -> higher if more matches are inliers
+                elif cfg.reward_type == "inlier+inlier_ratio":
+                    ratio_inlier = ransac_inliers.sum() / len(abs_idx_matches)
+                    reward = (
+                        (1.0 - beta) * 1.0 + beta * ratio_inlier
+                    )  # reward is a combination of the ratio of inliers and the number of inliers -> gradually changes
+                else:
+                    raise ValueError(f"Unknown reward type: {cfg.reward_type}")
+
+                dense_rewards = get_rewards(
+                    reward,
+                    kpts1[b],
+                    kpts2[b],
+                    mask_selection_for_matching_1,
+                    mask_selection_for_matching_2,
+                    mask_padding_grid_1[b],
+                    mask_padding_grid_2[b],
+                    rel_idx_matches,
+                    abs_idx_matches,
+                    ransac_inliers,
+                    label[b],
+                    fp_penalty * alpha,
+                    use_whitening=cfg.use_whitening,
+                    selected_only=cfg.selected_only,
+                    filter_mode=cfg.padding_filter_mode,
+                )
+
+                if descriptor_loss is not None:
+                    hard_loss = descriptor_loss(
+                        desc1=desc_1[b],
+                        desc2=desc_2[b],
+                        matches=abs_idx_matches,
+                        inliers=ransac_inliers,
+                        label=label[b],
+                        logits_1=None,
+                        logits_2=None,
+                    )
+                    loss_desc_stack = (
+                        hard_loss if loss_desc_stack is None else torch.hstack((loss_desc_stack, hard_loss))
+                    )
+
+                sum_reward_batch += dense_rewards.sum()
+
+                current_loss_policy = (dense_rewards * dense_logprobs).view(-1)
+
+                loss_policy_stack = (
+                    current_loss_policy
+                    if loss_policy_stack is None
+                    else torch.hstack((loss_policy_stack, current_loss_policy))
+                )
+
+                if kp_penalty != 0.0:
+                    # keypoints with low logprob are penalized
+                    # as they get large negative logprob values multiplying them with the penalty will make the loss larger
+                    loss_kp = (
+                        logprobs1[b][mask_selection_for_matching_1]
+                        * torch.full_like(
+                            logprobs1[b][mask_selection_for_matching_1],
+                            kp_penalty * alpha,
+                        )
+                    ).mean() + (
+                        logprobs2[b][mask_selection_for_matching_2]
+                        * torch.full_like(
+                            logprobs2[b][mask_selection_for_matching_2],
+                            kp_penalty * alpha,
+                        )
+                    ).mean()
+                    loss_kp_stack = loss_kp if loss_kp_stack is None else torch.hstack((loss_kp_stack, loss_kp))
+
+                sum_num_keypoints_1 += mask_selection_for_matching_1.sum()
+                sum_num_keypoints_2 += mask_selection_for_matching_2.sum()
+
+            loss = loss_policy_stack.mean()
+            if loss_kp_stack is not None:
+                loss += loss_kp_stack.mean()
+
+            loss = -loss
+
+            if descriptor_loss is not None:
+                loss += cfg.desc_loss_weight * loss_desc_stack.mean()
+
+            pbar.set_description(
+                f"LP: {loss.item():.4f} - Det: ({sum_num_keypoints_1 / batch_size:.4f}, {sum_num_keypoints_2 / batch_size:.4f}), #mRwd: {sum_reward_batch / batch_size:.1f}"
+            )
+            pbar.update()
+
+            # backward pass
+            loss /= num_grad_accs
+            fabric.backward(loss)
+
+            if i_step % num_grad_accs == 0:
+                opt_pi.step()
+                opt_pi.zero_grad()
+
+            if i_step % cfg.log_interval == 0:
+                wandb_logger.log(
+                    {
+                        # "loss": loss.item() if not use_amp else scaled_loss.item(),
+                        "loss": loss.item(),
+                        "loss_policy": -loss_policy_stack.mean().item(),
+                        "loss_kp": loss_kp_stack.mean().item() if loss_kp_stack is not None else 0.0,
+                        "loss_hard": (loss_desc_stack.mean().item() if loss_desc_stack is not None else 0.0),
+                        "mean_num_det_kpts1": sum_num_keypoints_1 / batch_size,
+                        "mean_num_det_kpts2": sum_num_keypoints_2 / batch_size,
+                        "mean_reward": sum_reward_batch / batch_size,
+                        "lr": opt_pi.param_groups[0]["lr"],
+                        "ma_skipped_batches": sum(ma_skipped_batches) / len(ma_skipped_batches),
+                        "inl_th": inl_th,
+                    },
+                    step=i_step,
+                )
+
+            if i_step % cfg.val_interval == 0:
+                val_benchmark.evaluate(net, fabric.device, progress_bar=False)
+                val_benchmark.log_results(logger=wandb_logger, step=i_step)
+
+                # ensure that the model is in training mode again
+                net.train()
+
+    # save the model
+    torch.save(
+        net.state_dict(),
+        output_dir / ("model" + "_" + str(i_step + 1) + "_final" + ".pth"),
+    )
+
+
+if __name__ == "__main__":
+    train()

ripe/utils/__init__.py

@@ -0,0 +1,2 @@
+# from x_dd.utils import loggers
+from ripe.utils.pylogger import get_pylogger  # noqa: F401

ripe/utils/image_utils.py

@@ -0,0 +1,62 @@
+import h5py
+import numpy as np
+import torch
+
+
+class Camera:
+    def __init__(self, K, R, t):
+        self.K = K
+        self.R = R
+        self.t = t
+
+    @classmethod
+    def from_calibration_file(cls, path: str):
+        with h5py.File(path, "r") as f:
+            K = torch.tensor(np.array(f["K"]), dtype=torch.float32)
+            R = torch.tensor(np.array(f["R"]), dtype=torch.float32)
+            T = torch.tensor(np.array(f["T"]), dtype=torch.float32)
+
+        return cls(K, R, T)
+
+    @property
+    def K_inv(self):
+        return self.K.inverse()
+
+    def to_cameradict(self):
+        fx = self.K[0, 0].item()
+        fy = self.K[1, 1].item()
+        cx = self.K[0, 2].item()
+        cy = self.K[1, 2].item()
+
+        params = {
+            "model": "PINHOLE",
+            "width": int(cx * 2),
+            "height": int(cy * 2),
+            "params": [fx, fy, cx, cy],
+        }
+
+        return params
+
+    def __repr__(self):
+        return f"ImageData(K={self.K}, R={self.R}, t={self.t})"
+
+
+def cameras2F(cam1: Camera, cam2: Camera) -> torch.Tensor:
+    E = cameras2E(cam1, cam2)
+    return cam2.K_inv.T @ E @ cam1.K_inv
+
+
+def cameras2E(cam1: Camera, cam2: Camera) -> torch.Tensor:
+    R = cam2.R @ cam1.R.T
+    T = cam2.t - R @ cam1.t
+    return cross_product_matrix(T) @ R
+
+
+def cross_product_matrix(v) -> torch.Tensor:
+    """Following en.wikipedia.org/wiki/Cross_product#Conversion_to_matrix_multiplication."""
+
+    return torch.tensor(
+        [[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]],
+        dtype=v.dtype,
+        device=v.device,
+    )

ripe/utils/pose_error.py

@@ -0,0 +1,62 @@
+# mostly from: https://github.com/cvg/glue-factory/blob/main/gluefactory/geometry/epipolar.py
+
+import numpy as np
+import torch
+
+
+def angle_error_mat(R1, R2):
+    cos = (torch.trace(torch.einsum("...ij, ...jk -> ...ik", R1.T, R2)) - 1) / 2
+    cos = torch.clip(cos, -1.0, 1.0)  # numerical errors can make it out of bounds
+    return torch.rad2deg(torch.abs(torch.arccos(cos)))
+
+
+def angle_error_vec(v1, v2, eps=1e-10):
+    n = torch.clip(v1.norm(dim=-1) * v2.norm(dim=-1), min=eps)
+    v1v2 = (v1 * v2).sum(dim=-1)  # dot product in the last dimension
+    return torch.rad2deg(torch.arccos(torch.clip(v1v2 / n, -1.0, 1.0)))
+
+
+def relative_pose_error(R_gt, t_gt, R, t, ignore_gt_t_thr=0.0, eps=1e-10):
+    # angle error between 2 vectors
+    t_err = angle_error_vec(t, t_gt, eps)
+    t_err = torch.minimum(t_err, 180 - t_err)  # handle E ambiguity
+    if t_gt.norm() < ignore_gt_t_thr:  # pure rotation is challenging
+        t_err = torch.zeros_like(t_err)
+
+    # angle error between 2 rotation matrices
+    r_err = angle_error_mat(R, R_gt)
+
+    return t_err, r_err
+
+
+def cal_error_auc(errors, thresholds):
+    sort_idx = np.argsort(errors)
+    errors = np.array(errors.copy())[sort_idx]
+    recall = (np.arange(len(errors)) + 1) / len(errors)
+    errors = np.r_[0.0, errors]
+    recall = np.r_[0.0, recall]
+    aucs = []
+    for t in thresholds:
+        last_index = np.searchsorted(errors, t)
+        r = np.r_[recall[:last_index], recall[last_index - 1]]
+        e = np.r_[errors[:last_index], t]
+        aucs.append(np.round((np.trapz(r, x=e) / t), 4))
+    return aucs
+
+
+class AUCMetric:
+    def __init__(self, thresholds, elements=None):
+        self._elements = elements
+        self.thresholds = thresholds
+        if not isinstance(thresholds, list):
+            self.thresholds = [thresholds]
+
+    def update(self, tensor):
+        assert tensor.dim() == 1
+        self._elements += tensor.cpu().numpy().tolist()
+
+    def compute(self):
+        if len(self._elements) == 0:
+            return np.nan
+        else:
+            return cal_error_auc(self._elements, self.thresholds)

ripe/utils/pylogger.py

@@ -0,0 +1,32 @@
+import logging
+
+# from pytorch_lightning.utilities import rank_zero_only
+
+
+def init_base_pylogger():
+    """Initializes base python command line logger."""
+
+    logging.basicConfig(
+        level=logging.WARNING,
+        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+        handlers=[logging.StreamHandler()],
+    )
+
+
+def get_pylogger(name=__name__) -> logging.Logger:
+    """Initializes multi-GPU-friendly python command line logger."""
+
+    if not logging.root.handlers:
+        init_base_pylogger()
+
+    logger = logging.getLogger(name)
+
+    logger.setLevel(logging.DEBUG)
+
+    # this ensures all logging levels get marked with the rank zero decorator
+    # otherwise logs would get multiplied for each GPU process in multi-GPU setup
+    # logging_levels = ("debug", "info", "warning", "error", "exception", "fatal", "critical")
+    # for level in logging_levels:
+    #     setattr(logger, level, rank_zero_only(getattr(logger, level)))
+
+    return logger

ripe/utils/utils.py

@@ -0,0 +1,192 @@
+import random
+from typing import List
+
+import cv2
+import numpy as np
+import torch
+from torchvision.transforms.functional import resize
+
+from ripe import utils
+
+log = utils.get_pylogger(__name__)
+
+
+def gridify(x, window_size):
+    """Turn a tensor of BxCxHxW into a tensor of
+    BxCx(H//window_size)x(W//window_size)x(window_size**2)
+
+    Params:
+        x: Input tensor of shape BxCxHxW
+        window_size: Size of the window
+
+    Returns:
+        x: Output tensor of shape BxCx(H//window_size)x(W//window_size)x(window_size**2)
+    """
+
+    assert x.dim() == 4, "Input tensor x must have 4 dimensions"
+
+    B, C, H, W = x.shape
+    x = (
+        x.unfold(2, window_size, window_size)
+        .unfold(3, window_size, window_size)
+        .reshape(B, C, H // window_size, W // window_size, window_size**2)
+    )
+
+    return x
+
+
+def get_grid(B, H, W, device):
+    x1_n = torch.meshgrid(*[torch.linspace(-1 + 1 / n, 1 - 1 / n, n, device=device) for n in (B, H, W)])
+    x1_n = torch.stack((x1_n[2], x1_n[1]), dim=-1).reshape(B, H * W, 2)
+    return x1_n
+
+
+def cv2_matches_from_kornia(match_dists: torch.Tensor, match_idxs: torch.Tensor) -> List[cv2.DMatch]:
+    return [cv2.DMatch(idx[0].item(), idx[1].item(), d.item()) for idx, d in zip(match_idxs, match_dists)]
+
+
+def to_cv_kpts(kpts, scores):
+    kp = kpts.cpu().numpy().astype(np.int16)
+    s = scores.cpu().numpy()
+
+    cv_kp = [cv2.KeyPoint(kp[i][0], kp[i][1], 6, 0, s[i]) for i in range(len(kp))]
+
+    return cv_kp
+
+
+def resize_image(image, min_size=512, max_size=768):
+    """Resize image to a new size while maintaining the aspect ratio.
+
+    Params:
+        image (torch.tensor): Image to be resized.
+        min_size (int): Minimum size of the smaller dimension.
+        max_size (int): Maximum size of the larger dimension.
+
+    Returns:
+        image: Resized image.
+    """
+
+    h, w = image.shape[-2:]
+
+    aspect_ratio = w / h
+
+    if w > h:
+        new_w = max(min_size, min(max_size, w))
+        new_h = int(new_w / aspect_ratio)
+    else:
+        new_h = max(min_size, min(max_size, h))
+        new_w = int(new_h * aspect_ratio)
+
+    new_size = (new_h, new_w)
+
+    image = resize(image, new_size)
+
+    return image
+
+
+def get_rewards(
+    reward,
+    kps1,
+    kps2,
+    selected_mask1,
+    selected_mask2,
+    padding_mask1,
+    padding_mask2,
+    rel_idx_matches,
+    abs_idx_matches,
+    ransac_inliers,
+    label,
+    penalty=0.0,
+    use_whitening=False,
+    selected_only=False,
+    filter_mode=None,
+):
+    with torch.no_grad():
+        reward *= 1.0 if label else -1.0
+
+        dense_returns = torch.zeros((len(kps1), len(kps2)), device=kps1.device)
+
+        if filter_mode == "ignore":
+            dense_returns[
+                abs_idx_matches[:, 0][ransac_inliers],
+                abs_idx_matches[:, 1][ransac_inliers],
+            ] = reward
+        elif filter_mode == "punish":
+            in_padding_area = (
+                padding_mask1[abs_idx_matches[:, 0]] & padding_mask2[abs_idx_matches[:, 1]]
+            )  # both in the image area (not in padding area)
+
+            dense_returns[
+                abs_idx_matches[:, 0][ransac_inliers & in_padding_area],
+                abs_idx_matches[:, 1][ransac_inliers & in_padding_area],
+            ] = reward
+            dense_returns[
+                abs_idx_matches[:, 0][ransac_inliers & ~in_padding_area],
+                abs_idx_matches[:, 1][ransac_inliers & ~in_padding_area],
+            ] = -1.0
+        else:
+            raise ValueError(f"Unknown filter mode: {filter_mode}")
+
+        if selected_only:
+            dense_returns = dense_returns[selected_mask1, :][:, selected_mask2]
+        if filter_mode == "ignore" and not selected_only:
+            dense_returns = dense_returns[padding_mask1, :][:, padding_mask2]
+
+        if penalty != 0.0:
+            # pos. pair: small penalty for not finding a match
+            # neg. pair: small reward for not finding a match
+            penalty_val = penalty if label else -penalty
+
+            dense_returns[dense_returns == 0.0] = penalty_val
+
+        if use_whitening:
+            dense_returns = (dense_returns - dense_returns.mean()) / (dense_returns.std() + 1e-6)
+
+    return dense_returns
+
+
+def get_other_random_id(idx: int, len_dataset: int, min_dist: int = 20):
+    for _ in range(10):
+        tgt_id = random.randint(0, len_dataset - 1)
+        if abs(idx - tgt_id) >= min_dist:
+            return tgt_id
+
+    raise ValueError(f"Could not find target image with distance >= {min_dist} from source image {idx}")
+
+
+def cv_resize_and_pad_to_shape(image, new_shape, padding_color=(0, 0, 0)):
+    """Resizes image to new_shape with maintaining the aspect ratio and pads with padding_color if
+    needed.
+
+    Params:
+        image: Image to be resized.
+        new_shape: Expected (height, width) of new image.
+        padding_color: Tuple in BGR of padding color
+    Returns:
+        image: Resized image with padding
+    """
+    h, w = image.shape[:2]
+
+    scale_h = new_shape[0] / h
+    scale_w = new_shape[1] / w
+
+    scale = None
+    if scale_w * h > new_shape[0]:
+        scale = scale_h
+    elif scale_h * w > new_shape[1]:
+        scale = scale_w
+    else:
+        scale = max(scale_h, scale_w)
+
+    new_w, new_h = int(round(w * scale)), int(round(h * scale))
+
+    image = cv2.resize(image, (new_w, new_h))
+
+    missing_h = new_shape[0] - new_h
+    missing_w = new_shape[1] - new_w
+
+    top, bottom = missing_h // 2, missing_h - (missing_h // 2)
+    left, right = missing_w // 2, missing_w - (missing_w // 2)
+
+    image = cv2.copyMakeBorder(image, top, bottom, left, right, cv2.BORDER_CONSTANT, value=padding_color)
+    return image

ripe/utils/wandb_utils.py

@@ -0,0 +1,16 @@
+import omegaconf
+
+
+def get_flattened_wandb_cfg(conf_dict):
+    flattened = {}
+
+    def _flatten(cfg, prefix=""):
+        for k, v in cfg.items():
+            new_key = f"{prefix}.{k}" if prefix else k
+            if isinstance(v, omegaconf.dictconfig.DictConfig):
+                _flatten(v, new_key)
+            else:
+                flattened[new_key] = v
+
+    _flatten(conf_dict)
+    return flattened