from typing import Dict, Optional, Tuple, Union, cast import torch from torch.distributions import Uniform from kornia.augmentation.random_generator.base import RandomGeneratorBase from kornia.augmentation.utils import ( _adapted_rsampling, _adapted_uniform, _common_param_check, _joint_range_check, _range_bound, ) from kornia.utils.helpers import _deprecated, _extract_device_dtype class AffineGenerator(RandomGeneratorBase): r"""Get parameters for ``affine`` for a random affine transform. Args: degrees: Range of degrees to select from like (min, max). translate: 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: 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: Range of degrees to select from. If float, a shear parallel to the x axis in the range (-shear, +shear) will be applied. If (a, b), a shear parallel to the x axis in the range (-shear, +shear) will be applied. If (a, b, c, d), then 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. If tensor, shear is a 2x2 tensor, a x-axis shear in (shear[0][0], shear[0][1]) and y-axis shear in (shear[1][0], shear[1][1]) will be applied. Will not apply shear by default. Returns: A dict of parameters to be passed for transformation. - translations (torch.Tensor): element-wise translations with a shape of (B, 2). - center (torch.Tensor): element-wise center with a shape of (B, 2). - scale (torch.Tensor): element-wise scales with a shape of (B, 2). - angle (torch.Tensor): element-wise rotation angles with a shape of (B,). - sx (torch.Tensor): element-wise x-axis shears with a shape of (B,). - sy (torch.Tensor): element-wise y-axis shears with a shape of (B,). Note: The generated random numbers are not reproducible across different devices and dtypes. By default, the parameters will be generated on CPU in float32. This can be changed by calling ``self.set_rng_device_and_dtype(device="cuda", dtype=torch.float64)``. """ def __init__( self, degrees: Union[torch.Tensor, float, Tuple[float, float]], translate: Optional[Union[torch.Tensor, Tuple[float, float]]] = None, scale: Optional[Union[torch.Tensor, Tuple[float, float], Tuple[float, float, float, float]]] = None, shear: Optional[Union[torch.Tensor, float, Tuple[float, float]]] = None, ) -> None: super().__init__() self.degrees = degrees self.translate = translate self.scale = scale self.shear = shear def __repr__(self) -> str: repr = f"degrees={self.degrees}, translate={self.translate}, scale={self.scale}, shear={self.shear}" return repr def make_samplers(self, device: torch.device, dtype: torch.dtype) -> None: _degrees = _range_bound(self.degrees, 'degrees', 0, (-360, 360)).to(device=device, dtype=dtype) _translate = ( self.translate if self.translate is None else _range_bound(self.translate, 'translate', bounds=(0, 1), check='singular').to( device=device, dtype=dtype ) ) _scale: Optional[torch.Tensor] = None if self.scale is not None: if len(self.scale) == 2: _scale = _range_bound(self.scale[:2], 'scale', bounds=(0, float('inf')), check='singular').to( device=device, dtype=dtype ) elif len(self.scale) == 4: _scale = torch.cat( [ _range_bound(self.scale[:2], 'scale_x', bounds=(0, float('inf')), check='singular'), _range_bound( self.scale[2:], 'scale_y', bounds=(0, float('inf')), check='singular' # type:ignore ), ] ).to(device=device, dtype=dtype) else: raise ValueError(f"'scale' expected to be either 2 or 4 elements. Got {self.scale}") _shear: Optional[torch.Tensor] = None if self.shear is not None: shear = torch.as_tensor(self.shear, device=device, dtype=dtype) if shear.shape == torch.Size([2, 2]): _shear = shear else: _shear = torch.stack( [ _range_bound(shear if shear.dim() == 0 else shear[:2], 'shear-x', 0, (-360, 360)), torch.tensor([0, 0], device=device, dtype=dtype) if shear.dim() == 0 or len(shear) == 2 else _range_bound(shear[2:], 'shear-y', 0, (-360, 360)), ] ) translate_x_sampler: Optional[Uniform] = None translate_y_sampler: Optional[Uniform] = None scale_2_sampler: Optional[Uniform] = None scale_4_sampler: Optional[Uniform] = None shear_x_sampler: Optional[Uniform] = None shear_y_sampler: Optional[Uniform] = None if _translate is not None: translate_x_sampler = Uniform(-_translate[0], _translate[0], validate_args=False) translate_y_sampler = Uniform(-_translate[1], _translate[1], validate_args=False) if _scale is not None: if len(_scale) == 2: scale_2_sampler = Uniform(_scale[0], _scale[1], validate_args=False) elif len(_scale) == 4: scale_2_sampler = Uniform(_scale[0], _scale[1], validate_args=False) scale_4_sampler = Uniform(_scale[2], _scale[3], validate_args=False) else: raise ValueError(f"'scale' expected to be either 2 or 4 elements. Got {self.scale}") if _shear is not None: _joint_range_check(cast(torch.Tensor, _shear)[0], "shear") _joint_range_check(cast(torch.Tensor, _shear)[1], "shear") shear_x_sampler = Uniform(_shear[0][0], _shear[0][1], validate_args=False) shear_y_sampler = Uniform(_shear[1][0], _shear[1][1], validate_args=False) self.degree_sampler = Uniform(_degrees[0], _degrees[1], validate_args=False) self.translate_x_sampler = translate_x_sampler self.translate_y_sampler = translate_y_sampler self.scale_2_sampler = scale_2_sampler self.scale_4_sampler = scale_4_sampler self.shear_x_sampler = shear_x_sampler self.shear_y_sampler = shear_y_sampler def forward(self, batch_shape: torch.Size, same_on_batch: bool = False) -> Dict[str, torch.Tensor]: # type: ignore batch_size = batch_shape[0] height = batch_shape[-2] width = batch_shape[-1] _device, _dtype = _extract_device_dtype([self.degrees, self.translate, self.scale, self.shear]) _common_param_check(batch_size, same_on_batch) if not (isinstance(width, (int,)) and isinstance(height, (int,)) and width > 0 and height > 0): raise AssertionError(f"`width` and `height` must be positive integers. Got {width}, {height}.") angle = _adapted_rsampling((batch_size,), self.degree_sampler, same_on_batch).to(device=_device, dtype=_dtype) # compute tensor ranges if self.scale_2_sampler is not None: _scale = _adapted_rsampling((batch_size,), self.scale_2_sampler, same_on_batch).unsqueeze(1).repeat(1, 2) if self.scale_4_sampler is not None: _scale[:, 1] = _adapted_rsampling((batch_size,), self.scale_4_sampler, same_on_batch) _scale = _scale.to(device=_device, dtype=_dtype) else: _scale = torch.ones((batch_size, 2), device=_device, dtype=_dtype) if self.translate_x_sampler is not None and self.translate_y_sampler is not None: translations = torch.stack( [ _adapted_rsampling((batch_size,), self.translate_x_sampler, same_on_batch) * width, _adapted_rsampling((batch_size,), self.translate_y_sampler, same_on_batch) * height, ], dim=-1, ) translations = translations.to(device=_device, dtype=_dtype) else: translations = torch.zeros((batch_size, 2), device=_device, dtype=_dtype) center: torch.Tensor = torch.tensor([width, height], device=_device, dtype=_dtype).view(1, 2) / 2.0 - 0.5 center = center.expand(batch_size, -1) if self.shear_x_sampler is not None and self.shear_y_sampler is not None: sx = _adapted_rsampling((batch_size,), self.shear_x_sampler, same_on_batch) sy = _adapted_rsampling((batch_size,), self.shear_y_sampler, same_on_batch) sx = sx.to(device=_device, dtype=_dtype) sy = sy.to(device=_device, dtype=_dtype) else: sx = sy = torch.tensor([0] * batch_size, device=_device, dtype=_dtype) return dict(translations=translations, center=center, scale=_scale, angle=angle, sx=sx, sy=sy) @_deprecated(replace_with=AffineGenerator.__name__) 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, same_on_batch: bool = False, device: torch.device = torch.device('cpu'), dtype: torch.dtype = torch.float32, ) -> 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 (torch.Tensor): Range of degrees to select from like (min, max). translate (tensor, 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 (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 (tensor, optional): Range of degrees to select from. Shear is a 2x2 tensor, a x-axis shear in (shear[0][0], shear[0][1]) and y-axis shear in (shear[1][0], shear[1][1]) will be applied. Will not apply shear by default. same_on_batch (bool): apply the same transformation across the batch. Default: False. device (torch.device): the device on which the random numbers will be generated. Default: cpu. dtype (torch.dtype): the data type of the generated random numbers. Default: float32. Returns: params Dict[str, torch.Tensor]: parameters to be passed for transformation. - translations (torch.Tensor): element-wise translations with a shape of (B, 2). - center (torch.Tensor): element-wise center with a shape of (B, 2). - scale (torch.Tensor): element-wise scales with a shape of (B, 2). - angle (torch.Tensor): element-wise rotation angles with a shape of (B,). - sx (torch.Tensor): element-wise x-axis shears with a shape of (B,). - sy (torch.Tensor): element-wise y-axis shears with a shape of (B,). Note: The generated random numbers are not reproducible across different devices and dtypes. """ _common_param_check(batch_size, same_on_batch) _joint_range_check(degrees, "degrees") if not (isinstance(width, (int,)) and isinstance(height, (int,)) and width > 0 and height > 0): raise AssertionError(f"`width` and `height` must be positive integers. Got {width}, {height}.") _device, _dtype = _extract_device_dtype([degrees, translate, scale, shear]) degrees = degrees.to(device=device, dtype=dtype) angle = _adapted_uniform((batch_size,), degrees[0], degrees[1], same_on_batch) angle = angle.to(device=_device, dtype=_dtype) # compute tensor ranges if scale is not None: scale = scale.to(device=device, dtype=dtype) if not (len(scale.shape) == 1 and len(scale) in (2, 4)): raise AssertionError(f"`scale` shall have 2 or 4 elements. Got {scale}.") _joint_range_check(cast(torch.Tensor, scale[:2]), "scale") _scale = _adapted_uniform((batch_size,), scale[0], scale[1], same_on_batch).unsqueeze(1).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) _scale = _scale.to(device=_device, dtype=_dtype) else: _scale = torch.ones((batch_size, 2), device=_device, dtype=_dtype) if translate is not None: translate = translate.to(device=device, dtype=dtype) if not (0.0 <= translate[0] <= 1.0 and 0.0 <= translate[1] <= 1.0 and translate.shape == torch.Size([2])): raise AssertionError(f"Expect translate contains two elements and ranges are in [0, 1]. Got {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, ) translations = translations.to(device=_device, dtype=_dtype) else: translations = torch.zeros((batch_size, 2), device=_device, dtype=_dtype) center: torch.Tensor = torch.tensor([width, height], device=_device, dtype=_dtype).view(1, 2) / 2.0 - 0.5 center = center.expand(batch_size, -1) if shear is not None: shear = shear.to(device=device, dtype=dtype) _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) sx = sx.to(device=_device, dtype=_dtype) sy = sy.to(device=_device, dtype=_dtype) else: sx = sy = torch.tensor([0] * batch_size, device=_device, dtype=_dtype) return dict(translations=translations, center=center, scale=_scale, angle=angle, sx=sx, sy=sy)