def __init__(
self,
degrees: Union[Tensor, float, Tuple[float, float]],
translate: Optional[Union[Tensor, Tuple[float, float]]] = None,
scale: Optional[Union[Tensor, Tuple[float, float],
Tuple[float, float, float, float]]] = None,
shear: Optional[Union[Tensor, float, Tuple[float, float]]] = 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,
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.AffineGenerator,
rg.AffineGenerator(degrees, translate, scale, shear))
self.flags = dict(resample=Resample.get(resample),
padding_mode=SamplePadding.get(padding_mode),
align_corners=align_corners)
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 (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.
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
"""
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(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 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))