def compute_rotate_tranformation3d(input: torch.Tensor,
params: Dict[str, torch.Tensor]):
r"""Compute the applied transformation matrix :math: `(*, 4, 4)`.
Args:
input (torch.Tensor): Tensor to be transformed with shape (D, H, W), (C, D, H, W), (B, C, D, H, W).
params (Dict[str, torch.Tensor]):
- params['yaw']: degree to be applied.
- params['pitch']: degree to be applied.
- params['roll']: degree to be applied.
Returns:
torch.Tensor: The applied transformation matrix :math: `(*, 4, 4)`
"""
input = _transform_input3d(input)
_validate_input_dtype(
input, accepted_dtypes=[torch.float16, torch.float32, torch.float64])
yaw: torch.Tensor = params["yaw"].type_as(input)
pitch: torch.Tensor = params["pitch"].type_as(input)
roll: torch.Tensor = params["roll"].type_as(input)
center: torch.Tensor = _compute_tensor_center3d(input)
rotation_mat: torch.Tensor = _compute_rotation_matrix3d(
yaw, pitch, roll, center.expand(yaw.shape[0], -1))
# rotation_mat is B x 3 x 4 and we need a B x 4 x 4 matrix
trans_mat: torch.Tensor = torch.eye(4,
device=input.device,
dtype=input.dtype).repeat(
input.shape[0], 1, 1)
trans_mat[:, 0] = rotation_mat[:, 0]
trans_mat[:, 1] = rotation_mat[:, 1]
trans_mat[:, 2] = rotation_mat[:, 2]
return trans_mat
def compute_rotate_tranformation3d(input: torch.Tensor, params: Dict[str, torch.Tensor]):
r"""Compute the rotation transformation matrix :math: `(*, 4, 4)`.
Args:
input (torch.Tensor): Tensor to be transformed with shape :math:`(*, C, D, H, W)`.
params (Dict[str, torch.Tensor]):
- params['yaw']: degree to be applied.
- params['pitch']: degree to be applied.
- params['roll']: degree to be applied.
Returns:
torch.Tensor: The rotation transformation matrix :math: `(*, 4, 4)`.
"""
yaw: torch.Tensor = params["yaw"].to(input)
pitch: torch.Tensor = params["pitch"].to(input)
roll: torch.Tensor = params["roll"].to(input)
center: torch.Tensor = _compute_tensor_center3d(input)
rotation_mat: torch.Tensor = _compute_rotation_matrix3d(yaw, pitch, roll, center.expand(yaw.shape[0], -1))
# rotation_mat is B x 3 x 4 and we need a B x 4 x 4 matrix
trans_mat: torch.Tensor = torch.eye(4, device=input.device, dtype=input.dtype).repeat(input.shape[0], 1, 1)
trans_mat[:, 0] = rotation_mat[:, 0]
trans_mat[:, 1] = rotation_mat[:, 1]
trans_mat[:, 2] = rotation_mat[:, 2]
return trans_mat
def compute_transformation(self, input: torch.Tensor, params: Dict[str, torch.Tensor]) -> torch.Tensor:
yaw: torch.Tensor = params["yaw"].to(input)
pitch: torch.Tensor = params["pitch"].to(input)
roll: torch.Tensor = params["roll"].to(input)
center: torch.Tensor = _compute_tensor_center3d(input)
rotation_mat: torch.Tensor = _compute_rotation_matrix3d(yaw, pitch, roll, center.expand(yaw.shape[0], -1))
# rotation_mat is B x 3 x 4 and we need a B x 4 x 4 matrix
trans_mat: torch.Tensor = torch.eye(4, device=input.device, dtype=input.dtype).repeat(input.shape[0], 1, 1)
trans_mat[:, 0] = rotation_mat[:, 0]
trans_mat[:, 1] = rotation_mat[:, 1]
trans_mat[:, 2] = rotation_mat[:, 2]
return trans_mat
def compute_transformation(self, input: Tensor,
params: Dict[str, Tensor]) -> Tensor:
yaw: Tensor = params["yaw"].to(input)
pitch: Tensor = params["pitch"].to(input)
roll: Tensor = params["roll"].to(input)
center: Tensor = _compute_tensor_center3d(input)
rotation_mat: Tensor = _compute_rotation_matrix3d(
yaw, pitch, roll, center.expand(yaw.shape[0], -1))
# rotation_mat is B x 3 x 4 and we need a B x 4 x 4 matrix
trans_mat: Tensor = kornia.eye_like(4, input)
trans_mat[:, 0] = rotation_mat[:, 0]
trans_mat[:, 1] = rotation_mat[:, 1]
trans_mat[:, 2] = rotation_mat[:, 2]
return trans_mat