def get_affine_matrix3d(translations: torch.Tensor, center: torch.Tensor, scale: torch.Tensor, angles: torch.Tensor,
sxy: Optional[torch.Tensor] = None, sxz: Optional[torch.Tensor] = None,
syx: Optional[torch.Tensor] = None, syz: Optional[torch.Tensor] = None,
szx: Optional[torch.Tensor] = None, szy: Optional[torch.Tensor] = None) -> torch.Tensor:
r"""Composes 3d affine matrix from the components.
Args:
translations (torch.Tensor): tensor containing the translation vector (dx,dy,dz) with shape :math:`(B, 3)`.
center (torch.Tensor): tensor containing the center vector (x,y,z) with shape :math:`(B, 3)`.
scale (torch.Tensor): tensor containing the scale factor with shape :math:`(B)`.
angle: (torch.Tensor): angle axis vector containing the rotation angles in degrees in the form
of (rx, ry, rz) with shape :math:`(B, 3)`. Internally it calls Rodrigues to compute
the rotation matrix from axis-angle.
sxy (torch.Tensor, optional): tensor containing the shear factor in the xy-direction with shape :math:`(B)`.
sxz (torch.Tensor, optional): tensor containing the shear factor in the xz-direction with shape :math:`(B)`.
syx (torch.Tensor, optional): tensor containing the shear factor in the yx-direction with shape :math:`(B)`.
syz (torch.Tensor, optional): tensor containing the shear factor in the yz-direction with shape :math:`(B)`.
szx (torch.Tensor, optional): tensor containing the shear factor in the zx-direction with shape :math:`(B)`.
szy (torch.Tensor, optional): tensor containing the shear factor in the zy-direction with shape :math:`(B)`.
Returns:
torch.Tensor: the 3d affine transformation matrix :math:`(B, 4, 4)`.
"""
transform: torch.Tensor = get_projective_transform(center, -angles, scale)
transform[..., 3] += translations # tx/ty/tz
# pad transform to get Bx3x3
transform_h = convert_affinematrix_to_homography3d(transform)
if any([s is not None for s in [sxy, sxz, syx, syz, szx, szy]]):
shear_mat = get_shear_matrix3d(center, sxy, sxz, syx, syz, szx, szy)
transform_h = transform_h @ shear_mat
return transform_h
def warp_affine3d(
src: torch.Tensor,
M: torch.Tensor,
dsize: Tuple[int, int, int],
flags: str = 'bilinear',
padding_mode: str = 'zeros',
align_corners: bool = True,
) -> torch.Tensor:
r"""Apply a projective transformation a to 3d tensor.
.. warning::
This API signature it is experimental and might suffer some changes in the future.
Args:
src : input tensor of shape :math:`(B, C, D, H, W)`.
M: projective transformation matrix of shape :math:`(B, 3, 4)`.
dsize: size of the output image (depth, height, width).
mode: interpolation mode to calculate output values
``'bilinear'`` | ``'nearest'``.
padding_mode: padding mode for outside grid values
``'zeros'`` | ``'border'`` | ``'reflection'``.
align_corners : mode for grid_generation.
Returns:
torch.Tensor: the warped 3d tensor with shape :math:`(B, C, D, H, W)`.
.. note::
This function is often used in conjunction with :func:`get_perspective_transform3d`.
"""
if len(src.shape) != 5:
raise AssertionError(src.shape)
if not (len(M.shape) == 3 and M.shape[-2:] == (3, 4)):
raise AssertionError(M.shape)
if len(dsize) != 3:
raise AssertionError(dsize)
B, C, D, H, W = src.size()
size_src: Tuple[int, int, int] = (D, H, W)
size_out: Tuple[int, int, int] = dsize
M_4x4 = convert_affinematrix_to_homography3d(M) # Bx4x4
# we need to normalize the transformation since grid sample needs -1/1 coordinates
dst_norm_trans_src_norm: torch.Tensor = normalize_homography3d(
M_4x4, size_src, size_out) # Bx4x4
src_norm_trans_dst_norm = _torch_inverse_cast(dst_norm_trans_src_norm)
P_norm: torch.Tensor = src_norm_trans_dst_norm[:, :3] # Bx3x4
# compute meshgrid and apply to input
dsize_out: List[int] = [B, C] + list(size_out)
grid = torch.nn.functional.affine_grid(P_norm,
dsize_out,
align_corners=align_corners)
return torch.nn.functional.grid_sample(src,
grid,
align_corners=align_corners,
mode=flags,
padding_mode=padding_mode)
def get_projective_transform(center: torch.Tensor, angles: torch.Tensor,
scales: torch.Tensor) -> torch.Tensor:
r"""Calculate the projection matrix for a 3D rotation.
.. warning::
This API signature it is experimental and might suffer some changes in the future.
The function computes the projection matrix given the center and angles per axis.
Args:
center: center of the rotation (x,y,z) in the source with shape :math:`(B, 3)`.
angles: angle axis vector containing the rotation angles in degrees in the form
of (rx, ry, rz) with shape :math:`(B, 3)`. Internally it calls Rodrigues to compute
the rotation matrix from axis-angle.
scales: scale factor for x-y-z-directions with shape :math:`(B, 3)`.
Returns:
the projection matrix of 3D rotation with shape :math:`(B, 3, 4)`.
.. note::
This function is often used in conjunction with :func:`warp_affine3d`.
"""
if not (len(center.shape) == 2 and center.shape[-1] == 3):
raise AssertionError(center.shape)
if not (len(angles.shape) == 2 and angles.shape[-1] == 3):
raise AssertionError(angles.shape)
if center.device != angles.device:
raise AssertionError(center.device, angles.device)
if center.dtype != angles.dtype:
raise AssertionError(center.dtype, angles.dtype)
# create rotation matrix
angle_axis_rad: torch.Tensor = K.deg2rad(angles)
rmat: torch.Tensor = K.angle_axis_to_rotation_matrix(
angle_axis_rad) # Bx3x3
scaling_matrix: torch.Tensor = K.eye_like(3, rmat)
scaling_matrix = scaling_matrix * scales.unsqueeze(dim=1)
rmat = rmat @ scaling_matrix.to(rmat)
# define matrix to move forth and back to origin
from_origin_mat = torch.eye(4)[None].repeat(rmat.shape[0], 1,
1).type_as(center) # Bx4x4
from_origin_mat[..., :3, -1] += center
to_origin_mat = from_origin_mat.clone()
to_origin_mat = _torch_inverse_cast(from_origin_mat)
# append translation with zeros
proj_mat = projection_from_Rt(rmat,
torch.zeros_like(center)[..., None]) # Bx3x4
# chain 4x4 transforms
proj_mat = convert_affinematrix_to_homography3d(proj_mat) # Bx4x4
proj_mat = from_origin_mat @ proj_mat @ to_origin_mat
return proj_mat[..., :3, :] # Bx3x4
def warp_projective(src: torch.Tensor,
M: torch.Tensor,
dsize: Tuple[int, int, int],
flags: str = 'bilinear',
padding_mode: str = 'zeros',
align_corners: bool = True) -> torch.Tensor:
r"""Applies a projective transformation a to 3d tensor.
.. warning::
This API signature it is experimental and might suffer some changes in the future.
Args:
src (torch.Tensor): input tensor of shape :math:`(B, C, D, H, W)`.
M (torch.Tensor): projective transformation matrix of shape :math:`(B, 3, 4)`.
dsize (Tuple[int, int, int]): size of the output image (depth, height, width).
mode (str): interpolation mode to calculate output values
'bilinear' | 'nearest'. Default: 'bilinear'.
padding_mode (str): padding mode for outside grid values
'zeros' | 'border' | 'reflection'. Default: 'zeros'.
align_corners (bool): mode for grid_generation. Default: True.
Returns:
torch.Tensor: the warped 3d tensor with shape :math:`(B, C, D, H, W)`.
"""
assert len(src.shape) == 5, src.shape
assert len(M.shape) == 3 and M.shape[-2:] == (3, 4), M.shape
assert len(dsize) == 3, dsize
B, C, D, H, W = src.size()
size_src: Tuple[int, int, int] = (D, H, W)
size_out: Tuple[int, int, int] = dsize
M_4x4 = convert_affinematrix_to_homography3d(M) # Bx4x4
# we need to normalize the transformation since grid sample needs -1/1 coordinates
dst_norm_trans_src_norm: torch.Tensor = normalize_homography3d(
M_4x4, size_src, size_out) # Bx4x4
src_norm_trans_dst_norm = torch.inverse(dst_norm_trans_src_norm)
P_norm: torch.Tensor = src_norm_trans_dst_norm[:, :3] # Bx3x4
# compute meshgrid and apply to input
dsize_out: List[int] = [B, C] + list(size_out)
grid = torch.nn.functional.affine_grid(P_norm,
dsize_out,
align_corners=align_corners)
return torch.nn.functional.grid_sample(src,
grid,
align_corners=align_corners,
mode=flags,
padding_mode=padding_mode)
def get_affine_matrix3d(translations: torch.Tensor,
center: torch.Tensor,
scale: torch.Tensor,
angles: torch.Tensor,
sxy: Optional[torch.Tensor] = None,
sxz: Optional[torch.Tensor] = None,
syx: Optional[torch.Tensor] = None,
syz: Optional[torch.Tensor] = None,
szx: Optional[torch.Tensor] = None,
szy: Optional[torch.Tensor] = None) -> torch.Tensor:
r"""Composes affine matrix Bx4x4 from the components
Returns:
torch.Tensor: params to be passed to the affine transformation.
"""
transform: torch.Tensor = get_projective_transform(center, -angles, scale)
transform[..., 3] += translations # tx/ty/tz
# pad transform to get Bx3x3
transform_h = convert_affinematrix_to_homography3d(transform)
shear_mat = get_shear_matrix3d(center, sxy, sxz, syx, syz, szx, szy)
transform_h = transform_h @ shear_mat
return transform_h
def get_shear_matrix3d(
center: torch.Tensor,
sxy: Optional[torch.Tensor] = None,
sxz: Optional[torch.Tensor] = None,
syx: Optional[torch.Tensor] = None,
syz: Optional[torch.Tensor] = None,
szx: Optional[torch.Tensor] = None,
szy: Optional[torch.Tensor] = None,
):
r"""Composes shear matrix Bx4x4 from the components.
Note: Ordered shearing, shear x-axis then y-axis then z-axis.
.. math::
\begin{bmatrix}
1 & o & r & oy + rz \\
m & p & s & mx + py + sz -y \\
n & q & t & nx + qy + tz -z \\
0 & 0 & 0 & 1 \\
\end{bmatrix}
Where:
m = S_{xy}
n = S_{xz}
o = S_{yx}
p = S_{xy}S_{yx} + 1
q = S_{xz}S_{yx} + S_{yz}
r = S_{zx} + S_{yx}S_{zy}
s = S_{xy}S_{zx} + (S_{xy}S_{yx} + 1)S_{zy}
t = S_{xz}S_{zx} + (S_{xz}S_{yx} + S_{yz})S_{zy} + 1
Params:
center: shearing center coordinates of (x, y, z).
sxy: shearing degree along x axis, towards y plane.
sxz: shearing degree along x axis, towards z plane.
syx: shearing degree along y axis, towards x plane.
syz: shearing degree along y axis, towards z plane.
szx: shearing degree along z axis, towards x plane.
szy: shearing degree along z axis, towards y plane.
Returns:
params to be passed to the affine transformation.
Examples:
>>> rng = torch.manual_seed(0)
>>> sxy, sxz, syx, syz = torch.randn(4, 1)
>>> sxy, sxz, syx, syz
(tensor([1.5410]), tensor([-0.2934]), tensor([-2.1788]), tensor([0.5684]))
>>> center = torch.tensor([[0., 0., 0.]]) # Bx3
>>> get_shear_matrix3d(center, sxy=sxy, sxz=sxz, syx=syx, syz=syz)
tensor([[[ 1.0000, -1.4369, 0.0000, 0.0000],
[-33.5468, 49.2039, 0.0000, 0.0000],
[ 0.3022, -1.0729, 1.0000, 0.0000],
[ 0.0000, 0.0000, 0.0000, 1.0000]]])
.. note::
This function is often used in conjuntion with :func:`warp_perspective3d`.
"""
sxy = torch.tensor([0.0]).repeat(center.size(0)) if sxy is None else sxy
sxz = torch.tensor([0.0]).repeat(center.size(0)) if sxz is None else sxz
syx = torch.tensor([0.0]).repeat(center.size(0)) if syx is None else syx
syz = torch.tensor([0.0]).repeat(center.size(0)) if syz is None else syz
szx = torch.tensor([0.0]).repeat(center.size(0)) if szx is None else szx
szy = torch.tensor([0.0]).repeat(center.size(0)) if szy is None else szy
x, y, z = torch.split(center, 1, dim=-1)
x, y, z = x.view(-1), y.view(-1), z.view(-1)
# Prepare parameters
sxy_tan = torch.tan(sxy) # type: ignore
sxz_tan = torch.tan(sxz) # type: ignore
syx_tan = torch.tan(syx) # type: ignore
syz_tan = torch.tan(syz) # type: ignore
szx_tan = torch.tan(szx) # type: ignore
szy_tan = torch.tan(szy) # type: ignore
# compute translation matrix
m00, m10, m20, m01, m11, m21, m02, m12, m22 = _compute_shear_matrix_3d(
sxy_tan, sxz_tan, syx_tan, syz_tan, szx_tan, szy_tan)
m03 = m01 * y + m02 * z
m13 = m10 * x + m11 * y + m12 * z - y
m23 = m20 * x + m21 * y + m22 * z - z
# shear matrix is implemented with negative values
sxy_tan, sxz_tan, syx_tan, syz_tan, szx_tan, szy_tan = -sxy_tan, -sxz_tan, -syx_tan, -syz_tan, -szx_tan, -szy_tan
m00, m10, m20, m01, m11, m21, m02, m12, m22 = _compute_shear_matrix_3d(
sxy_tan, sxz_tan, syx_tan, syz_tan, szx_tan, szy_tan)
shear_mat = torch.stack(
[m00, m01, m02, m03, m10, m11, m12, m13, m20, m21, m22, m23],
dim=-1).view(-1, 3, 4)
shear_mat = convert_affinematrix_to_homography3d(shear_mat)
return shear_mat
def warp_affine3d(
src: torch.Tensor,
M: torch.Tensor,
dsize: Tuple[int, int, int],
flags: str = 'bilinear',
padding_mode: str = 'zeros',
align_corners: Optional[bool] = None,
) -> torch.Tensor:
r"""Applies a projective transformation a to 3d tensor.
.. warning::
This API signature it is experimental and might suffer some changes in the future.
Args:
src : input tensor of shape :math:`(B, C, D, H, W)`.
M: projective transformation matrix of shape :math:`(B, 3, 4)`.
dsize: size of the output image (depth, height, width).
mode: interpolation mode to calculate output values
``'bilinear'`` | ``'nearest'``.
padding_mode: padding mode for outside grid values
``'zeros'`` | ``'border'`` | ``'reflection'``.
align_corners : mode for grid_generation.
Returns:
torch.Tensor: the warped 3d tensor with shape :math:`(B, C, D, H, W)`.
.. note::
This function is often used in conjuntion with :func:`get_perspective_transform3d`.
"""
assert len(src.shape) == 5, src.shape
assert len(M.shape) == 3 and M.shape[-2:] == (3, 4), M.shape
assert len(dsize) == 3, dsize
B, C, D, H, W = src.size()
# TODO: remove the statement below in kornia v0.6
if align_corners is None:
message: str = (
"The align_corners default value has been changed. By default now is set True "
"in order to match cv2.warpAffine. In case you want to keep your previous "
"behaviour set it to False. This warning will disappear in kornia > v0.6."
)
warnings.warn(message)
# set default value for align corners
align_corners = True
size_src: Tuple[int, int, int] = (D, H, W)
size_out: Tuple[int, int, int] = dsize
M_4x4 = convert_affinematrix_to_homography3d(M) # Bx4x4
# we need to normalize the transformation since grid sample needs -1/1 coordinates
dst_norm_trans_src_norm: torch.Tensor = normalize_homography3d(
M_4x4, size_src, size_out) # Bx4x4
src_norm_trans_dst_norm = _torch_inverse_cast(dst_norm_trans_src_norm)
P_norm: torch.Tensor = src_norm_trans_dst_norm[:, :3] # Bx3x4
# compute meshgrid and apply to input
dsize_out: List[int] = [B, C] + list(size_out)
grid = torch.nn.functional.affine_grid(P_norm,
dsize_out,
align_corners=align_corners)
return torch.nn.functional.grid_sample(src,
grid,
align_corners=align_corners,
mode=flags,
padding_mode=padding_mode)
def get_shear_matrix3d(
center: torch.Tensor,
sxy: Optional[torch.Tensor] = None,
sxz: Optional[torch.Tensor] = None,
syx: Optional[torch.Tensor] = None,
syz: Optional[torch.Tensor] = None,
szx: Optional[torch.Tensor] = None,
szy: Optional[torch.Tensor] = None,
):
r"""Composes shear matrix Bx4x4 from the components.
.. math::
\begin{bmatrix}
1 & o & r & oy + rz \\
m & p & s & mx + py + sz -y \\
n & q & t & nx + qy + tz -z \\
0 & 0 & 0 & 1 \\
\end{bmatrix}
Where:
m = S_{xy}
n = S_{xz}
o = S_{yx}
p = S_{xy}S_{yx} + 1
q = S_{xz}S_{yx} + S_{yz}
r = S_{zx} + S_{yx}S_{zy}
s = S_{xy}S_{zx} + (S_{xy}S_{yx} + 1)S_{zy}
t = S_{xz}S_{zx} + (S_{xz}S_{yx} + S_{yz})S_{zy} + 1
Returns:
torch.Tensor: params to be passed to the affine transformation.
"""
sxy = torch.tensor(0) if sxy is None else sxy
sxz = torch.tensor(0) if sxz is None else sxz
syx = torch.tensor(0) if syx is None else syx
syz = torch.tensor(0) if syz is None else syz
szx = torch.tensor(0) if szx is None else szx
szy = torch.tensor(0) if szy is None else szy
x, y, z = torch.split(center, 1, dim=-1)
x, y, z = x.view(-1), y.view(-1), z.view(-1)
# Prepare parameters
sxy_tan = torch.tan(sxy) # type: ignore
sxz_tan = torch.tan(sxz) # type: ignore
syx_tan = torch.tan(syx) # type: ignore
syz_tan = torch.tan(syz) # type: ignore
szx_tan = torch.tan(szx) # type: ignore
szy_tan = torch.tan(szy) # type: ignore
# compute translation matrix
m00, m10, m20, m01, m11, m21, m02, m12, m22 = _computer_shear_matrix(
sxy_tan, sxz_tan, syx_tan, syz_tan, szx_tan, szy_tan)
m03 = m01 * y + m02 * z
m13 = m10 * x + m11 * y + m12 * z - y
m23 = m20 * x + m21 * y + m22 * z - z
# shear matrix is implemented with negative values
sxy_tan, sxz_tan, syx_tan, syz_tan, szx_tan, szy_tan = \
- sxy_tan, - sxz_tan, - syx_tan, - syz_tan, - szx_tan, - szy_tan
m00, m10, m20, m01, m11, m21, m02, m12, m22 = _computer_shear_matrix(
sxy_tan, sxz_tan, syx_tan, syz_tan, szx_tan, szy_tan)
shear_mat = torch.stack(
[m00, m01, m02, m03, m10, m11, m12, m13, m20, m21, m22, m23],
dim=-1).view(-1, 3, 4)
shear_mat = convert_affinematrix_to_homography3d(shear_mat)
return shear_mat
