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))