vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/RoMaV2/src/romav2/romav2.py

@@ -0,0 +1,533 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from collections import OrderedDict
+
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image
+
+import logging
+from romav2.device import device
+from romav2.features import Descriptor, FineFeatures
+from romav2.geometry import (
+    bhwc_grid_sample,
+    bhwc_interpolate,
+    get_normalized_grid,
+    prec_mat_from_prec_params,
+    to_pixel,
+)
+from romav2.io import check_not_i16
+from romav2.matcher import Matcher
+from romav2.refiner import Refiners
+from romav2.types import Setting, ImageLike
+
+logger = logging.getLogger(__name__)
+
+
+def _interpolate_warp_and_confidence(
+    *,
+    warp: torch.Tensor,
+    confidence: torch.Tensor,
+    H: int,
+    W: int,
+    patch_size: int,
+    zero_out_precision: bool,
+):
+    warp = bhwc_interpolate(
+        warp.detach(),
+        size=(H // patch_size, W // patch_size),
+        mode="bilinear",
+        align_corners=False,
+    )
+    if zero_out_precision:
+        # delta at 4 is absolute, and if we
+        # for the second pass we therefore can't use first pred.
+        # overlap is fine since it's relative to matcher pred.
+        confidence[..., 1:] = 0.0
+
+    confidence = bhwc_interpolate(
+        confidence.detach(),
+        size=(H // patch_size, W // patch_size),
+        mode="bilinear",
+        align_corners=False,
+    )
+    return warp, confidence
+
+
+def _map_confidence(*, confidence: torch.Tensor, threshold: float | None):
+    overlap = confidence[..., :1].sigmoid()
+    if threshold is not None:
+        overlap[overlap > threshold] = 1.0
+    precision = prec_mat_from_prec_params(confidence[..., 1:4])
+    return overlap, precision
+
+
+class RoMaV2(nn.Module):
+    @dataclass(frozen=True)
+    class Cfg:
+        descriptor: Descriptor.Cfg = Descriptor.Cfg()
+        matcher: Matcher.Cfg = Matcher.Cfg()
+        refiners: Refiners.Cfg = Refiners.Cfg()
+        refiner_features: FineFeatures.Cfg = FineFeatures.Cfg()
+        anchor_width: int = 512
+        anchor_height: int = 512
+        setting: Setting = "precise"
+        compile: bool = True
+        name: str = "RoMa v2"
+
+    # settings
+    H_lr: int
+    W_lr: int
+    H_hr: int | None
+    W_hr: int | None
+    bidirectional: bool
+    threshold: float | None
+    balanced_sampling: bool
+
+    def __init__(self, cfg: Cfg | None = None):
+        super().__init__()
+        if cfg is None:
+            # default
+            cfg = RoMaV2.Cfg()
+            weights = torch.hub.load_state_dict_from_url(
+                "https://github.com/Parskatt/RoMaV2/releases/download/weights/romav2.pt"
+            )
+        else:
+            weights = None
+        self.f = Descriptor(cfg.descriptor)
+        self.matcher = Matcher(cfg.matcher)
+        self.cfg = cfg
+        self.anchor_width = cfg.anchor_width
+        self.anchor_height = cfg.anchor_height
+        self.refiners = Refiners(cfg.refiners)
+        self.refiner_features = FineFeatures(cfg.refiner_features)
+        self.to(device)
+        self.eval()
+        self.apply_setting(cfg.setting)
+        self.name = cfg.name
+        if weights is not None:
+            self.load_state_dict(weights)
+        if cfg.compile:
+            logger.info(f"Compiling {self.name}...")
+            self.compile()
+        logger.info(f"{self.name} initialized.")
+
+    def apply_setting(self, setting: Setting):
+        if setting in ["mega1500", "scannet1500", "wxbs", "satast"]:
+            self.H_lr = 800
+            self.W_lr = 800
+            self.H_hr = 1024
+            self.W_hr = 1024
+            self.bidirectional = True
+            self.threshold = 0.05
+            self.balanced_sampling = True
+        elif setting == "turbo":
+            self.H_lr = 320
+            self.W_lr = 320
+            self.H_hr = None
+            self.W_hr = None
+            self.bidirectional = False
+            self.threshold = None
+            self.balanced_sampling = True
+        elif setting == "fast":
+            self.H_lr = 512
+            self.W_lr = 512
+            self.H_hr = None
+            self.W_hr = None
+            self.bidirectional = False
+            self.threshold = None
+            self.balanced_sampling = True
+        elif setting == "base":
+            self.H_lr = 640
+            self.W_lr = 640
+            self.H_hr = None
+            self.W_hr = None
+            self.bidirectional = False
+            self.threshold = None
+            self.balanced_sampling = True
+        elif setting == "precise":
+            self.H_lr = 800
+            self.W_lr = 800
+            self.H_hr = 1280
+            self.W_hr = 1280
+            self.bidirectional = True
+            self.threshold = None
+            self.balanced_sampling = True
+        else:
+            raise TypeError(f"Invalid setting: {setting}")
+
+    @torch.inference_mode()
+    def forward(
+        self,
+        img_A_lr: torch.Tensor,
+        img_B_lr: torch.Tensor,
+        img_A_hr: torch.Tensor | None = None,
+        img_B_hr: torch.Tensor | None = None,
+    ) -> dict[str, tuple[torch.Tensor, torch.Tensor] | torch.Tensor]:
+        if torch.get_float32_matmul_precision() != "highest":
+            raise RuntimeError("Float32 matmul precision must be set to highest")
+        assert not self.training, "Currently only inference mode released"
+        # assumes images between [0, 1]
+        # init preds
+        predictions = OrderedDict()
+        # extract feats
+        f_A = self.f(img_A_lr)
+        f_B = self.f(img_B_lr)
+        # match feats
+        matcher_output = self.matcher(
+            f_A, f_B, img_A=img_A_lr, img_B=img_B_lr, bidirectional=self.bidirectional
+        )
+        # return matcher_output
+        predictions["matcher"] = matcher_output
+        warp_AB, confidence_AB = (
+            matcher_output["warp_AB"],
+            matcher_output["confidence_AB"],
+        )
+        if self.bidirectional:
+            warp_BA, confidence_BA = (
+                matcher_output["warp_BA"],
+                matcher_output["confidence_BA"],
+            )
+        else:
+            warp_BA = None
+            confidence_BA = None
+        # refine warp, maybe twice (if hr is available)
+        for stage, (img_A, img_B) in enumerate(
+            zip([img_A_lr, img_A_hr], [img_B_lr, img_B_hr])
+        ):
+            if img_A is None or img_B is None:
+                continue
+            B, C, H, W = img_A.shape
+            scale_factor = torch.tensor(
+                (W / self.anchor_width, H / self.anchor_height), device=device
+            )
+            refiner_features_A = self.refiner_features(img_A)
+            refiner_features_B = self.refiner_features(img_B)
+            for patch_size_str, refiner in self.refiners.items():
+                patch_size = int(patch_size_str)
+                zero_out_precision = (
+                    img_A_hr is not None and patch_size == 4 and stage == 1
+                )
+                warp_AB, confidence_AB = _interpolate_warp_and_confidence(
+                    warp=warp_AB,
+                    confidence=confidence_AB,
+                    H=H,
+                    W=W,
+                    patch_size=patch_size,
+                    zero_out_precision=zero_out_precision,
+                )
+                if self.bidirectional:
+                    warp_BA, confidence_BA = _interpolate_warp_and_confidence(
+                        warp=warp_BA,
+                        confidence=confidence_BA,
+                        H=H,
+                        W=W,
+                        patch_size=patch_size,
+                        zero_out_precision=zero_out_precision,
+                    )
+
+                f_patch_A = refiner_features_A[patch_size]
+                f_patch_B = refiner_features_B[patch_size]
+                refiner_output_AB = refiner(
+                    f_A=f_patch_A,
+                    f_B=f_patch_B,
+                    prev_warp=warp_AB,
+                    prev_confidence=confidence_AB,
+                    scale_factor=scale_factor,
+                )
+                if self.bidirectional:
+                    refiner_output_BA = refiner(
+                        f_A=f_patch_B,
+                        f_B=f_patch_A,
+                        prev_warp=warp_BA,
+                        prev_confidence=confidence_BA,
+                        scale_factor=scale_factor,
+                    )
+                else:
+                    refiner_output_BA = None
+                predictions[f"refiner_{patch_size}_AB"] = refiner_output_AB
+                predictions[f"refiner_{patch_size}_BA"] = refiner_output_BA
+                warp_AB, confidence_AB = (
+                    refiner_output_AB["warp"],
+                    refiner_output_AB["confidence"],
+                )
+                if self.bidirectional:
+                    warp_BA, confidence_BA = (
+                        refiner_output_BA["warp"],
+                        refiner_output_BA["confidence"],
+                    )
+            predictions["warp_AB"] = warp_AB
+            predictions["confidence_AB"] = confidence_AB
+            if self.bidirectional:
+                predictions["warp_BA"] = warp_BA
+                predictions["confidence_BA"] = confidence_BA
+            else:
+                predictions["warp_BA"] = None
+                predictions["confidence_BA"] = None
+        return predictions
+
+    def _load_image(self, img_like: ImageLike) -> torch.Tensor:
+        if isinstance(img_like, str) or isinstance(img_like, Path):
+            img_pil = Image.open(img_like)
+            check_not_i16(img_pil)
+            img_pil = img_pil.convert("RGB")
+            img = torch.from_numpy(np.array(img_pil)).permute(2, 0, 1).to(device)
+        elif isinstance(img_like, Image.Image):
+            img = torch.from_numpy(np.array(img_like)).permute(2, 0, 1).to(device)
+        elif isinstance(img_like, np.ndarray):
+            assert img_like.shape[-1] == 3, (
+                f"Image must have 3 channels, but got shape {img_like.shape=}"
+            )
+            img = torch.from_numpy(img_like).permute(2, 0, 1).to(device)
+        elif isinstance(img_like, torch.Tensor):
+            assert img_like.shape[1] == 3, (
+                f"Image must have 3 channels, but got shape {img_like.shape=}"
+            )
+            img = img_like
+        else:
+            raise ValueError(f"Unsupported image type: {type(img_like)}")
+
+        if img.dtype == torch.uint8:
+            img = img.float() / 255.0
+        if len(img.shape) == 3:
+            img = img[None]
+        return img
+
+    @torch.inference_mode()
+    def match(
+        self,
+        img_like_A: ImageLike,
+        img_like_B: ImageLike,
+    ) -> dict[str, torch.Tensor]:
+        self.eval()
+        img_A = self._load_image(img_like_A)
+        img_B = self._load_image(img_like_B)
+
+        img_A_lr = F.interpolate(
+            img_A,
+            size=(self.H_lr, self.W_lr),
+            mode="bicubic",
+            align_corners=False,
+            antialias=True,
+        )
+        img_B_lr = F.interpolate(
+            img_B,
+            size=(self.H_lr, self.W_lr),
+            mode="bicubic",
+            align_corners=False,
+            antialias=True,
+        )
+        if self.H_hr is not None and self.W_hr is not None:
+            img_A_hr = F.interpolate(
+                img_A,
+                size=(self.H_hr, self.W_hr),
+                mode="bicubic",
+                align_corners=False,
+                antialias=True,
+            )
+            img_B_hr = F.interpolate(
+                img_B,
+                size=(self.H_hr, self.W_hr),
+                mode="bicubic",
+                align_corners=False,
+                antialias=True,
+            )
+        else:
+            img_A_hr = None
+            img_B_hr = None
+
+        preds = self(img_A_lr, img_B_lr, img_A_hr=img_A_hr, img_B_hr=img_B_hr)
+        
+        warp_AB = preds["warp_AB"]
+        confidence_AB = preds["confidence_AB"]
+        warp_BA = preds["warp_BA"]
+        confidence_BA = preds["confidence_BA"]
+        overlap_AB, precision_AB = _map_confidence(
+            confidence=confidence_AB, threshold=self.threshold
+        )
+        if self.bidirectional:
+            overlap_BA, precision_BA = _map_confidence(
+                confidence=confidence_BA, threshold=self.threshold
+            )
+        else:
+            overlap_BA = None
+            precision_BA = None
+
+        preds = {
+            "warp_AB": warp_AB.clone(),
+            "confidence_AB": confidence_AB.clone(),
+            "overlap_AB": overlap_AB.clone(),
+            "precision_AB": precision_AB.clone(),
+            "warp_BA": warp_BA.clone() if warp_BA is not None else None,
+            "confidence_BA": confidence_BA.clone() if confidence_BA is not None else None,
+            "overlap_BA": overlap_BA.clone() if overlap_BA is not None else None,
+            "precision_BA": precision_BA.clone() if precision_BA is not None else None,
+        }
+        return preds
+
+    def sample(self, preds: dict[str, torch.Tensor], num_corresp: int):
+        warp = preds["warp_AB"]
+        confidence_AB = preds["overlap_AB"]
+        precision_AB = preds["precision_AB"] if "precision_AB" in preds else None
+
+        warp = warp[0]
+        confidence_AB = confidence_AB[0].reshape(-1)
+        if precision_AB is not None:
+            precision_AB = precision_AB[0]
+
+        H_A, W_A, two = warp.shape
+        grid = get_normalized_grid(1, H_A, W_A)[0]
+        matches_AB = torch.cat((grid, warp), dim=-1).reshape(-1, 4)
+        if self.bidirectional:
+            confidence_BA = preds["overlap_BA"]
+            warp_BA = preds["warp_BA"]
+
+            precision_BA = (
+                preds["precision_BA"] if "precision_BA" in preds else None
+            )
+            warp_BA = warp_BA[0]
+            confidence_BA = confidence_BA[0]
+            if precision_BA is not None:
+                precision_BA = precision_BA[0]
+
+            if precision_BA is not None and precision_AB is not None:
+                precision_A = bhwc_grid_sample(
+                    precision_BA[None].reshape(1, H_A, W_A, -1),
+                    warp[None],
+                    mode="bilinear",
+                    align_corners=False,
+                ).reshape(H_A, W_A, 2, 2)
+                precision_B = bhwc_grid_sample(
+                    precision_AB[None].reshape(1, H_A, W_A, -1),
+                    warp_BA[None],
+                    mode="bilinear",
+                    align_corners=False,
+                ).reshape(H_A, W_A, 2, 2)
+                precision_fwd = torch.stack(
+                    (precision_A, precision_AB), dim=-3
+                ).reshape(-1, 2, 2, 2)
+                precision_bwd = torch.stack(
+                    (precision_BA, precision_B), dim=-3
+                ).reshape(-1, 2, 2, 2)
+                precision = torch.cat((precision_fwd, precision_bwd), dim=0)
+            else:
+                precision = None
+            # let's hope H_A is equal to H_B
+            grid = get_normalized_grid(1, H_A, W_A)[0]
+            matches_BA = torch.cat((warp_BA, grid), dim=-1).reshape(-1, 4)
+            confidence = torch.cat(
+                (confidence_AB.reshape(-1), confidence_BA.reshape(-1)), dim=0
+            )
+            matches = torch.cat((matches_AB, matches_BA), dim=0)
+        else:
+            matches = matches_AB
+            confidence = confidence_AB.reshape(-1)
+            precision = precision_AB.reshape(-1, 2, 2)
+
+        expansion_factor = 4
+        confidence = confidence * matches.abs().amax(dim=-1).le(1 - 1 / H_A).float()
+        corresp_inds = torch.multinomial(
+            confidence, expansion_factor * num_corresp, replacement=False
+        )
+        sampled_matches = matches[corresp_inds]
+        sampled_confidence = confidence[corresp_inds]
+        if precision is not None:
+            sampled_precision = precision[corresp_inds]
+        else:
+            sampled_precision = None
+        # return sampled_matches, sampled_confidence
+        density = kde(sampled_matches)
+
+        p = 1 / (density + 1)
+        p[density < 10] = (
+            1e-7  # Basically should have at least 10 perfect neighbours, or around 100 ok ones
+        )
+        balanced_samples = torch.multinomial(
+            p, num_samples=min(num_corresp, len(sampled_confidence)), replacement=False
+        )
+        return (
+            sampled_matches[balanced_samples],
+            sampled_confidence[balanced_samples],
+            sampled_precision[balanced_samples][:, 0]
+            if sampled_precision is not None
+            else None,
+            sampled_precision[balanced_samples][:, 1]
+            if sampled_precision is not None
+            else None,
+        )
+
+    @classmethod
+    def prec_map_coordinates(
+        cls, precision: torch.Tensor, *, H_in: int, W_in: int, H_out: int, W_out: int
+    ):
+        W_ratio = W_in / W_out
+        H_ratio = H_in / H_out
+        ratio = torch.tensor([W_ratio, H_ratio], device=precision.device)
+        precision = precision * ratio[None, :, None] * ratio[None, None, :]
+        return precision
+
+    @classmethod
+    def to_pixel_coordinates(
+        cls, warp: torch.Tensor, H_A: int, W_A: int, H_B: int, W_B: int
+    ):
+        return to_pixel(warp[..., :2], H=H_A, W=W_A), to_pixel(
+            warp[..., 2:], H=H_B, W=W_B
+        )
+
+    @classmethod
+    def match_keypoints(
+        cls,
+        x_A: torch.Tensor,
+        x_B: torch.Tensor,
+        warp: torch.Tensor,
+        certainty: torch.Tensor,
+        return_tuple: bool = True,
+        return_inds: bool = False,
+        max_dist: float = 0.005,
+        cert_th: float = 0.0,
+    ):
+        x_A_to_B = F.grid_sample(
+            warp[..., -2:].permute(2, 0, 1)[None],
+            x_A[None, None],
+            align_corners=False,
+            mode="bilinear",
+        )[0, :, 0].mT
+        cert_A_to_B = F.grid_sample(
+            certainty[None, None, ...],
+            x_A[None, None],
+            align_corners=False,
+            mode="bilinear",
+        )[0, 0, 0]
+        D = torch.cdist(x_A_to_B, x_B)
+        inds_A, inds_B = torch.nonzero(
+            (D == D.min(dim=-1, keepdim=True).values)
+            * (D == D.min(dim=-2, keepdim=True).values)
+            * (cert_A_to_B[:, None] > cert_th)
+            * (D < max_dist),
+            as_tuple=True,
+        )
+
+        if return_tuple:
+            if return_inds:
+                return inds_A, inds_B
+            else:
+                return x_A[inds_A], x_B[inds_B]
+        else:
+            if return_inds:
+                return torch.cat((inds_A, inds_B), dim=-1)
+            else:
+                return torch.cat((x_A[inds_A], x_B[inds_B]), dim=-1)
+
+
+def kde(x: torch.Tensor, std: float = 0.1, half: bool = True) -> torch.Tensor:
+    # use a gaussian kernel to estimate density
+    if half:
+        x = x.half()
+    scores = (-(torch.cdist(x, x) ** 2) / (2 * std**2)).exp()
+    density = scores.sum(dim=-1)
+    return density

vismatch/third_party/RoMaV2/src/romav2/types.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+from typing import Callable
+from dataclasses import dataclass
+from pathlib import Path
+import torch
+from typing import Literal
+import numpy as np
+from PIL import Image
+
+
+@dataclass
+class Batch:
+    img_A: torch.Tensor
+    img_B: torch.Tensor
+    depth_A: torch.Tensor
+    depth_B: torch.Tensor
+    K_A: torch.Tensor
+    K_B: torch.Tensor
+    pose_A: torch.Tensor
+    pose_B: torch.Tensor
+    T_AB: torch.Tensor
+    img_A_path: Path
+    img_B_path: Path
+    source: GTSource | list[GTSource]
+    warp_A_to_B: torch.Tensor
+    warp_B_to_A: torch.Tensor
+    mask_A_to_B: torch.Tensor
+    mask_B_to_A: torch.Tensor
+
+    def to(self, device: torch.device) -> Batch:
+        return Batch(
+            img_A=self.img_A.to(device),
+            img_B=self.img_B.to(device),
+            K_A=self.K_A.to(device),
+            K_B=self.K_B.to(device),
+            pose_A=self.pose_A.to(device),
+            pose_B=self.pose_B.to(device),
+            depth_A=self.depth_A.to(device),
+            depth_B=self.depth_B.to(device),
+            warp_A_to_B=self.warp_A_to_B.to(device),
+            warp_B_to_A=self.warp_B_to_A.to(device),
+            T_AB=self.T_AB.to(device),
+            img_A_path=self.img_A_path,
+            img_B_path=self.img_B_path,
+            source=self.source,
+            mask_A_to_B=self.mask_A_to_B.to(device),
+            mask_B_to_A=self.mask_B_to_A.to(device),
+        )
+
+    @classmethod
+    def collate(cls, samples: list[Batch]) -> Batch:
+        keys = samples[0].__dict__.keys()
+        batch = {}
+        for k in keys:
+            if isinstance(samples[0].__dict__[k], torch.Tensor):
+                batch[k] = torch.stack([s.__dict__[k] for s in samples])
+            else:
+                batch[k] = [s.__dict__[k] for s in samples]
+        return Batch(**batch)  # type: ignore[missing-arguments]
+
+
+GTSource = Literal["depth", "warp"]
+HeadType = Literal["dpt-no-pos"]
+SampleMode = Literal["frame_distance", "overlap"]
+ConfidenceMode = Literal["covis", "frame", "positive"]
+NormType = Literal["batch"]
+RefinersType = Literal["roma-4-pow2"]
+MatcherStyle = Literal["romav2"]
+DescriptorName = Literal["dinov3_vitl16", "dinov2_vitl14"]
+Normalizer = Callable[[torch.Tensor], torch.Tensor]
+OptimizerName = Literal["adamw"]
+ImageLike = torch.Tensor | np.ndarray | str | Path | Image.Image
+Setting = Literal[
+    "mega1500", "scannet1500", "wxbs", "satast", "base", "precise", "turbo", "fast"
+]

vismatch/third_party/RoMaV2/src/romav2/vis.py

@@ -0,0 +1,36 @@
+import torch
+import torch.nn.functional as F
+import numpy as np
+from romav2.device import device
+
+
+def vis(img1, img2, warp_AB, warp_BA, certainty_AB, certainty_BA):
+    H, W = warp_AB.shape[1:3]
+    x1 = F.interpolate(
+        (torch.tensor(np.array(img1)) / 255).to(device).permute(2, 0, 1)[None],
+        size=(H, W),
+        mode="bicubic",
+        align_corners=False,
+        antialias=True,
+    )
+    x2 = F.interpolate(
+        (torch.tensor(np.array(img2)) / 255).to(device).permute(2, 0, 1)[None],
+        size=(H, W),
+        mode="bicubic",
+        align_corners=False,
+        antialias=True,
+    )
+    im2_transfer_rgb = F.grid_sample(
+        x2, warp_AB, mode="bilinear", align_corners=False
+    )[0]
+    im1_transfer_rgb = F.grid_sample(
+        x1, warp_BA, mode="bilinear", align_corners=False
+    )[0]
+    warp_im = torch.cat((im2_transfer_rgb, im1_transfer_rgb), dim=2)  # .permute(1,2,0)
+    white_im = torch.ones((H, 2 * W), device=device)
+    certainty = torch.cat((certainty_AB, certainty_BA), dim=2)
+
+    vis_im = certainty[0] * warp_im + (1 - certainty[0]) * white_im
+    x = torch.cat((x1[0], x2[0]), dim=2)
+    vis_im = torch.cat((x, vis_im), dim=1)
+    return vis_im