vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/RIPE/demo.py

@@ -0,0 +1,51 @@
+import cv2
+import kornia.feature as KF
+import kornia.geometry as KG
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from torchvision.io import decode_image
+
+from ripe import vgg_hyper
+from ripe.utils.utils import cv2_matches_from_kornia, resize_image, to_cv_kpts
+
+dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+model = vgg_hyper().to(dev)
+model.eval()
+
+image1 = resize_image(decode_image("assets/all_souls_000013.jpg").float().to(dev) / 255.0)
+image2 = resize_image(decode_image("assets/all_souls_000055.jpg").float().to(dev) / 255.0)
+
+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)
+
+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]]
+
+H, mask = KG.ransac.RANSAC(model_type="fundamental", inl_th=1.0)(matched_pts_1, matched_pts_2)
+matchesMask = mask.int().ravel().tolist()
+
+result_ransac = cv2.drawMatches(
+    (image1.cpu().permute(1, 2, 0).numpy() * 255.0).astype(np.uint8),
+    to_cv_kpts(kpts_1, score_1),
+    (image2.cpu().permute(1, 2, 0).numpy() * 255.0).astype(np.uint8),
+    to_cv_kpts(kpts_2, score_2),
+    cv2_matches_from_kornia(match_dists, match_idxs),
+    None,
+    matchColor=(0, 255, 0),
+    matchesMask=matchesMask,
+    # matchesMask=None, # without RANSAC filtering
+    singlePointColor=(0, 0, 255),
+    flags=cv2.DrawMatchesFlags_DEFAULT,
+)
+
+plt.imshow(result_ransac)
+plt.axis("off")
+plt.tight_layout()
+
+# plt.show()
+plt.savefig("result_ransac.png")

vismatch/third_party/RIPE/ripe/__init__.py

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

vismatch/third_party/RIPE/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

vismatch/third_party/RIPE/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