Shortcuts

# Source code for rising.utils.affine

import itertools
from math import pi
from typing import Optional, Union

import torch

[docs]def points_to_homogeneous(batch: torch.Tensor) -> torch.Tensor:
"""
Transforms points from cartesian to homogeneous coordinates

Args:
batch: the batch of points to transform. Should be of shape
BATCHSIZE x NUMPOINTS x DIM.

Returns:
torch.Tensor: the batch of points in homogeneous coordinates

"""

[docs]def matrix_to_homogeneous(batch: torch.Tensor) -> torch.Tensor:
"""
Transforms a given transformation matrix to a homogeneous
transformation matrix.

Args:
batch: the batch of matrices to convert [N, dim, dim]

Returns:
torch.Tensor: the converted batch of matrices

"""
if batch.size(-1) == batch.size(-2):
missing = batch.new_zeros(size=(*batch.shape[:-1], 1))
batch = torch.cat([batch, missing], dim=-1)

missing = torch.zeros(
(batch.size(0), *[1 for tmp in batch.shape[1:-1]], batch.size(-1)), device=batch.device, dtype=batch.dtype
)

missing[..., -1] = 1

[docs]def matrix_to_cartesian(batch: torch.Tensor, keep_square: bool = False) -> torch.Tensor:
"""
Transforms a matrix for a homogeneous transformation back to cartesian
coordinates.

Args:
batch: the batch oif matrices to convert back
keep_square: if False: returns a NDIM x NDIM+1 matrix to keep the
translation part
if True: returns a NDIM x NDIM matrix but looses the translation
part. defaults to False.

Returns:
torch.Tensor: the given matrix in cartesian coordinates

"""
batch = batch[:, :-1, ...]
if keep_square:
batch = batch[..., :-1]

return batch

[docs]def points_to_cartesian(batch: torch.Tensor) -> torch.Tensor:
"""
Transforms a batch of points in homogeneous coordinates back to cartesian
coordinates.

Args:
batch: batch of points in homogeneous coordinates. Should be of shape
BATCHSIZE x NUMPOINTS x NDIM+1

Returns:
torch.Tensor: the batch of points in cartesian coordinates

"""

return batch[..., :-1] / batch[..., -1, None]

[docs]def matrix_revert_coordinate_order(batch: torch.Tensor) -> torch.Tensor:
"""
Reverts the coordinate order of a matrix (e.g. from xyz to zyx).

Args:
batch: the batched transformation matrices; Should be of shape
BATCHSIZE x NDIM x NDIM

Returns:
torch.Tensor: the matrix performing the same transformation on vectors with a
reversed coordinate order
"""
batch[:, :-1, :] = batch[:, :-1, :].flip(1).clone()
batch[:, :-1, :-1] = batch[:, :-1, :-1].flip(2).clone()
return batch

[docs]def get_batched_eye(
batchsize: int,
ndim: int,
device: Optional[Union[torch.device, str]] = None,
dtype: Optional[Union[torch.dtype, str]] = None,
) -> torch.Tensor:
"""
Produces a batched matrix containing 1s on the diagonal

Args:
batchsize : int
the batchsize (first dimension)
ndim : int
the dimensionality of the eyes (second and third dimension)
device : torch.device, str, optional
the device to put the resulting tensor to. Defaults to the default
device
dtype : torch.dtype, str, optional
the dtype of the resulting trensor. Defaults to the default dtype

Returns:
torch.Tensor: batched eye matrix

"""

[docs]def deg_to_rad(angles: Union[torch.Tensor, float, int]) -> Union[torch.Tensor, float, int]:
"""

Args:
angles: the (vectorized) angles to convert

Returns:
torch.Tensor: the transformed (vectorized) angles

"""
return angles * pi / 180

[docs]def unit_box(n: int, scale: Optional[torch.Tensor] = None) -> torch.Tensor:
"""
Create a (scaled) version of a unit box

Args:
n: number of dimensions
scale: scaling of each dimension

Returns:
torch.Tensor: scaled unit box
"""
box = torch.tensor([list(i) for i in itertools.product([0, 1], repeat=n)])
if scale is not None:
box = box.to(scale) * scale[None]
return box


© Copyright Copyright (c) 2019-2020, Justus Schock, Michael Baumgartner.. Revision b9cd7e8f.

Versions
latest
stable
v0.3.0
v0.2.2
v0.2.1
v0.2.0post0
v0.2.0 