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 """ return[batch, batch.new_ones((*batch.size()[:-1], 1))], dim=-1)
[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 =[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 return[batch, missing], dim=-2)
[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 """ return torch.eye(ndim, device=device, dtype=dtype).view(1, ndim, ndim).expand(batchsize, -1, -1).clone()
[docs]def deg_to_rad(angles: Union[torch.Tensor, float, int]) -> Union[torch.Tensor, float, int]: """ Converts from degree to radians. 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 = * scale[None] return box

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

Read the Docs v: stable
On Read the Docs
Project Home

Free document hosting provided by Read the Docs.