Add utiities

Author: Phil26AT

gluefactory/geometry/reconstruction.py

@@ -7,11 +7,23 @@
 import logging
 import math
 from pathlib import Path
-from typing import Dict, List, NamedTuple, Optional, Sequence, Tuple, Union
+from typing import (
+    Any,
+    Dict,
+    List,
+    NamedTuple,
+    Optional,
+    Sequence,
+    Tuple,
+    TypeAlias,
+    Union,
+)
 
 import h5py
 import numpy as np
 
+Self: TypeAlias = Any
+
 try:
     import pycolmap
 except ImportError:
@@ -199,6 +211,44 @@ def angular_error(self, other: "Pose", agg="max") -> torch.Tensor:
             "dt": lambda x: x[..., 1],
         }[agg](self.angular_drdt(other, stack=True))
 
+    def _opening_angle(self, return_cos: bool = False) -> float:
+        v0 = torch.zeros_like(self.t)
+        v0[..., -1] = 1.0
+        v1 = self.R[..., 2, :]
+        v01 = self.inv().t
+        v01 = v01 / v01.norm()
+        n = v0.cross(v01, dim=-1)
+        n = n / n.norm()
+        cos = n.dot(v1)
+        return cos if return_cos else (90 - torch.acos(cos) * 180.0 / 3.1415).abs()
+
+    def opening_angle(self, return_cos: bool = False) -> float:
+        return 0.5 * (
+            self._opening_angle(return_cos=return_cos)
+            + self.inv()._opening_angle(return_cos=return_cos)
+        )
+
+    @classmethod
+    def exp(cls, delta: torch.Tensor):
+        dr, dt = delta.split(3, dim=-1)
+        return cls.from_aa(dr, dt)
+
+    def manifold(self, delta: torch.Tensor | None = None) -> torch.Tensor:
+        if delta is None:
+            delta = torch.zeros_like(self._data[..., :6])
+        return delta
+
+    def update(self, delta: torch.Tensor | Self, inplace: bool = False) -> "Pose":
+        if not isinstance(delta, self.__class__):
+            delta = Pose.exp(delta)
+
+        updated_pose = delta @ self
+        if inplace:
+            self._data = updated_pose._data
+            return self
+        else:
+            return updated_pose
+
     def __repr__(self):
         return f"Pose: {self.shape} {self.dtype} {self.device}"
 
@@ -377,13 +427,32 @@ def J_world2image(self, p3d: torch.Tensor):
         J = self.J_denormalize() @ self.J_distort(p2d_dist) @ self.J_project(p3d)
         return J, valid
 
-    @tensor.autocast
-    def image2cam(self, p2d: torch.Tensor) -> torch.Tensor:
+    def image2cam(self, p2d: torch.Tensor, homogeneous: bool = True) -> torch.Tensor:
         """Convert 2D pixel corrdinates to 3D points with z=1"""
         assert self._data.shape
         p2d = self.normalize(p2d)
         # iterative undistortion
-        return gtr.to_homogeneous(p2d)
+        if homogeneous:
+            return gtr.to_homogeneous(p2d)
+        else:
+            return p2d
+
+    def manifold(self, delta: torch.Tensor | None = None) -> torch.Tensor:
+        if delta is None:
+            delta = torch.zeros_like(self._data[..., 2:4])
+        return delta
+
+    def update(self, delta: torch.Tensor | Self, inplace=False) -> "Camera":
+        if isinstance(delta, self.__class__):
+            delta = delta._data[..., 2:4]
+
+        if inplace:
+            self._data[..., 2:4] += delta
+            return self
+        else:
+            cam_copy = self.clone()
+            cam_copy._data[..., 2:4] += delta
+            return cam_copy
 
     def to_cameradict(self, camera_model: Optional[str] = None) -> List[Dict]:
         data = self._data.clone()

gluefactory/geometry/transforms.py

@@ -41,6 +41,14 @@ def batched_eye_like(x: torch.Tensor, n: int):
     return torch.eye(n).to(x)[None].repeat(len(x), 1, 1)
 
 
+def eye_like(tensor: torch.Tensor) -> torch.Tensor:
+    batch_dims = tensor.shape[:-2]
+    identity = torch.eye(*tensor.shape[-2:], dtype=tensor.dtype, device=tensor.device)
+    for _ in batch_dims:
+        identity = identity[None]
+    return identity.repeat(*batch_dims, 1, 1)
+
+
 def skew_symmetric(v):
     """Create a skew-symmetric matrix from a (batched) vector of size (..., 3)."""
     z = torch.zeros_like(v[..., 0])

gluefactory/utils/misc.py

@@ -11,8 +11,8 @@
 
 
 def map_tensor(input_, func):
-    string_classes = (str, bytes)
-    if isinstance(input_, string_classes):
+    scalar_classes = (str, bytes, bool, float, int)
+    if isinstance(input_, scalar_classes):
         return input_
     elif isinstance(input_, Mapping):
         return {k: map_tensor(sample, func) for k, sample in input_.items()}

gluefactory/utils/tensor.py

@@ -125,6 +125,10 @@ def squeeze(self, *args, **kwargs):
         """Create a new tensor of the same type and device."""
         return self.__class__(self._data.squeeze(*args, **kwargs))
 
+    def clone(self, *args, **kwargs):
+        """Create a new tensor of the same type and device."""
+        return self.__class__(self._data.clone(*args, **kwargs))
+
     @classmethod
     def stack(cls, objects: list, dim=0, *, out=None):
         """Stack a list of objects with the same type and shape."""