def apply_affine(input: torch.Tensor,
params: Dict[str, torch.Tensor]) -> torch.Tensor:
r"""Random affine transformation of the image keeping center invariant.
Args:
input (torch.Tensor): Tensor to be transformed with shape (H, W), (C, H, W), (B, C, H, W).
params (Dict[str, torch.Tensor]):
- params['angle']: Degrees of rotation.
- params['translations']: Horizontal and vertical translations.
- params['center']: Rotation center.
- params['scale']: Scaling params.
- params['sx']: Shear param toward x-axis.
- params['sy']: Shear param toward y-axis.
- params['resample']: Integer tensor. NEAREST = 0, BILINEAR = 1.
- params['align_corners']: Boolean tensor.
Returns:
torch.Tensor: The transfromed input
"""
if not torch.is_tensor(input):
raise TypeError(f"Input type is not a torch.Tensor. Got {type(input)}")
input = _transform_input(input)
_validate_input_dtype(
input, accepted_dtypes=[torch.float16, torch.float32, torch.float64])
# arrange input data
x_data: torch.Tensor = input.view(-1, *input.shape[-3:])
height, width = x_data.shape[-2:]
# concatenate transforms
transform: torch.Tensor = compute_affine_transformation(input, params)
resample_name: str = Resample(params['resample'].item()).name.lower()
align_corners: bool = cast(bool, params['align_corners'].item())
out_data: torch.Tensor = warp_affine(x_data,
transform[:, :2, :], (height, width),
resample_name,
align_corners=align_corners)
return out_data.view_as(input)
def apply_affine(input: torch.Tensor, params: Dict[str, torch.Tensor],
flags: Dict[str, torch.Tensor]) -> torch.Tensor:
r"""Random affine transformation of the image keeping center invariant.
Args:
input (torch.Tensor): Tensor to be transformed with shape :math:`(*, C, H, W)`.
params (Dict[str, torch.Tensor]):
- params['angle']: Degrees of rotation.
- params['translations']: Horizontal and vertical translations.
- params['center']: Rotation center.
- params['scale']: Scaling params.
- params['sx']: Shear param toward x-axis.
- params['sy']: Shear param toward y-axis.
flags (Dict[str, torch.Tensor]):
- params['resample']: Integer tensor. NEAREST = 0, BILINEAR = 1.
- params['padding_mode']: Integer tensor, see SamplePadding enum.
- params['align_corners']: Boolean tensor.
Returns:
torch.Tensor: The transfromed input with shape :math:`(B, C, H, W)`.
"""
# arrange input data
x_data: torch.Tensor = input.view(-1, *input.shape[-3:])
height, width = x_data.shape[-2:]
# concatenate transforms
transform: torch.Tensor = compute_affine_transformation(input, params)
resample_name: str = Resample(flags['resample'].item()).name.lower()
padding_mode: str = SamplePadding(
flags['padding_mode'].item()).name.lower()
align_corners: bool = cast(bool, flags['align_corners'].item())
out_data: torch.Tensor = warp_affine(
x_data,
transform[:, :2, :],
(height, width),
resample_name,
align_corners=align_corners,
padding_mode=padding_mode,
)
return out_data.view_as(input)
def __init__(
self, degrees: Union[torch.Tensor, float, Tuple[float, float], Tuple[float, float, float],
Tuple[Tuple[float, float], Tuple[float, float], Tuple[float, float]]],
translate: Optional[Union[torch.Tensor, Tuple[float, float, float]]] = None,
scale: Optional[Union[torch.Tensor, Tuple[float, float],
Tuple[Tuple[float, float], Tuple[float, float], Tuple[float, float]]]] = None,
shears: Union[torch.Tensor, float, Tuple[float, float], Tuple[float, float, float, float, float, float],
Tuple[Tuple[float, float], Tuple[float, float], Tuple[float, float], Tuple[float, float],
Tuple[float, float], Tuple[float, float]]] = None,
resample: Union[str, int, Resample] = Resample.BILINEAR.name,
return_transform: bool = False, same_on_batch: bool = False, align_corners: bool = False, p: float = 0.5
) -> None:
super(RandomAffine3D, self).__init__(p=p, return_transform=return_transform, same_on_batch=same_on_batch)
self.degrees = _tuple_range_reader(degrees, 3)
self.shear: Optional[torch.Tensor] = None
if shears is not None:
self.shear = _tuple_range_reader(shears, 6)
# check translation range
self.translate: Optional[torch.Tensor] = None
if translate is not None:
self.translate = translate if isinstance(translate, torch.Tensor) else torch.tensor(translate)
_singular_range_check(self.translate, 'translate', bounds=(0, 1), mode='3d')
# check scale range
self.scale: Optional[torch.Tensor] = None
if scale is not None:
self.scale = scale if isinstance(scale, torch.Tensor) else torch.tensor(scale)
if self.scale.shape == torch.Size([2]):
self.scale = self.scale.unsqueeze(0).repeat(3, 1)
elif self.scale.shape != torch.Size([3, 2]):
raise ValueError("'scale' shall be either shape (2) or (3, 2). Got {self.scale}")
_singular_range_check(self.scale[0], 'scale-x', bounds=(0, float('inf')), mode='2d')
_singular_range_check(self.scale[1], 'scale-y', bounds=(0, float('inf')), mode='2d')
_singular_range_check(self.scale[2], 'scale-z', bounds=(0, float('inf')), mode='2d')
self.resample = Resample.get(resample)
self.align_corners = align_corners
self.flags: Dict[str, torch.Tensor] = dict(
resample=torch.tensor(self.resample.value),
align_corners=torch.tensor(align_corners)
)
def apply_affine3d(input: torch.Tensor, params: Dict[str, torch.Tensor],
flags: Dict[str, torch.Tensor]) -> torch.Tensor:
r"""Random affine transformation of the image keeping center invariant.
Args:
input (torch.Tensor): Tensor to be transformed with shape :math:`(*, C, D, H, W)`.
params (Dict[str, torch.Tensor]):
- params['angles']: Degrees of rotation with the shape of :math: `(*, 3)` for yaw, pitch, roll.
- params['translations']: Horizontal, vertical and depthical translations (dx,dy,dz).
- params['center']: Rotation center (x,y,z).
- params['scale']: Isotropic scaling params.
- params['sxy']: Shear param toward x-y-axis.
- params['sxz']: Shear param toward x-z-axis.
- params['syx']: Shear param toward y-x-axis.
- params['syz']: Shear param toward y-z-axis.
- params['szx']: Shear param toward z-x-axis.
- params['szy']: Shear param toward z-y-axis.
flags (Dict[str, torch.Tensor]):
- params['resample']: Integer tensor. NEAREST = 0, BILINEAR = 1.
- params['align_corners']: Boolean tensor.
Returns:
torch.Tensor: Affine transfromed input with shape :math:`(*, C, D, H, W)`.
"""
# arrange input data
x_data: torch.Tensor = input.view(-1, *input.shape[-4:])
depth, height, width = x_data.shape[-3:]
# concatenate transforms
transform: torch.Tensor = compute_affine_transformation3d(input, params)
resample_name: str = Resample(flags['resample'].item()).name.lower()
align_corners: bool = cast(bool, flags['align_corners'].item())
out_data: torch.Tensor = warp_affine3d(x_data,
transform[:, :3, :],
(depth, height, width),
resample_name,
align_corners=align_corners)
return out_data.view_as(input)
def __init__(
self,
degrees: Union[Tensor, float, Tuple[float, float], List[float]],
resample: Union[str, int, Resample] = Resample.BILINEAR.name,
same_on_batch: bool = False,
align_corners: bool = True,
p: float = 0.5,
keepdim: bool = False,
return_transform: Optional[bool] = None,
) -> None:
super().__init__(p=p,
return_transform=return_transform,
same_on_batch=same_on_batch,
keepdim=keepdim)
self._param_generator = cast(
rg.PlainUniformGenerator,
rg.PlainUniformGenerator(
(degrees, "degrees", 0.0, (-360.0, 360.0))))
self.flags = dict(resample=Resample.get(resample),
align_corners=align_corners)
def __init__(
self,
distortion_scale: Union[torch.Tensor, float] = 0.5,
resample: Union[str, int, Resample] = Resample.BILINEAR.name,
return_transform: bool = False,
same_on_batch: bool = False,
align_corners: bool = False,
p: float = 0.5,
keepdim: bool = False,
) -> None:
super(RandomPerspective3D, self).__init__(
p=p, return_transform=return_transform, same_on_batch=same_on_batch, keepdim=keepdim
)
self._device, self._dtype = _extract_device_dtype([distortion_scale])
self.distortion_scale = distortion_scale
self.resample = Resample.get(resample)
self.align_corners = align_corners
self.flags: Dict[str, torch.Tensor] = dict(
interpolation=torch.tensor(self.resample.value), align_corners=torch.tensor(align_corners)
)
def __init__(
self,
degrees: Union[Tensor, float, Tuple[float, float, float],
Tuple[Tuple[float, float], Tuple[float, float],
Tuple[float, float]], ],
resample: Union[str, int, Resample] = Resample.BILINEAR.name,
same_on_batch: bool = False,
align_corners: bool = False,
p: float = 0.5,
keepdim: bool = False,
return_transform: Optional[bool] = None,
) -> None:
super().__init__(p=p,
return_transform=return_transform,
same_on_batch=same_on_batch,
keepdim=keepdim)
self.flags = dict(resample=Resample.get(resample),
align_corners=align_corners)
self._param_generator = cast(rg.RotationGenerator3D,
rg.RotationGenerator3D(degrees))
def random_rotation_generator3d(
batch_size: int,
degrees: Union[torch.Tensor, float, Tuple[float, float],
Tuple[float, float, float], Tuple[Tuple[float, float],
Tuple[float, float],
Tuple[float, float]]],
interpolation: Union[str, int, Resample] = Resample.BILINEAR.name,
same_on_batch: bool = False,
align_corners: bool = False) -> Dict[str, torch.Tensor]:
r"""Get parameters for ``rotate`` for a random rotate transform.
Args:
batch_size (int): the tensor batch size.
degrees (float or tuple or list): Range of degrees to select from.
If degrees is a number, then yaw, pitch, roll will be generated from the range of (-degrees, +degrees).
If degrees is a tuple of (min, max), then yaw, pitch, roll will be generated from the range of (min, max).
If degrees is a list of floats [a, b, c], then yaw, pitch, roll will be generated from (-a, a), (-b, b)
and (-c, c).
If degrees is a list of tuple ((a, b), (m, n), (x, y)), then yaw, pitch, roll will be generated from
(a, b), (m, n) and (x, y).
interpolation (int, str or kornia.Resample): Default: Resample.BILINEAR
same_on_batch (bool): apply the same transformation across the batch. Default: False
align_corners (bool): interpolation flag. Default: False.
Returns:
params Dict[str, torch.Tensor]: parameters to be passed for transformation.
"""
_degrees = _tuple_range_reader(degrees, 3)
yaw = _adapted_uniform((batch_size, ), _degrees[0][0], _degrees[0][1],
same_on_batch)
pitch = _adapted_uniform((batch_size, ), _degrees[1][0], _degrees[1][1],
same_on_batch)
roll = _adapted_uniform((batch_size, ), _degrees[2][0], _degrees[2][1],
same_on_batch)
return dict(yaw=yaw,
pitch=pitch,
roll=roll,
interpolation=torch.tensor(Resample.get(interpolation).value),
align_corners=torch.tensor(align_corners))
def apply_perspective3d(input: torch.Tensor, params: Dict[str, torch.Tensor],
flags: Dict[str, torch.Tensor]) -> torch.Tensor:
r"""Perform perspective transform of the given torch.Tensor or batch of tensors.
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['start_points']: Tensor containing [top-left, top-right, bottom-right,
bottom-left] of the orignal image with shape Bx8x3.
- params['end_points']: Tensor containing [top-left, top-right, bottom-right,
bottom-left] of the transformed image with shape Bx8x3.
flags (Dict[str, torch.Tensor]):
- params['interpolation']: Integer tensor. NEAREST = 0, BILINEAR = 1.
- params['align_corners']: Boolean tensor.
Returns:
torch.Tensor: Perspectively transformed tensor.
"""
input = _transform_input3d(input)
_validate_input_dtype(
input, accepted_dtypes=[torch.float16, torch.float32, torch.float64])
_, _, depth, height, width = input.shape
# compute the homography between the input points
transform: torch.Tensor = compute_perspective_transformation3d(
input, params)
out_data: torch.Tensor = input.clone()
# apply the computed transform
depth, height, width = input.shape[-3:]
resample_name: str = Resample(flags['interpolation'].item()).name.lower()
align_corners: bool = cast(bool, flags['align_corners'].item())
out_data = warp_perspective3d(input,
transform, (depth, height, width),
flags=resample_name,
align_corners=align_corners)
return out_data.view_as(input)
def random_rotation_generator(
batch_size: int,
degrees: FloatUnionType,
interpolation: Union[str, int, Resample] = Resample.BILINEAR.name,
same_on_batch: bool = False) -> Dict[str, torch.Tensor]:
r"""Get parameters for ``rotate`` for a random rotate transform.
Args:
batch_size (int): the tensor batch size.
degrees (sequence or float or tensor): range of degrees to select from. If degrees is a number the
range of degrees to select from will be (-degrees, +degrees)
interpolation (int, str or kornia.Resample): Default: Resample.BILINEAR
same_on_batch (bool): apply the same transformation across the batch. Default: False
"""
if not torch.is_tensor(degrees):
if isinstance(degrees, (float, int)):
if degrees < 0:
raise ValueError(
f"If Degrees is only one number it must be a positive number. Got{degrees}"
)
degrees = torch.tensor([-degrees, degrees]).to(torch.float32)
elif isinstance(degrees, (tuple, list)):
degrees = torch.tensor(degrees).to(torch.float32)
else:
raise TypeError(
f"Degrees should be a float number a sequence or a tensor. Got {type(degrees)}"
)
# https://mypy.readthedocs.io/en/latest/casts.html cast to please mypy gods
degrees = cast(torch.Tensor, degrees)
if degrees.numel() != 2:
raise ValueError("If degrees is a sequence it must be of length 2")
degrees = _adapted_uniform((batch_size, ), degrees[0], degrees[1],
same_on_batch)
return dict(degrees=degrees,
interpolation=torch.tensor(Resample.get(interpolation).value))
def __init__(
self,
size: Union[int, Tuple[int, int]],
side: str = "short",
resample: Union[str, int, Resample] = Resample.BILINEAR.name,
align_corners: bool = True,
p: float = 1.0,
return_transform: Optional[bool] = None,
keepdim: bool = False,
) -> None:
super().__init__(p=1.,
return_transform=return_transform,
same_on_batch=True,
p_batch=p,
keepdim=keepdim)
self._param_generator = cast(
rg.ResizeGenerator, rg.ResizeGenerator(resize_to=size, side=side))
self.flags = dict(size=size,
side=side,
resample=Resample.get(resample),
align_corners=align_corners)
def __init__(
self,
distortion_scale: Union[Tensor, float] = 0.5,
resample: Union[str, int, Resample] = Resample.BILINEAR.name,
same_on_batch: bool = False,
align_corners: bool = False,
p: float = 0.5,
keepdim: bool = False,
sampling_method: str = "basic",
return_transform: Optional[bool] = None,
) -> None:
super().__init__(p=p,
return_transform=return_transform,
same_on_batch=same_on_batch,
keepdim=keepdim)
self._param_generator = cast(
rg.PerspectiveGenerator,
rg.PerspectiveGenerator(distortion_scale,
sampling_method=sampling_method))
self.flags: Dict[str, Any] = dict(align_corners=align_corners,
resample=Resample.get(resample))
def apply_rotation(input: torch.Tensor,
params: Dict[str, torch.Tensor],
return_transform: bool = False):
r"""Rotate a tensor image or a batch of tensor images a random amount of degrees.
Input should be a tensor of shape (C, H, W) or a batch of tensors :math:`(*, C, H, W)`.
Args:
params (dict): A dict that must have {'degrees': torch.Tensor}. Can be generated from
kornia.augmentation.random_generator.random_rotation_generator
return_transform (bool): if ``True`` return the matrix describing the transformation applied to each
input tensor. If ``False`` and the input is a tuple the applied transformation
wont be concatenated
"""
input = _transform_input(input)
_validate_input_dtype(
input, accepted_dtypes=[torch.float16, torch.float32, torch.float64])
angles: torch.Tensor = params["degrees"].type_as(input)
transformed: torch.Tensor = rotate(
input,
angles,
mode=Resample(params['interpolation'].item()).name.lower())
if return_transform:
# TODO: This part should be inferred from rotate directly
center: torch.Tensor = _compute_tensor_center(input)
rotation_mat: torch.Tensor = _compute_rotation_matrix(
angles, center.expand(angles.shape[0], -1))
# rotation_mat is B x 2 x 3 and we need a B x 3 x 3 matrix
trans_mat: torch.Tensor = torch.eye(3,
device=input.device,
dtype=input.dtype).repeat(
input.shape[0], 1, 1)
trans_mat[:, 0] = rotation_mat[:, 0]
trans_mat[:, 1] = rotation_mat[:, 1]
return transformed, trans_mat
return transformed
def apply_rotation(input: torch.Tensor, params: Dict[str, torch.Tensor]) -> torch.Tensor:
r"""Rotate a tensor image or a batch of tensor images a random amount of degrees.
Input should be a tensor of shape (C, H, W) or a batch of tensors :math:`(*, C, H, W)`.
Args:
input (torch.Tensor): input image.
params (Dict[str, torch.Tensor]):
- params['degrees']: degree to be applied.
Returns:
torch.Tensor: The cropped input
"""
input = _transform_input(input)
_validate_input_dtype(input, accepted_dtypes=[torch.float16, torch.float32, torch.float64])
angles: torch.Tensor = params["degrees"].type_as(input)
resample_mode: str = Resample(params['interpolation'].item()).name.lower()
align_corners: bool = cast(bool, params['align_corners'].item())
transformed: torch.Tensor = rotate(input, angles, mode=resample_mode, align_corners=align_corners)
return transformed
def __init__(
self, size: Tuple[int, int, int],
padding: Optional[Union[int, Tuple[int, int, int], Tuple[int, int, int, int, int, int]]] = None,
pad_if_needed: Optional[bool] = False, fill: int = 0, padding_mode: str = 'constant',
resample: Union[str, int, Resample] = Resample.BILINEAR.name,
return_transform: bool = False, same_on_batch: bool = False,
align_corners: bool = True, p: float = 1.0, keepdim: bool = False
) -> None:
# Since PyTorch does not support ragged tensor. So cropping function happens batch-wisely.
super(RandomCrop3D, self).__init__(
p=1., return_transform=return_transform, same_on_batch=same_on_batch, p_batch=p, keepdim=keepdim)
self.size = size
self.padding = padding
self.pad_if_needed = pad_if_needed
self.fill = fill
self.padding_mode = padding_mode
self.resample = Resample.get(resample)
self.align_corners = align_corners
self.flags: Dict[str, torch.Tensor] = dict(
interpolation=torch.tensor(self.resample.value),
align_corners=torch.tensor(align_corners)
)
def apply_crop(input: torch.Tensor, params: Dict[str, torch.Tensor], flags: Dict[str, torch.Tensor]) -> torch.Tensor:
r"""Apply cropping by src bounding box and dst bounding box.
Order: top-left, top-right, bottom-right and bottom-left. The coordinates must be in the x, y order.
Args:
input (torch.Tensor): Tensor to be transformed with shape :math:`(*, C, H, W)`.
params (Dict[str, torch.Tensor]):
- params['src']: The applied cropping src matrix :math: `(*, 4, 2)`.
- params['dst']: The applied cropping dst matrix :math: `(*, 4, 2)`.
flags (Dict[str, torch.Tensor]):
- params['interpolation']: Integer tensor. NEAREST = 0, BILINEAR = 1.
- params['align_corners']: Boolean tensor.
Returns:
torch.Tensor: The cropped input.
"""
resample_mode: str = Resample.get(flags['interpolation'].item()).name.lower() # type: ignore
align_corners: bool = cast(bool, flags['align_corners'].item())
return crop_by_boxes(
input, params['src'], params['dst'], resample_mode, align_corners=align_corners)
def apply_rotation(input: torch.Tensor,
params: Dict[str, torch.Tensor]) -> torch.Tensor:
r"""Rotate a tensor image or a batch of tensor images a random amount of degrees.
Input should be a tensor of shape (C, H, W) or a batch of tensors :math:`(*, C, H, W)`.
Args:
params (dict): A dict that must have {'degrees': torch.Tensor}. Can be generated from
kornia.augmentation.random_generator.random_rotation_generator
return_transform (bool): if ``True`` return the matrix describing the transformation applied to each
input tensor. If ``False`` and the input is a tuple the applied transformation
wont be concatenated
"""
input = _transform_input(input)
_validate_input_dtype(
input, accepted_dtypes=[torch.float16, torch.float32, torch.float64])
angles: torch.Tensor = params["degrees"].type_as(input)
transformed: torch.Tensor = rotate(
input,
angles,
mode=Resample(params['interpolation'].item()).name.lower())
return transformed
def apply_crop(input: torch.Tensor, params: Dict[str, torch.Tensor]) -> torch.Tensor:
r"""Apply cropping by src bounding box and dst bounding box.
Order: top-left, top-right, bottom-right and bottom-left. The coordinates must be in the x, y order.
Args:
input (torch.Tensor): input image.
params (Dict[str, torch.Tensor]):
- params['src']: The applied cropping src matrix :math: `(*, 4, 2)`.
- params['dst']: The applied cropping dst matrix :math: `(*, 4, 2)`.
- params['interpolation']: Integer tensor. NEAREST = 0, BILINEAR = 1.
- params['align_corners']: Boolean tensor.
Returns:
torch.Tensor: The cropped input.
"""
input = _transform_input(input)
_validate_input_dtype(input, accepted_dtypes=[torch.float16, torch.float32, torch.float64])
resample_mode: str = Resample.get(params['interpolation'].item()).name.lower() # type: ignore
align_corners: bool = cast(bool, params['align_corners'].item())
return crop_by_boxes(
input, params['src'], params['dst'], resample_mode, align_corners=align_corners)
def __init__(
self,
size: Union[int, Tuple[int, int, int]],
align_corners: bool = True,
resample: Union[str, int, Resample] = Resample.BILINEAR.name,
return_transform: bool = False,
p: float = 1.0,
keepdim: bool = False,
) -> None:
# same_on_batch is always True for CenterCrop
# Since PyTorch does not support ragged tensor. So cropping function happens batch-wisely.
super().__init__(p=1.0, return_transform=return_transform, same_on_batch=True, p_batch=p, keepdim=keepdim)
if isinstance(size, tuple):
size = (size[0], size[1], size[2])
elif isinstance(size, int):
size = (size, size, size)
else:
raise Exception(f"Invalid size type. Expected (int, tuple(int, int int). Got: {size}.")
self.flags = dict(
size=size,
align_corners=align_corners,
resample=Resample.get(resample),
)
def apply_motion_blur3d(input: torch.Tensor, params: Dict[str, torch.Tensor],
flags: Dict[str, torch.Tensor]) -> torch.Tensor:
r"""Perform motion blur on an image.
The input image is expected to be in the range of [0, 1].
Args:
input (torch.Tensor): Tensor to be transformed with shape (H, W), (C, H, W), (B, C, H, W).
params (Dict[str, torch.Tensor]):
- params['ksize_factor']: motion kernel width and height (odd and positive).
- params['angle_factor']: yaw, pitch and roll range of the motion blur in degrees :math:`(B, 3)`.
- params['direction_factor']: forward/backward direction of the motion blur.
Lower values towards -1.0 will point the motion blur towards the back (with
angle provided via angle), while higher values towards 1.0 will point the motion
blur forward. A value of 0.0 leads to a uniformly (but still angled) motion blur.
flags (Dict[str, torch.Tensor]):
- flags['border_type']: the padding mode to be applied before convolving.
CONSTANT = 0, REFLECT = 1, REPLICATE = 2, CIRCULAR = 3. Default: BorderType.CONSTANT.
Returns:
torch.Tensor: Adjusted image with the shape as the inpute (\*, C, H, W).
"""
input = _transform_input3d(input)
_validate_input_dtype(
input, accepted_dtypes=[torch.float16, torch.float32, torch.float64])
kernel_size: int = cast(int, params['ksize_factor'].unique().item())
angle = params['angle_factor']
direction = params['direction_factor']
border_type: str = cast(
str,
BorderType(flags['border_type'].item()).name.lower())
mode: str = cast(str, Resample(flags['interpolation'].item()).name.lower())
return motion_blur3d(input, kernel_size, angle, direction, border_type,
mode)
def random_perspective_generator(
batch_size: int,
height: int,
width: int,
p: float,
distortion_scale: float,
interpolation: Union[str, int, Resample] = Resample.BILINEAR.name,
same_on_batch: bool = False,
align_corners: bool = False,
) -> Dict[str, torch.Tensor]:
r"""Get parameters for ``perspective`` for a random perspective transform.
Args:
batch_size (int): the tensor batch size.
height (int) : height of the image.
width (int): width of the image.
p (float): probability of the image being applied perspective.
distortion_scale (float): it controls the degree of distortion and ranges from 0 to 1. Default value is 0.5.
interpolation (int, str or kornia.Resample): Default: Resample.BILINEAR
same_on_batch (bool): apply the same transformation across the batch. Default: False
align_corners(bool): interpolation flag. Default: False.See
https://pytorch.org/docs/stable/nn.functional.html#torch.nn.functional.interpolate for detail
Returns:
List containing [top-left, top-right, bottom-right, bottom-left] of the original image,
List containing [top-left, top-right, bottom-right, bottom-left] of the transformed image.
The points are in -x order.
"""
params: Dict[str, torch.Tensor] = random_prob_generator(batch_size, p)
start_points, end_points = (_get_perspective_params(
batch_size, width, height, distortion_scale, same_on_batch))
params['start_points'] = start_points
params['end_points'] = end_points
params['interpolation'] = torch.tensor(Resample.get(interpolation).value)
params['align_corners'] = align_corners # type: ignore
return params
def __init__(
self,
kernel_size: Union[int, Tuple[int, int]],
angle: Union[Tensor, float, Tuple[float, float, float],
Tuple[Tuple[float, float], Tuple[float, float],
Tuple[float, float]], ],
direction: Union[Tensor, float, Tuple[float, float]],
border_type: Union[int, str, BorderType] = BorderType.CONSTANT.name,
resample: Union[str, int, Resample] = Resample.NEAREST.name,
same_on_batch: bool = False,
p: float = 0.5,
keepdim: bool = False,
return_transform: Optional[bool] = None,
) -> None:
super().__init__(p=p,
return_transform=return_transform,
same_on_batch=same_on_batch,
p_batch=1.0,
keepdim=keepdim)
self.flags = dict(border_type=BorderType.get(border_type),
resample=Resample.get(resample))
self._param_generator = cast(
rg.MotionBlurGenerator3D,
rg.MotionBlurGenerator3D(kernel_size, angle, direction))
def apply_crop(input: torch.Tensor,
params: Dict[str, torch.Tensor]) -> torch.Tensor:
"""
Args:
params (dict): A dict that must have {'src': torch.Tensor, 'dst': torch.Tensor}. Can be generated from
kornia.augmentation.random_generator.random_crop_generator
return_transform (bool): if ``True`` return the matrix describing the transformation applied to each
input tensor.
Returns:
torch.Tensor: The grayscaled input
torch.Tensor: The applied cropping matrix :math: `(*, 4, 2)` if return_transform flag
is set to ``True``
"""
input = _transform_input(input)
_validate_input_dtype(
input, accepted_dtypes=[torch.float16, torch.float32, torch.float64])
return crop_by_boxes( # type: ignore
input,
params['src'],
params['dst'],
Resample.get(
params['interpolation'].item()).name.lower(), # type: ignore
align_corners=params['align_corners'])
def apply_rotation(
input: torch.Tensor, params: Dict[str, torch.Tensor], flags: Dict[str, torch.Tensor]
) -> torch.Tensor:
r"""Rotate a tensor image or a batch of tensor images a random amount of degrees.
Args:
input (torch.Tensor): Tensor to be transformed with shape :math:`(*, C, H, W)`.
params (Dict[str, torch.Tensor]):
- params['degrees']: degree to be applied.
flags (Dict[str, torch.Tensor]):
- params['interpolation']: Integer tensor. NEAREST = 0, BILINEAR = 1.
- params['align_corners']: Boolean tensor.
Returns:
torch.Tensor: The cropped input.
"""
angles: torch.Tensor = params["degrees"].type_as(input)
resample_mode: str = Resample(flags['interpolation'].item()).name.lower()
align_corners: bool = cast(bool, flags['align_corners'].item())
transformed: torch.Tensor = rotate(input, angles, mode=resample_mode, align_corners=align_corners)
return transformed
def apply_perspective(input: torch.Tensor, params: Dict[str, torch.Tensor],
flags: Dict[str, torch.Tensor]) -> torch.Tensor:
r"""Perform perspective transform of the given torch.Tensor or batch of tensors.
Args:
input (torch.Tensor): Tensor to be transformed with shape :math:`(*, C, H, W)`.
params (Dict[str, torch.Tensor]):
- params['start_points']: Tensor containing [top-left, top-right, bottom-right,
bottom-left] of the original image with shape Bx4x2.
- params['end_points']: Tensor containing [top-left, top-right, bottom-right,
bottom-left] of the transformed image with shape Bx4x2.
flags (Dict[str, torch.Tensor]):
- params['interpolation']: Integer tensor. NEAREST = 0, BILINEAR = 1.
- params['align_corners']: Boolean tensor.
Returns:
torch.Tensor: Perspectively transformed tensor.
"""
_, _, height, width = input.shape
# compute the homography between the input points
transform: torch.Tensor = compute_perspective_transformation(input, params)
out_data: torch.Tensor = input.clone()
# apply the computed transform
height, width = input.shape[-2:]
resample_name: str = Resample(flags['interpolation'].item()).name.lower()
align_corners: bool = cast(bool, flags['align_corners'].item())
out_data = warp_perspective(input,
transform, (height, width),
mode=resample_name,
align_corners=align_corners)
return out_data.view_as(input)
def random_affine_generator3d(
batch_size: int,
depth: int,
height: int,
width: int,
degrees: torch.Tensor,
translate: Optional[torch.Tensor] = None,
scale: Optional[torch.Tensor] = None,
shears: Optional[torch.Tensor] = None,
resample: Union[str, int, Resample] = Resample.BILINEAR.name,
same_on_batch: bool = False,
align_corners: bool = False) -> Dict[str, torch.Tensor]:
r"""Get parameters for ```3d affine``` transformation random affine transform.
Args:
batch_size (int): the tensor batch size.
depth (int) : height of the image.
height (int) : height of the image.
width (int): width of the image.
degrees (torch.Tensor): Ranges of degrees with shape (3, 2) for yaw, pitch and roll.
translate (torch.Tensor, optional): maximum absolute fraction with shape (3,) for horizontal, vertical
and depthical translations. Will not translate by default.
scale (torch.Tensor, optional): scaling factor interval, e.g (a, b), then scale is
randomly sampled from the range a <= scale <= b. Will keep original scale by default.
shear (sequence or float, optional): Range of degrees to select from.
If shear is a number, a shear to the 6 facets in the range (-shear, +shear) will be apllied.
If shear is a tuple of 2 values, a shear to the 6 facets in the range (shear[0], shear[1]) will be applied.
If shear is a tuple of 6 values, a shear to the i-th facet in the range (-shear[i], shear[i])
will be applied.
If shear is a tuple of 6 tuples, a shear to the i-th facet in the range (-shear[i, 0], shear[i, 1])
will be applied.
resample (int, str or kornia.Resample): Default: Resample.BILINEAR
same_on_batch (bool): apply the same transformation across the batch. Default: False
align_corners(bool): interpolation flag. Default: False.See
https://pytorch.org/docs/stable/nn.functional.html#torch.nn.functional.interpolate for detail
Returns:
params Dict[str, torch.Tensor]: parameters to be passed for transformation.
"""
assert degrees.shape == torch.Size(
[3, 2]), f"'degrees' must be the shape of (3, 2). Got {degrees.shape}."
yaw = _adapted_uniform((batch_size, ), degrees[0][0], degrees[0][1],
same_on_batch)
pitch = _adapted_uniform((batch_size, ), degrees[1][0], degrees[1][1],
same_on_batch)
roll = _adapted_uniform((batch_size, ), degrees[2][0], degrees[2][1],
same_on_batch)
angles = torch.cat([yaw, pitch, roll], dim=-1).view((batch_size, -1))
# compute tensor ranges
if scale is not None:
assert scale.shape == torch.Size(
[3, 2]), f"'scale' must be the shape of (3, 2). Got {scale.shape}."
scale = torch.cat([
_adapted_uniform(
(batch_size, ), scale[0, 0], scale[0, 1], same_on_batch),
_adapted_uniform(
(batch_size, ), scale[1, 0], scale[1, 1], same_on_batch),
_adapted_uniform(
(batch_size, ), scale[2, 0], scale[2, 1], same_on_batch),
],
dim=1)
else:
scale = torch.ones(batch_size).repeat(1, 3)
if translate is not None:
assert translate.shape == torch.Size([
3
]), f"'translate' must be the shape of (2). Got {translate.shape}."
max_dx: torch.Tensor = translate[0] * depth
max_dy: torch.Tensor = translate[1] * width
max_dz: torch.Tensor = translate[2] * height
translations = torch.stack([
_adapted_uniform((batch_size, ), -max_dx, max_dx, same_on_batch),
_adapted_uniform((batch_size, ), -max_dy, max_dy, same_on_batch),
_adapted_uniform((batch_size, ), -max_dz, max_dz, same_on_batch)
],
dim=-1)
else:
translations = torch.zeros(batch_size, 3)
center: torch.Tensor = torch.tensor(
[depth, width, height], dtype=torch.float32).view(1, 3) / 2. - 0.5
center = center.expand(batch_size, -1)
if shears is not None:
assert shears.shape == torch.Size(
[6,
2]), f"'shears' must be the shape of (6, 2). Got {shears.shape}."
sxy = _adapted_uniform((batch_size, ), shears[0, 0], shears[0, 1],
same_on_batch)
sxz = _adapted_uniform((batch_size, ), shears[1, 0], shears[1, 1],
same_on_batch)
syx = _adapted_uniform((batch_size, ), shears[2, 0], shears[2, 1],
same_on_batch)
syz = _adapted_uniform((batch_size, ), shears[3, 0], shears[3, 1],
same_on_batch)
szx = _adapted_uniform((batch_size, ), shears[4, 0], shears[4, 1],
same_on_batch)
szy = _adapted_uniform((batch_size, ), shears[5, 0], shears[5, 1],
same_on_batch)
else:
sxy = sxz = syx = syz = szx = szy = torch.tensor([0] * batch_size)
return dict(translations=translations,
center=center,
scale=scale,
angles=angles,
sxy=sxy,
sxz=sxz,
syx=syx,
syz=syz,
szx=szx,
szy=szy,
resample=torch.tensor(Resample.get(resample).value),
align_corners=torch.tensor(align_corners))
def _get_random_affine_params(
batch_size: int,
depth: int,
height: int,
width: int,
degrees: torch.Tensor,
translate: Optional[Tuple[float, float, float]],
scales: Optional[Tuple[float, float]],
shears: Optional[torch.Tensor] = None,
resample: Union[str, int, Resample] = Resample.BILINEAR.name,
same_on_batch: bool = False,
align_corners: bool = False) -> Dict[str, torch.Tensor]:
r"""Get parameters for ```3d affine``` transformation random affine transform.
The returned matrix is Bx4x4.
Returns:
params Dict[str, torch.Tensor]: parameters to be passed for transformation.
"""
yaw = _adapted_uniform((batch_size, ), degrees[0][0], degrees[0][1],
same_on_batch)
pitch = _adapted_uniform((batch_size, ), degrees[1][0], degrees[1][1],
same_on_batch)
roll = _adapted_uniform((batch_size, ), degrees[2][0], degrees[2][1],
same_on_batch)
angles = torch.cat([yaw, pitch, roll], dim=-1).view((batch_size, -1))
# compute tensor ranges
if scales is not None:
scale = _adapted_uniform((batch_size, ), scales[0], scales[1],
same_on_batch)
else:
scale = torch.ones(batch_size)
if translate is not None:
max_dx: float = translate[0] * depth
max_dy: float = translate[1] * width
max_dz: float = translate[2] * height
translations = torch.stack([
_adapted_uniform((batch_size, ), -max_dx, max_dx, same_on_batch),
_adapted_uniform((batch_size, ), -max_dy, max_dy, same_on_batch),
_adapted_uniform((batch_size, ), -max_dz, max_dz, same_on_batch)
],
dim=-1)
else:
translations = torch.zeros(batch_size, 3)
center: torch.Tensor = torch.tensor(
[depth, width, height], dtype=torch.float32).view(1, 3) / 2. - 0.5
center = center.expand(batch_size, -1)
if shears is not None:
sxy = _adapted_uniform((batch_size, ), shears[0, 0], shears[0, 1],
same_on_batch)
sxz = _adapted_uniform((batch_size, ), shears[1, 0], shears[1, 1],
same_on_batch)
syx = _adapted_uniform((batch_size, ), shears[2, 0], shears[2, 1],
same_on_batch)
syz = _adapted_uniform((batch_size, ), shears[3, 0], shears[3, 1],
same_on_batch)
szx = _adapted_uniform((batch_size, ), shears[4, 0], shears[4, 1],
same_on_batch)
szy = _adapted_uniform((batch_size, ), shears[5, 0], shears[5, 1],
same_on_batch)
else:
sxy = sxz = syx = syz = szx = szy = torch.tensor([0] * batch_size)
return dict(translations=translations,
center=center,
scale=scale,
angles=angles,
sxy=sxy,
sxz=sxz,
syx=syx,
syz=syz,
szx=szx,
szy=szy,
resample=torch.tensor(Resample.get(resample).value),
align_corners=torch.tensor(align_corners))
def random_affine_generator(
batch_size: int,
height: int,
width: int,
degrees: torch.Tensor,
translate: Optional[torch.Tensor] = None,
scale: Optional[torch.Tensor] = None,
shear: Optional[torch.Tensor] = None,
resample: Union[str, int, Resample] = Resample.BILINEAR.name,
same_on_batch: bool = False,
align_corners: bool = False,
padding_mode: Union[str, int, SamplePadding] = SamplePadding.ZEROS.name,
) -> Dict[str, torch.Tensor]:
r"""Get parameters for ``affine`` for a random affine transform.
Args:
batch_size (int): the tensor batch size.
height (int) : height of the image.
width (int): width of the image.
degrees (float or tuple): Range of degrees to select from.
If degrees is a number instead of sequence like (min, max), the range of degrees
will be (-degrees, +degrees). Set to 0 to deactivate rotations.
translate (tuple, optional): tuple of maximum absolute fraction for horizontal
and vertical translations. For example translate=(a, b), then horizontal shift
is randomly sampled in the range -img_width * a < dx < img_width * a and vertical shift is
randomly sampled in the range -img_height * b < dy < img_height * b. Will not translate by default.
scale (tuple, optional): scaling factor interval, e.g (a, b), then scale is
randomly sampled from the range a <= scale <= b. Will keep original scale by default.
shear (sequence or float, optional): Range of degrees to select from.
If shear is a number, a shear parallel to the x axis in the range (-shear, +shear)
will be apllied. Else if shear is a tuple or list of 2 values a shear parallel to the x axis in the
range (shear[0], shear[1]) will be applied. Else if shear is a tuple or list of 4 values,
a x-axis shear in (shear[0], shear[1]) and y-axis shear in (shear[2], shear[3]) will be applied.
Will not apply shear by default
resample (int, str or kornia.Resample): Default: Resample.BILINEAR
same_on_batch (bool): apply the same transformation across the batch. Default: False
align_corners(bool): interpolation flag. Default: False.See
https://pytorch.org/docs/stable/nn.functional.html#torch.nn.functional.interpolate for detail
padding_mode (int, str or kornia.SamplePadding): Default: SamplePadding.ZEROS
Returns:
params Dict[str, torch.Tensor]: parameters to be passed for transformation.
"""
_joint_range_check(degrees, "degrees")
angle = _adapted_uniform((batch_size, ), degrees[0], degrees[1],
same_on_batch)
# compute tensor ranges
if scale is not None:
_joint_range_check(cast(torch.Tensor, scale[:2]), "scale")
scale = _adapted_uniform((batch_size, ), scale[0], scale[1],
same_on_batch).repeat(1, 2)
if len(scale) == 4:
_joint_range_check(cast(torch.Tensor, scale[2:]), "scale_y")
scale[:, 1] = _adapted_uniform((batch_size, ), scale[2], scale[3],
same_on_batch)
else:
scale = torch.ones((batch_size, 2))
if translate is not None:
_joint_range_check(cast(torch.Tensor, translate), "translate")
max_dx: torch.Tensor = translate[0] * width
max_dy: torch.Tensor = translate[1] * height
translations = torch.stack([
_adapted_uniform((batch_size, ), -max_dx, max_dx, same_on_batch),
_adapted_uniform((batch_size, ), -max_dy, max_dy, same_on_batch)
],
dim=-1)
else:
translations = torch.zeros(batch_size, 2)
center: torch.Tensor = torch.tensor(
[width, height], dtype=torch.float32).view(1, 2) / 2. - 0.5
center = center.expand(batch_size, -1)
if shear is not None:
_joint_range_check(cast(torch.Tensor, shear)[0], "shear")
_joint_range_check(cast(torch.Tensor, shear)[1], "shear")
sx = _adapted_uniform((batch_size, ), shear[0][0], shear[0][1],
same_on_batch)
sy = _adapted_uniform((batch_size, ), shear[1][0], shear[1][1],
same_on_batch)
else:
sx = sy = torch.tensor([0] * batch_size)
return dict(translations=translations,
center=center,
scale=scale,
angle=angle,
sx=sx,
sy=sy,
resample=torch.tensor(Resample.get(resample).value),
padding_mode=torch.tensor(
SamplePadding.get(padding_mode).value),
align_corners=torch.tensor(align_corners))
