vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/gim/gluefactory/geometry/wrappers.py

@@ -0,0 +1,425 @@
+"""
+Convenience classes for an SE3 pose and a pinhole Camera with lens distortion.
+Based on PyTorch tensors: differentiable, batched, with GPU support.
+"""
+
+import functools
+import inspect
+import math
+from typing import Dict, List, NamedTuple, Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from .utils import (
+    J_distort_points,
+    distort_points,
+    skew_symmetric,
+    so3exp_map,
+    to_homogeneous,
+)
+
+
+def autocast(func):
+    """Cast the inputs of a TensorWrapper method to PyTorch tensors
+    if they are numpy arrays. Use the device and dtype of the wrapper.
+    """
+
+    @functools.wraps(func)
+    def wrap(self, *args):
+        device = torch.device("cpu")
+        dtype = None
+        if isinstance(self, TensorWrapper):
+            if self._data is not None:
+                device = self.device
+                dtype = self.dtype
+        elif not inspect.isclass(self) or not issubclass(self, TensorWrapper):
+            raise ValueError(self)
+
+        cast_args = []
+        for arg in args:
+            if isinstance(arg, np.ndarray):
+                arg = torch.from_numpy(arg)
+                arg = arg.to(device=device, dtype=dtype)
+            cast_args.append(arg)
+        return func(self, *cast_args)
+
+    return wrap
+
+
+class TensorWrapper:
+    _data = None
+
+    @autocast
+    def __init__(self, data: torch.Tensor):
+        self._data = data
+
+    @property
+    def shape(self):
+        return self._data.shape[:-1]
+
+    @property
+    def device(self):
+        return self._data.device
+
+    @property
+    def dtype(self):
+        return self._data.dtype
+
+    def __getitem__(self, index):
+        return self.__class__(self._data[index])
+
+    def __setitem__(self, index, item):
+        self._data[index] = item.data
+
+    def to(self, *args, **kwargs):
+        return self.__class__(self._data.to(*args, **kwargs))
+
+    def cpu(self):
+        return self.__class__(self._data.cpu())
+
+    def cuda(self):
+        return self.__class__(self._data.cuda())
+
+    def pin_memory(self):
+        return self.__class__(self._data.pin_memory())
+
+    def float(self):
+        return self.__class__(self._data.float())
+
+    def double(self):
+        return self.__class__(self._data.double())
+
+    def detach(self):
+        return self.__class__(self._data.detach())
+
+    @classmethod
+    def stack(cls, objects: List, dim=0, *, out=None):
+        data = torch.stack([obj._data for obj in objects], dim=dim, out=out)
+        return cls(data)
+
+    @classmethod
+    def __torch_function__(self, func, types, args=(), kwargs=None):
+        if kwargs is None:
+            kwargs = {}
+        if func is torch.stack:
+            return self.stack(*args, **kwargs)
+        else:
+            return NotImplemented
+
+
+class Pose(TensorWrapper):
+    def __init__(self, data: torch.Tensor):
+        assert data.shape[-1] == 12
+        super().__init__(data)
+
+    @classmethod
+    @autocast
+    def from_Rt(cls, R: torch.Tensor, t: torch.Tensor):
+        """Pose from a rotation matrix and translation vector.
+        Accepts numpy arrays or PyTorch tensors.
+
+        Args:
+            R: rotation matrix with shape (..., 3, 3).
+            t: translation vector with shape (..., 3).
+        """
+        assert R.shape[-2:] == (3, 3)
+        assert t.shape[-1] == 3
+        assert R.shape[:-2] == t.shape[:-1]
+        data = torch.cat([R.flatten(start_dim=-2), t], -1)
+        return cls(data)
+
+    @classmethod
+    @autocast
+    def from_aa(cls, aa: torch.Tensor, t: torch.Tensor):
+        """Pose from an axis-angle rotation vector and translation vector.
+        Accepts numpy arrays or PyTorch tensors.
+
+        Args:
+            aa: axis-angle rotation vector with shape (..., 3).
+            t: translation vector with shape (..., 3).
+        """
+        assert aa.shape[-1] == 3
+        assert t.shape[-1] == 3
+        assert aa.shape[:-1] == t.shape[:-1]
+        return cls.from_Rt(so3exp_map(aa), t)
+
+    @classmethod
+    def from_4x4mat(cls, T: torch.Tensor):
+        """Pose from an SE(3) transformation matrix.
+        Args:
+            T: transformation matrix with shape (..., 4, 4).
+        """
+        assert T.shape[-2:] == (4, 4)
+        R, t = T[..., :3, :3], T[..., :3, 3]
+        return cls.from_Rt(R, t)
+
+    @classmethod
+    def from_colmap(cls, image: NamedTuple):
+        """Pose from a COLMAP Image."""
+        return cls.from_Rt(image.qvec2rotmat(), image.tvec)
+
+    @property
+    def R(self) -> torch.Tensor:
+        """Underlying rotation matrix with shape (..., 3, 3)."""
+        rvec = self._data[..., :9]
+        return rvec.reshape(rvec.shape[:-1] + (3, 3))
+
+    @property
+    def t(self) -> torch.Tensor:
+        """Underlying translation vector with shape (..., 3)."""
+        return self._data[..., -3:]
+
+    def inv(self) -> "Pose":
+        """Invert an SE(3) pose."""
+        R = self.R.transpose(-1, -2)
+        t = -(R @ self.t.unsqueeze(-1)).squeeze(-1)
+        return self.__class__.from_Rt(R, t)
+
+    def compose(self, other: "Pose") -> "Pose":
+        """Chain two SE(3) poses: T_B2C.compose(T_A2B) -> T_A2C."""
+        R = self.R @ other.R
+        t = self.t + (self.R @ other.t.unsqueeze(-1)).squeeze(-1)
+        return self.__class__.from_Rt(R, t)
+
+    @autocast
+    def transform(self, p3d: torch.Tensor) -> torch.Tensor:
+        """Transform a set of 3D points.
+        Args:
+            p3d: 3D points, numpy array or PyTorch tensor with shape (..., 3).
+        """
+        assert p3d.shape[-1] == 3
+        # assert p3d.shape[:-2] == self.shape  # allow broadcasting
+        return p3d @ self.R.transpose(-1, -2) + self.t.unsqueeze(-2)
+
+    def __mul__(self, p3D: torch.Tensor) -> torch.Tensor:
+        """Transform a set of 3D points: T_A2B * p3D_A -> p3D_B."""
+        return self.transform(p3D)
+
+    def __matmul__(
+        self, other: Union["Pose", torch.Tensor]
+    ) -> Union["Pose", torch.Tensor]:
+        """Transform a set of 3D points: T_A2B * p3D_A -> p3D_B.
+        or chain two SE(3) poses: T_B2C @ T_A2B -> T_A2C."""
+        if isinstance(other, self.__class__):
+            return self.compose(other)
+        else:
+            return self.transform(other)
+
+    @autocast
+    def J_transform(self, p3d_out: torch.Tensor):
+        # [[1,0,0,0,-pz,py],
+        #  [0,1,0,pz,0,-px],
+        #  [0,0,1,-py,px,0]]
+        J_t = torch.diag_embed(torch.ones_like(p3d_out))
+        J_rot = -skew_symmetric(p3d_out)
+        J = torch.cat([J_t, J_rot], dim=-1)
+        return J  # N x 3 x 6
+
+    def numpy(self) -> Tuple[np.ndarray]:
+        return self.R.numpy(), self.t.numpy()
+
+    def magnitude(self) -> Tuple[torch.Tensor]:
+        """Magnitude of the SE(3) transformation.
+        Returns:
+            dr: rotation anngle in degrees.
+            dt: translation distance in meters.
+        """
+        trace = torch.diagonal(self.R, dim1=-1, dim2=-2).sum(-1)
+        cos = torch.clamp((trace - 1) / 2, -1, 1)
+        dr = torch.acos(cos).abs() / math.pi * 180
+        dt = torch.norm(self.t, dim=-1)
+        return dr, dt
+
+    def __repr__(self):
+        return f"Pose: {self.shape} {self.dtype} {self.device}"
+
+
+class Camera(TensorWrapper):
+    eps = 1e-4
+
+    def __init__(self, data: torch.Tensor):
+        assert data.shape[-1] in {6, 8, 10}
+        super().__init__(data)
+
+    @classmethod
+    def from_colmap(cls, camera: Union[Dict, NamedTuple]):
+        """Camera from a COLMAP Camera tuple or dictionary.
+        We use the corner-convetion from COLMAP (center of top left pixel is (0.5, 0.5))
+        """
+        if isinstance(camera, tuple):
+            camera = camera._asdict()
+
+        model = camera["model"]
+        params = camera["params"]
+
+        if model in ["OPENCV", "PINHOLE", "RADIAL"]:
+            (fx, fy, cx, cy), params = np.split(params, [4])
+        elif model in ["SIMPLE_PINHOLE", "SIMPLE_RADIAL"]:
+            (f, cx, cy), params = np.split(params, [3])
+            fx = fy = f
+            if model == "SIMPLE_RADIAL":
+                params = np.r_[params, 0.0]
+        else:
+            raise NotImplementedError(model)
+
+        data = np.r_[camera["width"], camera["height"], fx, fy, cx, cy, params]
+        return cls(data)
+
+    @classmethod
+    @autocast
+    def from_calibration_matrix(cls, K: torch.Tensor):
+        cx, cy = K[..., 0, 2], K[..., 1, 2]
+        fx, fy = K[..., 0, 0], K[..., 1, 1]
+        data = torch.stack([2 * cx, 2 * cy, fx, fy, cx, cy], -1)
+        return cls(data)
+
+    @autocast
+    def calibration_matrix(self):
+        K = torch.zeros(
+            *self._data.shape[:-1],
+            3,
+            3,
+            device=self._data.device,
+            dtype=self._data.dtype,
+        )
+        K[..., 0, 2] = self._data[..., 4]
+        K[..., 1, 2] = self._data[..., 5]
+        K[..., 0, 0] = self._data[..., 2]
+        K[..., 1, 1] = self._data[..., 3]
+        K[..., 2, 2] = 1.0
+        return K
+
+    @property
+    def size(self) -> torch.Tensor:
+        """Size (width height) of the images, with shape (..., 2)."""
+        return self._data[..., :2]
+
+    @property
+    def f(self) -> torch.Tensor:
+        """Focal lengths (fx, fy) with shape (..., 2)."""
+        return self._data[..., 2:4]
+
+    @property
+    def c(self) -> torch.Tensor:
+        """Principal points (cx, cy) with shape (..., 2)."""
+        return self._data[..., 4:6]
+
+    @property
+    def dist(self) -> torch.Tensor:
+        """Distortion parameters, with shape (..., {0, 2, 4})."""
+        return self._data[..., 6:]
+
+    @autocast
+    def scale(self, scales: torch.Tensor):
+        """Update the camera parameters after resizing an image."""
+        s = scales
+        data = torch.cat([self.size * s, self.f * s, self.c * s, self.dist], -1)
+        return self.__class__(data)
+
+    def crop(self, left_top: Tuple[float], size: Tuple[int]):
+        """Update the camera parameters after cropping an image."""
+        left_top = self._data.new_tensor(left_top)
+        size = self._data.new_tensor(size)
+        data = torch.cat([size, self.f, self.c - left_top, self.dist], -1)
+        return self.__class__(data)
+
+    @autocast
+    def in_image(self, p2d: torch.Tensor):
+        """Check if 2D points are within the image boundaries."""
+        assert p2d.shape[-1] == 2
+        # assert p2d.shape[:-2] == self.shape  # allow broadcasting
+        size = self.size.unsqueeze(-2)
+        valid = torch.all((p2d >= 0) & (p2d <= (size - 1)), -1)
+        return valid
+
+    @autocast
+    def project(self, p3d: torch.Tensor) -> Tuple[torch.Tensor]:
+        """Project 3D points into the camera plane and check for visibility."""
+        z = p3d[..., -1]
+        valid = z > self.eps
+        z = z.clamp(min=self.eps)
+        p2d = p3d[..., :-1] / z.unsqueeze(-1)
+        return p2d, valid
+
+    def J_project(self, p3d: torch.Tensor):
+        x, y, z = p3d[..., 0], p3d[..., 1], p3d[..., 2]
+        zero = torch.zeros_like(z)
+        z = z.clamp(min=self.eps)
+        J = torch.stack([1 / z, zero, -x / z**2, zero, 1 / z, -y / z**2], dim=-1)
+        J = J.reshape(p3d.shape[:-1] + (2, 3))
+        return J  # N x 2 x 3
+
+    @autocast
+    def distort(self, pts: torch.Tensor) -> Tuple[torch.Tensor]:
+        """Distort normalized 2D coordinates
+        and check for validity of the distortion model.
+        """
+        assert pts.shape[-1] == 2
+        # assert pts.shape[:-2] == self.shape  # allow broadcasting
+        return distort_points(pts, self.dist)
+
+    def J_distort(self, pts: torch.Tensor):
+        return J_distort_points(pts, self.dist)  # N x 2 x 2
+
+    @autocast
+    def denormalize(self, p2d: torch.Tensor) -> torch.Tensor:
+        """Convert normalized 2D coordinates into pixel coordinates."""
+        return p2d * self.f.unsqueeze(-2) + self.c.unsqueeze(-2)
+
+    @autocast
+    def normalize(self, p2d: torch.Tensor) -> torch.Tensor:
+        """Convert normalized 2D coordinates into pixel coordinates."""
+        return (p2d - self.c.unsqueeze(-2)) / self.f.unsqueeze(-2)
+
+    def J_denormalize(self):
+        return torch.diag_embed(self.f).unsqueeze(-3)  # 1 x 2 x 2
+
+    @autocast
+    def cam2image(self, p3d: torch.Tensor) -> Tuple[torch.Tensor]:
+        """Transform 3D points into 2D pixel coordinates."""
+        p2d, visible = self.project(p3d)
+        p2d, mask = self.distort(p2d)
+        p2d = self.denormalize(p2d)
+        valid = visible & mask & self.in_image(p2d)
+        return p2d, valid
+
+    def J_world2image(self, p3d: torch.Tensor):
+        p2d_dist, valid = self.project(p3d)
+        J = self.J_denormalize() @ self.J_distort(p2d_dist) @ self.J_project(p3d)
+        return J, valid
+
+    @autocast
+    def image2cam(self, p2d: torch.Tensor) -> torch.Tensor:
+        """Convert 2D pixel corrdinates to 3D points with z=1"""
+        assert self._data.shape
+        p2d = self.normalize(p2d)
+        # iterative undistortion
+        return to_homogeneous(p2d)
+
+    def to_cameradict(self, camera_model: Optional[str] = None) -> List[Dict]:
+        data = self._data.clone()
+        if data.dim() == 1:
+            data = data.unsqueeze(0)
+        assert data.dim() == 2
+        b, d = data.shape
+        if camera_model is None:
+            camera_model = {6: "PINHOLE", 8: "RADIAL", 10: "OPENCV"}[d]
+        cameras = []
+        for i in range(b):
+            if camera_model.startswith("SIMPLE_"):
+                params = [x.item() for x in data[i, 3 : min(d, 7)]]
+            else:
+                params = [x.item() for x in data[i, 2:]]
+            cameras.append(
+                {
+                    "model": camera_model,
+                    "width": int(data[i, 0].item()),
+                    "height": int(data[i, 1].item()),
+                    "params": params,
+                }
+            )
+        return cameras if self._data.dim() == 2 else cameras[0]
+
+    def __repr__(self):
+        return f"Camera {self.shape} {self.dtype} {self.device}"

