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