Source code for kornia.geometry.transform.projwarp

"""Module to perform projective transformations to tensors."""
from typing import Tuple, List, Optional
import warnings

import torch

import kornia as K
from kornia.geometry.conversions import convert_affinematrix_to_homography3d
from kornia.geometry.transform.homography_warper import normalize_homography3d, homography_warp3d
from kornia.testing import check_is_tensor
from kornia.utils.helpers import _torch_solve_cast, _torch_inverse_cast

__all__ = [
    "warp_affine3d",
    "get_projective_transform",
    "projection_from_Rt",
    "get_perspective_transform3d",
    "warp_perspective3d"
]


[docs]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 (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)`. .. 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)
[docs]def projection_from_Rt(rmat: torch.Tensor, tvec: torch.Tensor) -> torch.Tensor: r"""Compute the projection matrix from Rotation and translation. .. warning:: This API signature it is experimental and might suffer some changes in the future. Concatenates the batch of rotations and translations such that :math:`P = [R | t]`. Args: rmat (torch.Tensor): the rotation matrix with shape :math:`(*, 3, 3)`. tvec (torch.Tensor): the translation vector with shape :math:`(*, 3, 1)`. Returns: torch.Tensor: the projection matrix with shape :math:`(*, 3, 4)`. """ assert len(rmat.shape) >= 2 and rmat.shape[-2:] == (3, 3), rmat.shape assert len(tvec.shape) >= 2 and tvec.shape[-2:] == (3, 1), tvec.shape return torch.cat([rmat, tvec], dim=-1) # Bx3x4
[docs]def get_projective_transform(center: torch.Tensor, angles: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: r"""Calculates 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 (torch.Tensor): center of the rotation (x,y,z) in the source with shape :math:`(B, 3)`. angles (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. scales (torch.Tensor): scale factor for x-y-z-directions with shape :math:`(B, 3)`. Returns: torch.Tensor: the projection matrix of 3D rotation with shape :math:`(B, 3, 4)`. .. note:: This function is often used in conjuntion with :func:`warp_affine3d`. """ assert len(center.shape) == 2 and center.shape[-1] == 3, center.shape assert len(angles.shape) == 2 and angles.shape[-1] == 3, angles.shape assert center.device == angles.device, (center.device, angles.device) assert center.dtype == angles.dtype, (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 tranlation 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
[docs]def get_perspective_transform3d(src: torch.Tensor, dst: torch.Tensor) -> torch.Tensor: r"""Calculate a 3d perspective transform from four pairs of the corresponding points. The function calculates the matrix of a perspective transform so that: .. math :: \begin{bmatrix} t_{i}x_{i}^{'} \\ t_{i}y_{i}^{'} \\ t_{i}z_{i}^{'} \\ t_{i} \\ \end{bmatrix} = \textbf{map_matrix} \cdot \begin{bmatrix} x_{i} \\ y_{i} \\ z_{i} \\ 1 \\ \end{bmatrix} where .. math :: dst(i) = (x_{i}^{'},y_{i}^{'},z_{i}^{'}), src(i) = (x_{i}, y_{i}, z_{i}), i = 0,1,2,5,7 Concrete math is as below: .. math :: \[ u_i =\frac{c_{00} * x_i + c_{01} * y_i + c_{02} * z_i + c_{03}} {c_{30} * x_i + c_{31} * y_i + c_{32} * z_i + c_{33}} \] \[ v_i =\frac{c_{10} * x_i + c_{11} * y_i + c_{12} * z_i + c_{13}} {c_{30} * x_i + c_{31} * y_i + c_{32} * z_i + c_{33}} \] \[ w_i =\frac{c_{20} * x_i + c_{21} * y_i + c_{22} * z_i + c_{23}} {c_{30} * x_i + c_{31} * y_i + c_{32} * z_i + c_{33}} \] .. math :: \begin{pmatrix} x_0 & y_0 & z_0 & 1 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & -x_0*u_0 & -y_0*u_0 & -z_0 * u_0 \\ x_1 & y_1 & z_1 & 1 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & -x_1*u_1 & -y_1*u_1 & -z_1 * u_1 \\ x_2 & y_2 & z_2 & 1 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & -x_2*u_2 & -y_2*u_2 & -z_2 * u_2 \\ x_5 & y_5 & z_5 & 1 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & -x_5*u_5 & -y_5*u_5 & -z_5 * u_5 \\ x_7 & y_7 & z_7 & 1 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & -x_7*u_7 & -y_7*u_7 & -z_7 * u_7 \\ 0 & 0 & 0 & 0 & x_0 & y_0 & z_0 & 1 & 0 & 0 & 0 & 0 & -x_0*v_0 & -y_0*v_0 & -z_0 * v_0 \\ 0 & 0 & 0 & 0 & x_1 & y_1 & z_1 & 1 & 0 & 0 & 0 & 0 & -x_1*v_1 & -y_1*v_1 & -z_1 * v_1 \\ 0 & 0 & 0 & 0 & x_2 & y_2 & z_2 & 1 & 0 & 0 & 0 & 0 & -x_2*v_2 & -y_2*v_2 & -z_2 * v_2 \\ 0 & 0 & 0 & 0 & x_5 & y_5 & z_5 & 1 & 0 & 0 & 0 & 0 & -x_5*v_5 & -y_5*v_5 & -z_5 * v_5 \\ 0 & 0 & 0 & 0 & x_7 & y_7 & z_7 & 1 & 0 & 0 & 0 & 0 & -x_7*v_7 & -y_7*v_7 & -z_7 * v_7 \\ 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & x_0 & y_0 & z_0 & 1 & -x_0*w_0 & -y_0*w_0 & -z_0 * w_0 \\ 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & x_1 & y_1 & z_1 & 1 & -x_1*w_1 & -y_1*w_1 & -z_1 * w_1 \\ 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & x_2 & y_2 & z_2 & 1 & -x_2*w_2 & -y_2*w_2 & -z_2 * w_2 \\ 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & x_5 & y_5 & z_5 & 1 & -x_5*w_5 & -y_5*w_5 & -z_5 * w_5 \\ 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & x_7 & y_7 & z_7 & 1 & -x_7*w_7 & -y_7*w_7 & -z_7 * w_7 \\ \end{pmatrix} Args: src (torch.Tensor): coordinates of quadrangle vertices in the source image with shape :math:`(B, 8, 3)`. dst (torch.Tensor): coordinates of the corresponding quadrangle vertices in the destination image with shape :math:`(B, 8, 3)`. Returns: torch.Tensor: the perspective transformation with shape :math:`(B, 4, 4)`. .. note:: This function is often used in conjuntion with :func:`warp_perspective3d`. """ if not isinstance(src, (torch.Tensor)): raise TypeError("Input type is not a torch.Tensor. Got {}" .format(type(src))) if not isinstance(dst, (torch.Tensor)): raise TypeError("Input type is not a torch.Tensor. Got {}" .format(type(dst))) if not src.shape[-2:] == (8, 3): raise ValueError("Inputs must be a Bx8x3 tensor. Got {}" .format(src.shape)) if not src.shape == dst.shape: raise ValueError("Inputs must have the same shape. Got {}" .format(dst.shape)) if not (src.shape[0] == dst.shape[0]): raise ValueError("Inputs must have same batch size dimension. Expect {} but got {}" .format(src.shape, dst.shape)) assert src.device == dst.device and src.dtype == dst.dtype, ( f"Expect `src` and `dst` to be in the same device (Got {src.dtype}, {dst.dtype}) " f"with the same dtype (Got {src.dtype}, {dst.dtype})." ) # we build matrix A by using only 4 point correspondence. The linear # system is solved with the least square method, so here # we could even pass more correspondence p = [] # 000, 100, 110, 101, 011 for i in [0, 1, 2, 5, 7]: p.append(_build_perspective_param3d(src[:, i], dst[:, i], 'x')) p.append(_build_perspective_param3d(src[:, i], dst[:, i], 'y')) p.append(_build_perspective_param3d(src[:, i], dst[:, i], 'z')) # A is Bx15x15 A = torch.stack(p, dim=1) # b is a Bx15x1 b = torch.stack([ dst[:, 0:1, 0], dst[:, 0:1, 1], dst[:, 0:1, 2], dst[:, 1:2, 0], dst[:, 1:2, 1], dst[:, 1:2, 2], dst[:, 2:3, 0], dst[:, 2:3, 1], dst[:, 2:3, 2], # dst[:, 3:4, 0], dst[:, 3:4, 1], dst[:, 3:4, 2], # dst[:, 4:5, 0], dst[:, 4:5, 1], dst[:, 4:5, 2], dst[:, 5:6, 0], dst[:, 5:6, 1], dst[:, 5:6, 2], # dst[:, 6:7, 0], dst[:, 6:7, 1], dst[:, 6:7, 2], dst[:, 7:8, 0], dst[:, 7:8, 1], dst[:, 7:8, 2], ], dim=1) # solve the system Ax = b X, LU = _torch_solve_cast(b, A) # create variable to return batch_size = src.shape[0] M = torch.ones(batch_size, 16, device=src.device, dtype=src.dtype) M[..., :15] = torch.squeeze(X, dim=-1) return M.view(-1, 4, 4) # Bx4x4
def _build_perspective_param3d(p: torch.Tensor, q: torch.Tensor, axis: str) -> torch.Tensor: ones = torch.ones_like(p)[..., 0:1] zeros = torch.zeros_like(p)[..., 0:1] if axis == 'x': return torch.cat([ p[:, 0:1], p[:, 1:2], p[:, 2:3], ones, zeros, zeros, zeros, zeros, zeros, zeros, zeros, zeros, -p[:, 0:1] * q[:, 0:1], -p[:, 1:2] * q[:, 0:1], -p[:, 2:3] * q[:, 0:1] ], dim=1) if axis == 'y': return torch.cat([ zeros, zeros, zeros, zeros, p[:, 0:1], p[:, 1:2], p[:, 2:3], ones, zeros, zeros, zeros, zeros, -p[:, 0:1] * q[:, 1:2], -p[:, 1:2] * q[:, 1:2], -p[:, 2:3] * q[:, 1:2] ], dim=1) if axis == 'z': return torch.cat([ zeros, zeros, zeros, zeros, zeros, zeros, zeros, zeros, p[:, 0:1], p[:, 1:2], p[:, 2:3], ones, -p[:, 0:1] * q[:, 2:3], -p[:, 1:2] * q[:, 2:3], -p[:, 2:3] * q[:, 2:3] ], dim=1) raise NotImplementedError(f"perspective params for axis `{axis}` is not implemented.")
[docs]def warp_perspective3d(src: torch.Tensor, M: torch.Tensor, dsize: Tuple[int, int, int], flags: str = 'bilinear', border_mode: str = 'zeros', align_corners: bool = False) -> torch.Tensor: r"""Applies a perspective transformation to an image. The function warp_perspective transforms the source image using the specified matrix: .. math:: \text{dst} (x, y) = \text{src} \left( \frac{M_{11} x + M_{12} y + M_{13}}{M_{31} x + M_{32} y + M_{33}} , \frac{M_{21} x + M_{22} y + M_{23}}{M_{31} x + M_{32} y + M_{33}} \right ) Args: src (torch.Tensor): input image with shape :math:`(B, C, D, H, W)`. M (torch.Tensor): transformation matrix with shape :math:`(B, 4, 4)`. dsize (tuple): size of the output image (height, width). flags (str): interpolation mode to calculate output values 'bilinear' | 'nearest'. Default: 'bilinear'. border_mode (str): padding mode for outside grid values 'zeros' | 'border' | 'reflection'. Default: 'zeros'. align_corners(bool): interpolation flag. Default: False. Returns: torch.Tensor: the warped input image :math:`(B, C, D, H, W)`. .. note:: This function is often used in conjuntion with :func:`get_perspective_transform3d`. """ check_is_tensor(src) check_is_tensor(M) if not len(src.shape) == 5: raise ValueError("Input src must be a BxCxDxHxW tensor. Got {}" .format(src.shape)) if not (len(M.shape) == 3 or M.shape[-2:] == (4, 4)): raise ValueError("Input M must be a Bx4x4 tensor. Got {}" .format(M.shape)) # launches the warper d, h, w = src.shape[-3:] return transform_warp_impl3d(src, M, (d, h, w), dsize, flags, border_mode, align_corners)
def transform_warp_impl3d(src: torch.Tensor, dst_pix_trans_src_pix: torch.Tensor, dsize_src: Tuple[int, int, int], dsize_dst: Tuple[int, int, int], grid_mode: str, padding_mode: str, align_corners: bool) -> torch.Tensor: """Compute the transform in normalized cooridnates and perform the warping. """ dst_norm_trans_src_norm: torch.Tensor = normalize_homography3d( dst_pix_trans_src_pix, dsize_src, dsize_dst) src_norm_trans_dst_norm = torch.inverse(dst_norm_trans_src_norm) return homography_warp3d(src, src_norm_trans_dst_norm, dsize_dst, grid_mode, padding_mode, align_corners, True)