vismatch/third_party/gim/gluefactory/models/__init__.py

@@ -0,0 +1,30 @@
+import importlib.util
+
+from ..utils.tools import get_class
+from .base_model import BaseModel
+
+
+def get_model(name):
+    import_paths = [
+        name,
+        f"{__name__}.{name}",
+        f"{__name__}.extractors.{name}",  # backward compatibility
+        f"{__name__}.matchers.{name}",  # backward compatibility
+    ]
+    for path in import_paths:
+        try:
+            spec = importlib.util.find_spec(path)
+        except ModuleNotFoundError:
+            spec = None
+        if spec is not None:
+            try:
+                return get_class(path, BaseModel)
+            except AssertionError:
+                mod = __import__(path, fromlist=[""])
+                try:
+                    return mod.__main_model__
+                except AttributeError as exc:
+                    print(exc)
+                    continue
+
+    raise RuntimeError(f'Model {name} not found in any of [{" ".join(import_paths)}]')

vismatch/third_party/gim/gluefactory/models/backbones/dinov2.py

@@ -0,0 +1,30 @@
+import torch
+import torch.nn.functional as F
+
+from ..base_model import BaseModel
+
+
+class DinoV2(BaseModel):
+    default_conf = {"weights": "dinov2_vits14", "allow_resize": False}
+    required_data_keys = ["image"]
+
+    def _init(self, conf):
+        self.net = torch.hub.load("facebookresearch/dinov2", conf.weights)
+        self.set_initialized()
+
+    def _forward(self, data):
+        img = data["image"]
+        if self.conf.allow_resize:
+            img = F.upsample(img, [int(x // 14 * 14) for x in img.shape[-2:]])
+        desc, cls_token = self.net.get_intermediate_layers(
+            img, n=1, return_class_token=True, reshape=True
+        )[0]
+
+        return {
+            "features": desc,
+            "global_descriptor": cls_token,
+            "descriptors": desc.flatten(-2).transpose(-2, -1),
+        }
+
+    def loss(self, pred, data):
+        raise NotImplementedError

vismatch/third_party/gim/gluefactory/models/base_model.py

@@ -0,0 +1,157 @@
+"""
+Base class for trainable models.
+"""
+
+from abc import ABCMeta, abstractmethod
+from copy import copy
+
+import omegaconf
+from omegaconf import OmegaConf
+from torch import nn
+
+
+class MetaModel(ABCMeta):
+    def __prepare__(name, bases, **kwds):
+        total_conf = OmegaConf.create()
+        for base in bases:
+            for key in ("base_default_conf", "default_conf"):
+                update = getattr(base, key, {})
+                if isinstance(update, dict):
+                    update = OmegaConf.create(update)
+                total_conf = OmegaConf.merge(total_conf, update)
+        return dict(base_default_conf=total_conf)
+
+
+class BaseModel(nn.Module, metaclass=MetaModel):
+    """
+    What the child model is expect to declare:
+        default_conf: dictionary of the default configuration of the model.
+        It recursively updates the default_conf of all parent classes, and
+        it is updated by the user-provided configuration passed to __init__.
+        Configurations can be nested.
+
+        required_data_keys: list of expected keys in the input data dictionary.
+
+        strict_conf (optional): boolean. If false, BaseModel does not raise
+        an error when the user provides an unknown configuration entry.
+
+        _init(self, conf): initialization method, where conf is the final
+        configuration object (also accessible with `self.conf`). Accessing
+        unknown configuration entries will raise an error.
+
+        _forward(self, data): method that returns a dictionary of batched
+        prediction tensors based on a dictionary of batched input data tensors.
+
+        loss(self, pred, data): method that returns a dictionary of losses,
+        computed from model predictions and input data. Each loss is a batch
+        of scalars, i.e. a torch.Tensor of shape (B,).
+        The total loss to be optimized has the key `'total'`.
+
+        metrics(self, pred, data): method that returns a dictionary of metrics,
+        each as a batch of scalars.
+    """
+
+    default_conf = {
+        "name": None,
+        "trainable": True,  # if false: do not optimize this model parameters
+        "freeze_batch_normalization": False,  # use test-time statistics
+        "timeit": False,  # time forward pass
+    }
+    required_data_keys = []
+    strict_conf = False
+
+    are_weights_initialized = False
+
+    def __init__(self, conf):
+        """Perform some logic and call the _init method of the child model."""
+        super().__init__()
+        default_conf = OmegaConf.merge(
+            self.base_default_conf, OmegaConf.create(self.default_conf)
+        )
+        if self.strict_conf:
+            OmegaConf.set_struct(default_conf, True)
+
+        # fixme: backward compatibility
+        if "pad" in conf and "pad" not in default_conf:  # backward compat.
+            with omegaconf.read_write(conf):
+                with omegaconf.open_dict(conf):
+                    conf["interpolation"] = {"pad": conf.pop("pad")}
+
+        if isinstance(conf, dict):
+            conf = OmegaConf.create(conf)
+        self.conf = conf = OmegaConf.merge(default_conf, conf)
+        OmegaConf.set_readonly(conf, True)
+        OmegaConf.set_struct(conf, True)
+        self.required_data_keys = copy(self.required_data_keys)
+        self._init(conf)
+
+        if not conf.trainable:
+            for p in self.parameters():
+                p.requires_grad = False
+
+    def train(self, mode=True):
+        super().train(mode)
+
+        def freeze_bn(module):
+            if isinstance(module, nn.modules.batchnorm._BatchNorm):
+                module.eval()
+
+        if self.conf.freeze_batch_normalization:
+            self.apply(freeze_bn)
+
+        return self
+
+    def forward(self, data):
+        """Check the data and call the _forward method of the child model."""
+
+        def recursive_key_check(expected, given):
+            for key in expected:
+                assert key in given, f"Missing key {key} in data"
+                if isinstance(expected, dict):
+                    recursive_key_check(expected[key], given[key])
+
+        recursive_key_check(self.required_data_keys, data)
+        return self._forward(data)
+
+    @abstractmethod
+    def _init(self, conf):
+        """To be implemented by the child class."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _forward(self, data):
+        """To be implemented by the child class."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def loss(self, pred, data):
+        """To be implemented by the child class."""
+        raise NotImplementedError
+
+    def load_state_dict(self, *args, **kwargs):
+        """Load the state dict of the model, and set the model to initialized."""
+        ret = super().load_state_dict(*args, **kwargs)
+        self.set_initialized()
+        return ret
+
+    def is_initialized(self):
+        """Recursively check if the model is initialized, i.e. weights are loaded"""
+        is_initialized = True  # initialize to true and perform recursive and
+        for _, w in self.named_children():
+            if isinstance(w, BaseModel):
+                # if children is BaseModel, we perform recursive check
+                is_initialized = is_initialized and w.is_initialized()
+            else:
+                # else, we check if self is initialized or the children has no params
+                n_params = len(list(w.parameters()))
+                is_initialized = is_initialized and (
+                    n_params == 0 or self.are_weights_initialized
+                )
+        return is_initialized
+
+    def set_initialized(self, to: bool = True):
+        """Recursively set the initialization state."""
+        self.are_weights_initialized = to
+        for _, w in self.named_parameters():
+            if isinstance(w, BaseModel):
+                w.set_initialized(to)