Source code for kornia.geometry.liegroup.se2

# kornia.geometry.se2 module inspired by Sophus-sympy.
# https://github.com/strasdat/Sophus/blob/master/sympy/sophus/se2.py
from typing import Optional, Tuple, overload

from kornia.core import Module, Parameter, Tensor, concatenate, pad, rand, stack, tensor, where, zeros_like
from kornia.geometry.liegroup._utils import check_se2_omega_shape, check_se2_r_t_shape, check_v_shape
from kornia.geometry.liegroup.so2 import So2
from kornia.testing import KORNIA_CHECK, KORNIA_CHECK_SAME_DEVICES, KORNIA_CHECK_TYPE


[docs]class Se2(Module): r"""Base class to represent the Se2 group. The SE(2) is the group of rigid body transformations about the origin of two-dimensional Euclidean space :math:`R^2` under the operation of composition. See more: Example: >>> so2 = So2.identity(1) >>> t = torch.ones((1, 2)) >>> se2 = Se2(so2, t) >>> se2 rotation: Parameter containing: tensor([1.+0.j], requires_grad=True) translation: Parameter containing: tensor([[1., 1.]], requires_grad=True) """
[docs] def __init__(self, rotation: So2, translation: Tensor) -> None: """Constructor for the base class. Internally represented by a complex number `z` and a translation 2-vector. Args: rotation: So2 group encompassing a rotation. translation: translation vector with the shape of :math:`(B, 2)`. Example: >>> so2 = So2.identity(1) >>> t = torch.ones((1, 2)) >>> se2 = Se2(so2, t) >>> se2 rotation: Parameter containing: tensor([1.+0.j], requires_grad=True) translation: Parameter containing: tensor([[1., 1.]], requires_grad=True) """ super().__init__() KORNIA_CHECK_TYPE(rotation, So2) # TODO change to KORNIA_CHECK_SHAPE once there is multiple shape support check_se2_r_t_shape(rotation, translation) self._rotation = rotation self._translation = Parameter(translation)
[docs] def __repr__(self) -> str: return f"rotation: {self.r}\ntranslation: {self.t}"
def __getitem__(self, idx) -> 'Se2': return Se2(self._rotation[idx], self._translation[idx]) @overload def __mul__(self, right: 'Se2') -> 'Se2': ... @overload def __mul__(self, right: Tensor) -> Tensor: ...
[docs] def __mul__(self, right): """Compose two Se2 transformations. Args: right: the other Se2 transformation. Return: The resulting Se2 transformation. """ so2 = self.so2 t = self.t if isinstance(right, Se2): KORNIA_CHECK_TYPE(right, Se2) _r = so2 * right.so2 _t = t + so2 * right.t return Se2(_r, _t) elif isinstance(right, Tensor): KORNIA_CHECK_TYPE(right, Tensor) # TODO change to KORNIA_CHECK_SHAPE once there is multiple shape support check_se2_r_t_shape(so2, right) return so2 * right + t else: raise TypeError(f"Unsupported type: {type(right)}")
@property def so2(self) -> So2: """Return the underlying rotation(So2).""" return self._rotation @property def r(self) -> So2: """Return the underlying rotation(So2).""" return self._rotation @property def t(self) -> Tensor: """Return the underlying translation vector of shape :math:`(B,2)`.""" return self._translation @property def rotation(self) -> So2: """Return the underlying rotation(So2).""" return self._rotation @property def translation(self) -> Tensor: """Return the underlying translation vector of shape :math:`(B,2)`.""" return self._translation
[docs] @staticmethod def exp(v) -> 'Se2': """Converts elements of lie algebra to elements of lie group. Args: v: vector of shape :math:`(B, 3)`. Example: >>> v = torch.ones((1, 3)) >>> s = Se2.exp(v) >>> s.r Parameter containing: tensor([0.5403+0.8415j], requires_grad=True) >>> s.t Parameter containing: tensor([[0.3818, 1.3012]], requires_grad=True) """ check_v_shape(v) theta = v[..., 2] so2 = So2.exp(theta) z = tensor(0.0, device=v.device, dtype=v.dtype) theta_nonzeros = theta != 0.0 a = where(theta_nonzeros, so2.z.imag / theta, z) b = where(theta_nonzeros, (1.0 - so2.z.real) / theta, z) x = v[..., 0] y = v[..., 1] t = stack((a * x - b * y, b * x + a * y), -1) return Se2(so2, t)
[docs] def log(self) -> Tensor: """Converts elements of lie group to elements of lie algebra. Example: >>> v = torch.ones((1, 3)) >>> s = Se2.exp(v).log() >>> s tensor([[1.0000, 1.0000, 1.0000]], grad_fn=<StackBackward0>) """ theta = self.so2.log() half_theta = 0.5 * theta denom = self.so2.z.real - 1 a = where( denom != 0, -(half_theta * self.so2.z.imag) / denom, tensor(0.0, device=theta.device, dtype=theta.dtype) ) row0 = stack((a, half_theta), -1) row1 = stack((-half_theta, a), -1) V_inv = stack((row0, row1), -2) upsilon = V_inv @ self.t[..., None] return stack((upsilon[..., 0, 0], upsilon[..., 1, 0], theta), -1)
[docs] @staticmethod def hat(v): """Converts elements from vector space to lie algebra. Returns matrix of shape :math:`(B, 3, 3)`. Args: v: vector of shape:math:`(B, 3)`. Example: >>> theta = torch.tensor(3.1415/2) >>> So2.hat(theta) tensor([[0.0000, 1.5707], [1.5707, 0.0000]]) """ # TODO change to KORNIA_CHECK_SHAPE once there is multiple shape support check_v_shape(v) upsilon = stack((v[..., 0], v[..., 1]), -1) theta = v[..., 2] col0 = concatenate((So2.hat(theta), upsilon.unsqueeze(-2)), -2) return pad(col0, (0, 1))
[docs] @staticmethod def vee(omega) -> Tensor: """Converts elements from lie algebra to vector space. Args: omega: 3x3-matrix representing lie algebra of shape :math:`(B, 3, 3)`. Returns: vector of shape :math:`(B, 3)`. Example: >>> v = torch.ones(3) >>> omega_hat = Se2.hat(v) >>> Se2.vee(omega_hat) tensor([1., 1., 1.]) """ # TODO change to KORNIA_CHECK_SHAPE once there is multiple shape support check_se2_omega_shape(omega) upsilon = omega[..., 2, :2] theta = So2.vee(omega[..., :2, :2]) return concatenate((upsilon, theta[..., None]), -1)
[docs] @classmethod def identity(cls, batch_size: Optional[int] = None, device=None, dtype=None) -> 'Se2': """Create a Se2 group representing an identity rotation and zero translation. Args: batch_size: the batch size of the underlying data. Example: >>> s = Se2.identity(1) >>> s.r Parameter containing: tensor([1.+0.j], requires_grad=True) >>> s.t Parameter containing: tensor([[0., 0.]], requires_grad=True) """ t: Tensor = tensor([0.0, 0.0], device=device, dtype=dtype) if batch_size is not None: KORNIA_CHECK(batch_size >= 1, msg="batch_size must be positive") t = t.repeat(batch_size, 1) return cls(So2.identity(batch_size, device, dtype), t)
[docs] def matrix(self) -> Tensor: """Returns the matrix representation of shape :math:`(B, 3, 3)`. Example: >>> s = Se2(So2.identity(1), torch.ones(1, 2)) >>> s.matrix() tensor([[[1., -0., 1.], [0., 1., 1.], [0., 0., 1.]]], grad_fn=<CopySlices>) """ rt = concatenate((self.r.matrix(), self.t[..., None]), -1) rt_3x3 = pad(rt, (0, 0, 0, 1)) # add last row zeros rt_3x3[..., -1, -1] = 1.0 return rt_3x3
[docs] def inverse(self) -> 'Se2': """Returns the inverse transformation. Example: >>> s = Se2(So2.identity(1), torch.ones(1,2)) >>> s_inv = s.inverse() >>> s_inv.r Parameter containing: tensor([1.+0.j], requires_grad=True) >>> s_inv.t Parameter containing: tensor([[-1., -1.]], requires_grad=True) """ r_inv: So2 = self.r.inverse() t_inv: Tensor = r_inv * (-1 * self.t) return Se2(r_inv, t_inv)
[docs] @classmethod def random(cls, batch_size: Optional[int] = None, device=None, dtype=None) -> 'Se2': """Create a Se2 group representing a random transformation. Args: batch_size: the batch size of the underlying data. Example: >>> s = Se2.random() >>> s = Se2.random(batch_size=3) """ r = So2.random(batch_size, device, dtype) shape: Tuple[int, ...] if batch_size is None: shape = (2,) else: KORNIA_CHECK(batch_size >= 1, msg="batch_size must be positive") shape = (batch_size, 2) return cls(r, rand(shape, device=device, dtype=dtype))
[docs] @classmethod def trans(cls, x: Tensor, y: Tensor) -> "Se2": """Construct a translation only Se2 instance. Args: x: the x-axis translation. y: the y-axis translation. """ KORNIA_CHECK(x.shape == y.shape) KORNIA_CHECK_SAME_DEVICES([x, y]) batch_size = x.shape[0] if len(x.shape) > 0 else None rotation = So2.identity(batch_size, x.device, x.dtype) return cls(rotation, stack((x, y), -1))
[docs] @classmethod def trans_x(cls, x: Tensor) -> "Se2": """Construct a x-axis translation. Args: x: the x-axis translation. """ zs = zeros_like(x) return cls.trans(x, zs)
[docs] @classmethod def trans_y(cls, y: Tensor) -> "Se2": """Construct a y-axis translation. Args: y: the y-axis translation. """ zs = zeros_like(y) return cls.trans(zs, y)
[docs] def adjoint(self) -> Tensor: """Returns the adjoint matrix of shape :math:`(B, 3, 3)`. Example: >>> s = Se2.identity() >>> s.adjoint() tensor([[1., -0., 0.], [0., 1., -0.], [0., 0., 1.]], grad_fn=<CopySlices>) """ rt = self.matrix() rt[..., 0:2, 2] = stack((self.t[..., 1], -self.t[..., 0]), -1) return rt