def test_extract_device_dtype(tensor_list, out_device, out_dtype, will_throw_error): # Add GPU tests when GPU testing avaliable if torch.cuda.is_available(): import warnings warnings.warn("Add GPU tests.") if will_throw_error: with pytest.raises(ValueError): _extract_device_dtype(tensor_list) else: device, dtype = _extract_device_dtype(tensor_list) assert device == out_device assert dtype == out_dtype
def test_extract_device_dtype(tensor_list, out_device, out_dtype, will_throw_error): # TODO: include the warning in another way - possibly loggers. # Add GPU tests when GPU testing available # if torch.cuda.is_available(): # import warnings # warnings.warn("Add GPU tests.") if will_throw_error: with pytest.raises(ValueError): _extract_device_dtype(tensor_list) else: device, dtype = _extract_device_dtype(tensor_list) assert device == out_device assert dtype == out_dtype
def __init__( self, kernel_size: Union[int, Tuple[int, int]], angle: Union[ torch.Tensor, float, Tuple[float, float, float], Tuple[Tuple[float, float], Tuple[float, float], Tuple[float, float]], ], direction: Union[torch.Tensor, float, Tuple[float, float]], border_type: Union[int, str, BorderType] = BorderType.CONSTANT.name, resample: Union[str, int, Resample] = Resample.NEAREST.name, return_transform: bool = False, same_on_batch: bool = False, p: float = 0.5, keepdim: bool = False, ) -> None: super(RandomMotionBlur3D, self).__init__( p=p, return_transform=return_transform, same_on_batch=same_on_batch, p_batch=1.0, keepdim=keepdim ) self._device, self._dtype = _extract_device_dtype([angle, direction]) self.kernel_size: Union[int, Tuple[int, int]] = kernel_size self.angle = angle self.direction = direction self.resample = Resample.get(resample) self.border_type = BorderType.get(border_type) self.flags: Dict[str, torch.Tensor] = { "border_type": torch.tensor(self.border_type.value), "interpolation": torch.tensor(self.resample.value), }
def _adapted_uniform( shape: Union[Tuple, torch.Size], low: Union[float, int, torch.Tensor], high: Union[float, int, torch.Tensor], same_on_batch: bool = False, epsilon: float = 1e-6, ) -> torch.Tensor: r"""The uniform sampling function that accepts 'same_on_batch'. If same_on_batch is True, all values generated will be exactly same given a batch_size (shape[0]). By default, same_on_batch is set to False. By default, sampling happens on the default device and dtype. If low/high is a tensor, sampling will happen in the same device/dtype as low/high tensor. """ device, dtype = _extract_device_dtype([ low if isinstance(low, torch.Tensor) else None, high if isinstance(high, torch.Tensor) else None ]) low = torch.as_tensor(low, device=device, dtype=dtype) high = torch.as_tensor(high, device=device, dtype=dtype) # validate_args=False to fix pytorch 1.7.1 error: # ValueError: Uniform is not defined when low>= high. dist = Uniform(low, high, validate_args=False) return _adapted_rsampling(shape, dist, same_on_batch)
def random_rotation_generator3d( batch_size: int, degrees: torch.Tensor, same_on_batch: bool = False, device: torch.device = torch.device('cpu'), dtype: torch.dtype = torch.float32 ) -> Dict[str, torch.Tensor]: r"""Get parameters for ``rotate`` for a random rotate transform. Args: batch_size (int): the tensor batch size. degrees (torch.Tensor): Ranges of degrees (3, 2) for yaw, pitch and roll. 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. - yaw (torch.Tensor): element-wise rotation yaws with a shape of (B,). - pitch (torch.Tensor): element-wise rotation pitches with a shape of (B,). - roll (torch.Tensor): element-wise rotation rolls with a shape of (B,). """ assert degrees.shape == torch.Size([3, 2]), f"'degrees' must be the shape of (3, 2). Got {degrees.shape}." _device, _dtype = _extract_device_dtype([degrees]) degrees = degrees.to(device=device, dtype=dtype) 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.to(device=_device, dtype=_dtype), pitch=pitch.to(device=_device, dtype=_dtype), roll=roll.to(device=_device, dtype=_dtype))
def __init__( self, degrees: Union[ torch.Tensor, float, Tuple[float, float, float], Tuple[Tuple[float, float], Tuple[float, float], Tuple[float, float]], ], 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(RandomRotation3D, self).__init__( p=p, return_transform=return_transform, same_on_batch=same_on_batch, keepdim=keepdim ) self._device, self._dtype = _extract_device_dtype([degrees]) self.degrees = degrees 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 _adapted_beta(shape: Union[Tuple, torch.Size], a: Union[float, int, torch.Tensor], b: Union[float, int, torch.Tensor], same_on_batch=False) -> torch.Tensor: r""" The beta sampling function that accepts 'same_on_batch'. If same_on_batch is True, all values generated will be exactly same given a batch_size (shape[0]). By default, same_on_batch is set to False. """ device, dtype = _extract_device_dtype([ a if isinstance(a, torch.Tensor) else None, b if isinstance(b, torch.Tensor) else None, ]) dtype = torch.float32 if dtype is None else dtype if not isinstance(a, torch.Tensor): a = torch.tensor(a, dtype=torch.float64) else: a = a.to(torch.float64) if not isinstance(b, torch.Tensor): b = torch.tensor(b, dtype=torch.float64) else: b = b.to(torch.float64) dist = Beta(a, b) return _adapted_rsampling(shape, dist, same_on_batch).to(device=device, dtype=dtype)
def __init__( self, angle: Optional[torch.Tensor] = None, translation: Optional[torch.Tensor] = None, scale_factor: Optional[torch.Tensor] = None, shear: Optional[torch.Tensor] = None, center: Optional[torch.Tensor] = None, mode: str = 'bilinear', padding_mode: str = 'zeros', align_corners: Optional[bool] = None, ) -> None: batch_sizes = [ arg.size()[0] for arg in (angle, translation, scale_factor, shear) if arg is not None ] if not batch_sizes: msg = ( "Affine was created without any affine parameter. At least one of angle, translation, scale_factor, or " "shear has to be set.") raise RuntimeError(msg) batch_size = batch_sizes[0] if not all(other == batch_size for other in batch_sizes[1:]): raise RuntimeError( f"The batch sizes of the affine parameters mismatch: {batch_sizes}" ) self._batch_size = batch_size super().__init__() device, dtype = _extract_device_dtype( [angle, translation, scale_factor]) if angle is None: angle = torch.zeros(batch_size, device=device, dtype=dtype) self.angle = angle if translation is None: translation = torch.zeros(batch_size, 2, device=device, dtype=dtype) self.translation = translation if scale_factor is None: scale_factor = torch.ones(batch_size, 2, device=device, dtype=dtype) self.scale_factor = scale_factor self.shear = shear self.center = center self.mode = mode self.padding_mode = padding_mode self.align_corners = align_corners
def random_color_jitter_generator( batch_size: int, brightness: Optional[torch.Tensor] = None, contrast: Optional[torch.Tensor] = None, saturation: Optional[torch.Tensor] = None, hue: Optional[torch.Tensor] = None, same_on_batch: bool = False) -> Dict[str, torch.Tensor]: r"""Generate random color jiter parameters for a batch of images. Args: batch_size (int): the number of images. brightness (torch.Tensor, optional): Brightness factor tensor of range (a, b). The provided range must follow 0 <= a <= b <= 2. Default value is [0., 0.]. contrast (torch.Tensor, optional): Contrast factor tensor of range (a, b). The provided range must follow 0 <= a <= b. Default value is [0., 0.]. saturation (torch.Tensor, optional): Saturation factor tensor of range (a, b). The provided range must follow 0 <= a <= b. Default value is [0., 0.]. hue (torch.Tensor, optional): Saturation factor tensor of range (a, b). The provided range must follow -0.5 <= a <= b < 0.5. Default value is [0., 0.]. same_on_batch (bool): apply the same transformation across the batch. Default: False. Returns: params Dict[str, torch.Tensor]: parameters to be passed for transformation. """ _common_param_check(batch_size, same_on_batch) device, dtype = _extract_device_dtype( [brightness, contrast, hue, saturation]) brightness = torch.tensor([0., 0.], device=device, dtype=dtype) \ if brightness is None else cast(torch.Tensor, brightness) contrast = torch.tensor([0., 0.], device=device, dtype=dtype) \ if contrast is None else cast(torch.Tensor, contrast) hue = torch.tensor([0., 0.], device=device, dtype=dtype) \ if hue is None else cast(torch.Tensor, hue) saturation = torch.tensor([0., 0.], device=device, dtype=dtype) \ if saturation is None else cast(torch.Tensor, saturation) _joint_range_check(brightness, "brightness", (0, 2)) _joint_range_check(contrast, "contrast", (0, float('inf'))) _joint_range_check(hue, "hue", (-0.5, 0.5)) _joint_range_check(saturation, "saturation", (0, float('inf'))) brightness_factor = _adapted_uniform((batch_size, ), brightness[0], brightness[1], same_on_batch) contrast_factor = _adapted_uniform((batch_size, ), contrast[0], contrast[1], same_on_batch) hue_factor = _adapted_uniform((batch_size, ), hue[0], hue[1], same_on_batch) saturation_factor = _adapted_uniform((batch_size, ), saturation[0], saturation[1], same_on_batch) return dict(brightness_factor=brightness_factor, contrast_factor=contrast_factor, hue_factor=hue_factor, saturation_factor=saturation_factor, order=torch.randperm(4))
def find_homography_dlt( points1: torch.Tensor, points2: torch.Tensor, weights: Optional[torch.Tensor] = None ) -> torch.Tensor: r"""Computes the homography matrix using the DLT formulation. The linear system is solved by using the Weighted Least Squares Solution for the 4 Points algorithm. Args: points1 (torch.Tensor): A set of points in the first image with a tensor shape :math:`(B, N, 2)`. points2 (torch.Tensor): A set of points in the second image with a tensor shape :math:`(B, N, 2)`. weights (torch.Tensor, optional): Tensor containing the weights per point correspondence with a shape of :math:`(B, N)`. Defaults to all ones. Returns: torch.Tensor: the computed homography matrix with shape :math:`(B, 3, 3)`. """ assert points1.shape == points2.shape, points1.shape assert len(points1.shape) >= 1 and points1.shape[-1] == 2, points1.shape assert points1.shape[1] >= 4, points1.shape device, dtype = _extract_device_dtype([points1, points2]) eps: float = 1e-8 points1_norm, transform1 = normalize_points(points1) points2_norm, transform2 = normalize_points(points2) x1, y1 = torch.chunk(points1_norm, dim=-1, chunks=2) # BxNx1 x2, y2 = torch.chunk(points2_norm, dim=-1, chunks=2) # BxNx1 ones, zeros = torch.ones_like(x1), torch.zeros_like(x1) # DIAPO 11: https://www.uio.no/studier/emner/matnat/its/nedlagte-emner/UNIK4690/v16/forelesninger/lecture_4_3-estimating-homographies-from-feature-correspondences.pdf # noqa: E501 ax = torch.cat([zeros, zeros, zeros, -x1, -y1, -ones, y2 * x1, y2 * y1, y2], dim=-1) ay = torch.cat([x1, y1, ones, zeros, zeros, zeros, -x2 * x1, -x2 * y1, -x2], dim=-1) A = torch.cat((ax, ay), dim=-1).reshape(ax.shape[0], -1, ax.shape[-1]) if weights is None: # All points are equally important A = A.transpose(-2, -1) @ A else: # We should use provided weights assert len(weights.shape) == 2 and weights.shape == points1.shape[:2], weights.shape w_diag = torch.diag_embed(weights.unsqueeze(dim=-1).repeat(1, 1, 2).reshape(weights.shape[0], -1)) A = A.transpose(-2, -1) @ w_diag @ A try: U, S, V = torch.svd(A) except: warnings.warn('SVD did not converge', RuntimeWarning) return torch.empty((points1_norm.size(0), 3, 3), device=device, dtype=dtype) H = V[..., -1].view(-1, 3, 3) H = transform2.inverse() @ (H @ transform1) H_norm = H / (H[..., -1:, -1:] + eps) return H_norm
def random_motion_blur_generator( batch_size: int, kernel_size: Union[int, Tuple[int, int]], angle: torch.Tensor, direction: torch.Tensor, same_on_batch: bool = False) -> Dict[str, torch.Tensor]: r"""Get parameters for motion blur. Args: batch_size (int): the tensor batch size. kernel_size (int or (int, int)): motion kernel size (odd and positive) or range. angle (torch.Tensor): angle of the motion blur in degrees (anti-clockwise rotation). direction (torch.Tensor): 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. same_on_batch (bool): apply the same transformation across the batch. Default: False. Returns: params Dict[str, torch.Tensor]: parameters to be passed for transformation. """ _common_param_check(batch_size, same_on_batch) _joint_range_check(angle, 'angle') _joint_range_check(direction, 'direction', (-1, 1)) device, dtype = _extract_device_dtype([angle, direction]) if isinstance(kernel_size, int): ksize_factor = torch.tensor([kernel_size] * batch_size, device=device, dtype=dtype) elif isinstance(kernel_size, tuple): # kernel_size is fixed across the batch assert len( kernel_size ) == 2, f"`kernel_size` must be (2,) if it is a tuple. Got {kernel_size}." ksize_factor = _adapted_uniform((batch_size, ), kernel_size[0] // 2, kernel_size[1] // 2, same_on_batch=True).int() * 2 + 1 else: raise TypeError(f"Unsupported type: {type(kernel_size)}") angle_factor = _adapted_uniform((batch_size, ), angle[0], angle[1], same_on_batch) direction_factor = _adapted_uniform((batch_size, ), direction[0], direction[1], same_on_batch) return dict(ksize_factor=ksize_factor.int(), angle_factor=angle_factor, direction_factor=direction_factor)
def random_motion_blur_generator3d( batch_size: int, kernel_size: Union[int, Tuple[int, int]], angle: torch.Tensor, direction: torch.Tensor, same_on_batch: bool = False) -> Dict[str, torch.Tensor]: r"""Get parameters for motion blur. Args: batch_size (int): the tensor batch size. kernel_size (int or (int, int)): motion kernel size (odd and positive) or range. angle (torch.Tensor): yaw, pitch and roll range of the motion blur in degrees :math:`(3, 2)`. direction (torch.Tensor): 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. same_on_batch (bool): apply the same transformation across the batch. Default: False. Returns: params Dict[str, torch.Tensor]: parameters to be passed for transformation. """ device, dtype = _extract_device_dtype([angle, direction]) if isinstance(kernel_size, int): ksize_factor = torch.tensor([kernel_size] * batch_size, device=device, dtype=dtype) elif isinstance(kernel_size, tuple): # kernel_size is fixed across the batch ksize_factor = _adapted_uniform((batch_size, ), kernel_size[0] // 2, kernel_size[1] // 2, same_on_batch=True).int() * 2 + 1 ksize_factor = ksize_factor.to(device=device, dtype=dtype) else: raise TypeError(f"Unsupported type: {type(kernel_size)}") assert angle.shape == torch.Size( [3, 2]), f"'angle' must be the shape of (3, 2). Got {angle.shape}." yaw = _adapted_uniform((batch_size, ), angle[0][0], angle[0][1], same_on_batch) pitch = _adapted_uniform((batch_size, ), angle[1][0], angle[1][1], same_on_batch) roll = _adapted_uniform((batch_size, ), angle[2][0], angle[2][1], same_on_batch) angle_factor = torch.stack([yaw, pitch, roll], dim=1) direction_factor = _adapted_uniform((batch_size, ), direction[0], direction[1], same_on_batch) return dict(ksize_factor=ksize_factor.int(), angle_factor=angle_factor, direction_factor=direction_factor)
def random_mixup_generator( batch_size: int, p: float = 0.5, lambda_val: 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"""Generate mixup indexes and lambdas for a batch of inputs. Args: batch_size (int): the number of images. If batchsize == 1, the output will be as same as the input. p (flot): probability of applying mixup. lambda_val (torch.Tensor, optional): min-max strength for mixup images, ranged from [0., 1.]. If None, it will be set to tensor([0., 1.]), which means no restrictions. 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. - mix_pairs (torch.Tensor): element-wise probabilities with a shape of (B,). - mixup_lambdas (torch.Tensor): element-wise probabilities with a shape of (B,). Note: The generated random numbers are not reproducible across different devices and dtypes. Examples: >>> rng = torch.manual_seed(0) >>> random_mixup_generator(5, 0.7) {'mixup_pairs': tensor([4, 0, 3, 1, 2]), 'mixup_lambdas': tensor([0.6323, 0.0000, 0.4017, 0.0223, 0.1689])} """ _common_param_check(batch_size, same_on_batch) _device, _dtype = _extract_device_dtype([lambda_val]) lambda_val = torch.as_tensor([0., 1.] if lambda_val is None else lambda_val, device=device, dtype=dtype) _joint_range_check(lambda_val, 'lambda_val', bounds=(0, 1)) batch_probs: torch.Tensor = random_prob_generator( batch_size, p, same_on_batch=same_on_batch, device=device, dtype=dtype ) mixup_pairs: torch.Tensor = torch.randperm(batch_size, device=device, dtype=dtype).long() mixup_lambdas: torch.Tensor = _adapted_uniform( (batch_size,), lambda_val[0], lambda_val[1], same_on_batch=same_on_batch) mixup_lambdas = mixup_lambdas * batch_probs return dict( mixup_pairs=mixup_pairs.to(device=_device, dtype=torch.long), mixup_lambdas=mixup_lambdas.to(device=_device, dtype=_dtype) )
def random_solarize_generator( batch_size: int, thresholds: torch.Tensor = torch.tensor([0.4, 0.6]), additions: torch.Tensor = torch.tensor([-0.1, 0.1]), same_on_batch: bool = False, device: torch.device = torch.device('cpu'), dtype: torch.dtype = torch.float32 ) -> Dict[str, torch.Tensor]: r"""Generate random solarize parameters for a batch of images. For each pixel in the image less than threshold, we add 'addition' amount to it and then clip the pixel value to be between 0 and 1.0 Args: batch_size (int): the number of images. thresholds (torch.Tensor): Pixels less than threshold will selected. Otherwise, subtract 1.0 from the pixel. Takes in a range tensor of (0, 1). Default value will be sampled from [0.4, 0.6]. additions (torch.Tensor): The value is between -0.5 and 0.5. Default value will be sampled from [-0.1, 0.1] 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. - thresholds_factor (torch.Tensor): element-wise thresholds factors with a shape of (B,). - additions_factor (torch.Tensor): element-wise additions factors 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(thresholds, 'thresholds', (0, 1)) _joint_range_check(additions, 'additions', (-0.5, 0.5)) _device, _dtype = _extract_device_dtype([thresholds, additions]) thresholds_factor = _adapted_uniform( (batch_size,), thresholds[0].to(device=device, dtype=dtype), thresholds[1].to(device=device, dtype=dtype), same_on_batch) additions_factor = _adapted_uniform( (batch_size,), additions[0].to(device=device, dtype=dtype), additions[1].to(device=device, dtype=dtype), same_on_batch) return dict( thresholds_factor=thresholds_factor.to(device=_device, dtype=_dtype), additions_factor=additions_factor.to(device=_device, dtype=_dtype) )
def _adapted_uniform(shape: Union[Tuple, torch.Size], low: Union[float, int, torch.Tensor], high: Union[float, int, torch.Tensor], same_on_batch=False) -> torch.Tensor: r"""The uniform sampling function that accepts 'same_on_batch'. If same_on_batch is True, all values generated will be exactly same given a batch_size (shape[0]). By default, same_on_batch is set to False. """ device, dtype = _extract_device_dtype([ low if isinstance(low, torch.Tensor) else None, high if isinstance(high, torch.Tensor) else None, ]) if not isinstance(low, torch.Tensor): low = torch.tensor(low, device=device, dtype=torch.float64) if not isinstance(high, torch.Tensor): high = torch.tensor(high, device=device, dtype=torch.float64) dist = Uniform(low.to(torch.float64), high.to(torch.float64)) return _adapted_rsampling(shape, dist, same_on_batch).to(device=device, dtype=dtype)
def _adapted_beta(shape: Union[Tuple, torch.Size], a: Union[float, int, torch.Tensor], b: Union[float, int, torch.Tensor], same_on_batch: bool = False) -> torch.Tensor: r""" The beta sampling function that accepts 'same_on_batch'. If same_on_batch is True, all values generated will be exactly same given a batch_size (shape[0]). By default, same_on_batch is set to False. By default, sampling happens on the default device and dtype. If a/b is a tensor, sampling will happen in the same device/dtype as a/b tensor. """ device, dtype = _extract_device_dtype([ a if isinstance(a, torch.Tensor) else None, b if isinstance(b, torch.Tensor) else None, ]) a = torch.as_tensor(a, device=device, dtype=dtype) b = torch.as_tensor(b, device=device, dtype=dtype) dist = Beta(a, b, validate_args=False) return _adapted_rsampling(shape, dist, same_on_batch)
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 random_cutmix_generator( batch_size: int, width: int, height: int, p: float = 0.5, num_mix: int = 1, beta: Optional[torch.Tensor] = None, cut_size: 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"""Generate cutmix indexes and lambdas for a batch of inputs. Args: batch_size (int): the number of images. If batchsize == 1, the output will be as same as the input. width (int): image width. height (int): image height. p (float): probability of applying cutmix. num_mix (int): number of images to mix with. Default is 1. beta (torch.Tensor, optional): hyperparameter for generating cut size from beta distribution. If None, it will be set to 1. cut_size (torch.Tensor, optional): controlling the minimum and maximum cut ratio from [0, 1]. If None, it will be set to [0, 1], which means no restriction. 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. - mix_pairs (torch.Tensor): element-wise probabilities with a shape of (num_mix, B). - crop_src (torch.Tensor): element-wise probabilities with a shape of (num_mix, B, 4, 2). Note: The generated random numbers are not reproducible across different devices and dtypes. Examples: >>> rng = torch.manual_seed(0) >>> random_cutmix_generator(3, 224, 224, p=0.5, num_mix=2) {'mix_pairs': tensor([[2, 0, 1], [1, 2, 0]]), 'crop_src': tensor([[[[ 35, 25], [208, 25], [208, 198], [ 35, 198]], <BLANKLINE> [[156, 137], [155, 137], [155, 136], [156, 136]], <BLANKLINE> [[ 3, 12], [210, 12], [210, 219], [ 3, 219]]], <BLANKLINE> <BLANKLINE> [[[ 83, 125], [177, 125], [177, 219], [ 83, 219]], <BLANKLINE> [[ 54, 8], [205, 8], [205, 159], [ 54, 159]], <BLANKLINE> [[ 97, 70], [ 96, 70], [ 96, 69], [ 97, 69]]]])} """ _device, _dtype = _extract_device_dtype([beta, cut_size]) beta = torch.as_tensor(1. if beta is None else beta, device=device, dtype=dtype) cut_size = torch.as_tensor([0., 1.] if cut_size is None else cut_size, device=device, dtype=dtype) assert num_mix >= 1 and isinstance(num_mix, (int,)), \ f"`num_mix` must be an integer greater than 1. Got {num_mix}." assert type(height) == int and height > 0 and type(width) == int and width > 0, \ f"'height' and 'width' must be integers. Got {height}, {width}." _joint_range_check(cut_size, 'cut_size', bounds=(0, 1)) _common_param_check(batch_size, same_on_batch) if batch_size == 0: return dict( mix_pairs=torch.zeros([0, 3], device=_device, dtype=torch.long), crop_src=torch.zeros([0, 4, 2], device=_device, dtype=torch.long) ) batch_probs: torch.Tensor = random_prob_generator( batch_size * num_mix, p, same_on_batch, device=device, dtype=dtype) mix_pairs: torch.Tensor = torch.rand(num_mix, batch_size, device=device, dtype=dtype).argsort(dim=1) cutmix_betas: torch.Tensor = _adapted_beta((batch_size * num_mix,), beta, beta, same_on_batch=same_on_batch) # Note: torch.clamp does not accept tensor, cutmix_betas.clamp(cut_size[0], cut_size[1]) throws: # Argument 1 to "clamp" of "_TensorBase" has incompatible type "Tensor"; expected "float" cutmix_betas = torch.min(torch.max(cutmix_betas, cut_size[0]), cut_size[1]) cutmix_rate = torch.sqrt(1. - cutmix_betas) * batch_probs cut_height = (cutmix_rate * height).long() cut_width = (cutmix_rate * width).long() _gen_shape = (1,) if same_on_batch: _gen_shape = (cut_height.size(0),) cut_height = cut_height[0] cut_width = cut_width[0] # Reserve at least 1 pixel for cropping. x_start = _adapted_uniform( _gen_shape, torch.zeros_like(cut_width, device=device, dtype=dtype), (width - cut_width - 1).to(device=device, dtype=dtype), same_on_batch).to(device=device, dtype=torch.long) y_start = _adapted_uniform( _gen_shape, torch.zeros_like(cut_height, device=device, dtype=dtype), (height - cut_height - 1).to(device=device, dtype=dtype), same_on_batch).to(device=device, dtype=torch.long) crop_src = bbox_generator(x_start.squeeze(), y_start.squeeze(), cut_width, cut_height) # (B * num_mix, 4, 2) => (num_mix, batch_size, 4, 2) crop_src = crop_src.view(num_mix, batch_size, 4, 2) return dict( mix_pairs=mix_pairs.to(device=_device, dtype=torch.long), crop_src=crop_src.to(device=_device, dtype=torch.long) )
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, same_on_batch: bool = False, device: torch.device = torch.device('cpu'), dtype: torch.dtype = torch.float32, ) -> Dict[str, torch.Tensor]: r"""Get parameters for ```3d affine``` transformation random affine transform. Args: batch_size (int): the tensor batch size. depth (int) : depth 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 (dx,dy,dz). 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. Shaped as (6, 2) for 6 facet (xy, xz, yx, yz, zx, zy). The shear to the i-th facet in the range (-shear[i, 0], shear[i, 1]) will be applied. same_on_batch (bool): apply the same transformation across the batch. Default: False Returns: params Dict[str, torch.Tensor]: parameters to be passed for transformation. - translations (torch.Tensor): element-wise translations with a shape of (B, 3). - center (torch.Tensor): element-wise center with a shape of (B, 3). - scale (torch.Tensor): element-wise scales with a shape of (B, 3). - angle (torch.Tensor): element-wise rotation angles with a shape of (B, 3). - sxy (torch.Tensor): element-wise x-y-facet shears with a shape of (B,). - sxz (torch.Tensor): element-wise x-z-facet shears with a shape of (B,). - syx (torch.Tensor): element-wise y-x-facet shears with a shape of (B,). - syz (torch.Tensor): element-wise y-z-facet shears with a shape of (B,). - szx (torch.Tensor): element-wise z-x-facet shears with a shape of (B,). - szy (torch.Tensor): element-wise z-y-facet shears with a shape of (B,). Note: The generated random numbers are not reproducible across different devices and dtypes. """ if not (type(depth) is int and depth > 0 and type(height) is int and height > 0 and type(width) is int and width > 0): raise AssertionError( f"'depth', 'height' and 'width' must be integers. Got {depth}, {height}, {width}." ) _device, _dtype = _extract_device_dtype( [degrees, translate, scale, shears]) if degrees.shape != torch.Size([3, 2]): raise AssertionError( f"'degrees' must be the shape of (3, 2). Got {degrees.shape}.") degrees = degrees.to(device=device, dtype=dtype) 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.stack([yaw, pitch, roll], dim=1) # compute tensor ranges if scale is not None: if scale.shape != torch.Size([3, 2]): raise AssertionError( f"'scale' must be the shape of (3, 2). Got {scale.shape}.") scale = scale.to(device=device, dtype=dtype) scale = torch.stack( [ _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, device=device, dtype=dtype).reshape(batch_size, 1).repeat(1, 3) if translate is not None: if translate.shape != torch.Size([3]): raise AssertionError( f"'translate' must be the shape of (2). Got {translate.shape}." ) translate = translate.to(device=device, dtype=dtype) max_dx: torch.Tensor = translate[0] * width max_dy: torch.Tensor = translate[1] * height max_dz: torch.Tensor = translate[2] * depth # translations should be in x,y,z 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), device=device, dtype=dtype) # center should be in x,y,z center: torch.Tensor = torch.tensor([width, height, depth], device=device, dtype=dtype).view(1, 3) / 2.0 - 0.5 center = center.expand(batch_size, -1) if shears is not None: if shears.shape != torch.Size([6, 2]): raise AssertionError( f"'shears' must be the shape of (6, 2). Got {shears.shape}.") shears = shears.to(device=device, dtype=dtype) 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, device=device, dtype=dtype) return dict( translations=translations.to(device=_device, dtype=_dtype), center=center.to(device=_device, dtype=_dtype), scale=scale.to(device=_device, dtype=_dtype), angles=angles.to(device=_device, dtype=_dtype), sxy=sxy.to(device=_device, dtype=_dtype), sxz=sxz.to(device=_device, dtype=_dtype), syx=syx.to(device=_device, dtype=_dtype), syz=syz.to(device=_device, dtype=_dtype), szx=szx.to(device=_device, dtype=_dtype), szy=szy.to(device=_device, dtype=_dtype), )
def random_perspective_generator3d( batch_size: int, depth: int, height: int, width: int, distortion_scale: torch.Tensor, same_on_batch: bool = False, device: torch.device = torch.device('cpu'), dtype: torch.dtype = torch.float32, ) -> Dict[str, torch.Tensor]: r"""Get parameters for ``perspective`` for a random perspective transform. Args: batch_size (int): the tensor batch size. depth (int) : depth of the image. height (int) : height of the image. width (int): width of the image. distortion_scale (torch.Tensor): it controls the degree of distortion and ranges from 0 to 1. 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. - src (torch.Tensor): perspective source bounding boxes with a shape of (B, 8, 3). - dst (torch.Tensor): perspective target bounding boxes with a shape (B, 8, 3). Note: The generated random numbers are not reproducible across different devices and dtypes. """ if not (distortion_scale.dim() == 0 and 0 <= distortion_scale <= 1): raise AssertionError( f"'distortion_scale' must be a scalar within [0, 1]. Got {distortion_scale}" ) _device, _dtype = _extract_device_dtype([distortion_scale]) distortion_scale = distortion_scale.to(device=device, dtype=dtype) start_points: torch.Tensor = torch.tensor( [[ [0.0, 0, 0], [width - 1, 0, 0], [width - 1, height - 1, 0], [0, height - 1, 0], [0.0, 0, depth - 1], [width - 1, 0, depth - 1], [width - 1, height - 1, depth - 1], [0, height - 1, depth - 1], ]], device=device, dtype=dtype, ).expand(batch_size, -1, -1) # generate random offset not larger than half of the image fx = distortion_scale * width / 2 fy = distortion_scale * height / 2 fz = distortion_scale * depth / 2 factor = torch.stack([fx, fy, fz], dim=0).view(-1, 1, 3) rand_val: torch.Tensor = _adapted_uniform( start_points.shape, torch.tensor(0, device=device, dtype=dtype), torch.tensor(1, device=device, dtype=dtype), same_on_batch, ) pts_norm = torch.tensor( [[[1, 1, 1], [-1, 1, 1], [-1, -1, 1], [1, -1, 1], [1, 1, -1], [-1, 1, -1], [-1, -1, -1], [1, -1, -1]]], device=device, dtype=dtype, ) end_points = start_points + factor * rand_val * pts_norm return dict( start_points=start_points.to(device=_device, dtype=_dtype), end_points=end_points.to(device=_device, dtype=_dtype), )
def random_crop_generator3d( batch_size: int, input_size: Tuple[int, int, int], size: Union[Tuple[int, int, int], torch.Tensor], resize_to: Optional[Tuple[int, int, int]] = 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 ```crop``` transformation for crop transform. Args: batch_size (int): the tensor batch size. input_size (tuple): Input image shape, like (d, h, w). size (tuple): Desired size of the crop operation, like (d, h, w). If tensor, it must be (B, 3). resize_to (tuple): Desired output size of the crop, like (d, h, w). If None, no resize will be performed. 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. - src (torch.Tensor): cropping bounding boxes with a shape of (B, 8, 3). - dst (torch.Tensor): output bounding boxes with a shape (B, 8, 3). Note: The generated random numbers are not reproducible across different devices and dtypes. """ _device, _dtype = _extract_device_dtype( [size if isinstance(size, torch.Tensor) else None]) if not isinstance(size, torch.Tensor): size = torch.tensor(size, device=device, dtype=dtype).repeat(batch_size, 1) else: size = size.to(device=device, dtype=dtype) if size.shape != torch.Size([batch_size, 3]): raise AssertionError( "If `size` is a tensor, it must be shaped as (B, 3). " f"Got {size.shape} while expecting {torch.Size([batch_size, 3])}.") if not (len(input_size) == 3 and isinstance(input_size[0], (int, )) and isinstance(input_size[1], (int, )) and isinstance(input_size[2], (int, )) and input_size[0] > 0 and input_size[1] > 0 and input_size[2] > 0): raise AssertionError( f"`input_size` must be a tuple of 3 positive integers. Got {input_size}." ) x_diff = input_size[2] - size[:, 2] + 1 y_diff = input_size[1] - size[:, 1] + 1 z_diff = input_size[0] - size[:, 0] + 1 if (x_diff < 0).any() or (y_diff < 0).any() or (z_diff < 0).any(): raise ValueError( f"input_size {str(input_size)} cannot be smaller than crop size {str(size)} in any dimension." ) if batch_size == 0: return dict( src=torch.zeros([0, 8, 3], device=_device, dtype=_dtype), dst=torch.zeros([0, 8, 3], device=_device, dtype=_dtype), ) if same_on_batch: # If same_on_batch, select the first then repeat. x_start = _adapted_uniform((batch_size, ), 0, x_diff[0], same_on_batch).floor() y_start = _adapted_uniform((batch_size, ), 0, y_diff[0], same_on_batch).floor() z_start = _adapted_uniform((batch_size, ), 0, z_diff[0], same_on_batch).floor() else: x_start = _adapted_uniform((1, ), 0, x_diff, same_on_batch).floor() y_start = _adapted_uniform((1, ), 0, y_diff, same_on_batch).floor() z_start = _adapted_uniform((1, ), 0, z_diff, same_on_batch).floor() crop_src = bbox_generator3d( x_start.to(device=_device, dtype=_dtype).view(-1), y_start.to(device=_device, dtype=_dtype).view(-1), z_start.to(device=_device, dtype=_dtype).view(-1), size[:, 2].to(device=_device, dtype=_dtype) - 1, size[:, 1].to(device=_device, dtype=_dtype) - 1, size[:, 0].to(device=_device, dtype=_dtype) - 1, ) if resize_to is None: crop_dst = bbox_generator3d( torch.tensor([0] * batch_size, device=_device, dtype=_dtype), torch.tensor([0] * batch_size, device=_device, dtype=_dtype), torch.tensor([0] * batch_size, device=_device, dtype=_dtype), size[:, 2].to(device=_device, dtype=_dtype) - 1, size[:, 1].to(device=_device, dtype=_dtype) - 1, size[:, 0].to(device=_device, dtype=_dtype) - 1, ) else: if not (len(resize_to) == 3 and isinstance(resize_to[0], (int, )) and isinstance(resize_to[1], (int, )) and isinstance(resize_to[2], (int, )) and resize_to[0] > 0 and resize_to[1] > 0 and resize_to[2] > 0): raise AssertionError( f"`resize_to` must be a tuple of 3 positive integers. Got {resize_to}." ) crop_dst = torch.tensor( [[ [0, 0, 0], [resize_to[-1] - 1, 0, 0], [resize_to[-1] - 1, resize_to[-2] - 1, 0], [0, resize_to[-2] - 1, 0], [0, 0, resize_to[-3] - 1], [resize_to[-1] - 1, 0, resize_to[-3] - 1], [resize_to[-1] - 1, resize_to[-2] - 1, resize_to[-3] - 1], [0, resize_to[-2] - 1, resize_to[-3] - 1], ]], device=_device, dtype=_dtype, ).repeat(batch_size, 1, 1) return dict(src=crop_src.to(device=_device), dst=crop_dst.to(device=_device))
def random_motion_blur_generator3d( batch_size: int, kernel_size: Union[int, Tuple[int, int]], angle: torch.Tensor, direction: torch.Tensor, same_on_batch: bool = False, device: torch.device = torch.device('cpu'), dtype: torch.dtype = torch.float32, ) -> Dict[str, torch.Tensor]: r"""Get parameters for motion blur. Args: batch_size (int): the tensor batch size. kernel_size (int or (int, int)): motion kernel size (odd and positive) or range. angle (torch.Tensor): yaw, pitch and roll range of the motion blur in degrees :math:`(3, 2)`. direction (torch.Tensor): 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. 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. - ksize_factor (torch.Tensor): element-wise kernel size factors with a shape of (B,). - angle_factor (torch.Tensor): element-wise center with a shape of (B,). - direction_factor (torch.Tensor): element-wise scales with a shape of (B,). Note: The generated random numbers are not reproducible across different devices and dtypes. """ _device, _dtype = _extract_device_dtype([angle, direction]) _joint_range_check(direction, 'direction', (-1, 1)) if isinstance(kernel_size, int): if not (kernel_size >= 3 and kernel_size % 2 == 1): raise AssertionError( f"`kernel_size` must be odd and greater than 3. Got {kernel_size}." ) ksize_factor = torch.tensor([kernel_size] * batch_size, device=device, dtype=dtype).int() elif isinstance(kernel_size, tuple): if not (len(kernel_size) == 2 and kernel_size[0] >= 3 and kernel_size[0] <= kernel_size[1]): raise AssertionError( f"`kernel_size` must be greater than 3. Got range {kernel_size}." ) # kernel_size is fixed across the batch ksize_factor = (_adapted_uniform((batch_size, ), kernel_size[0] // 2, kernel_size[1] // 2, same_on_batch=True).int() * 2 + 1) else: raise TypeError(f"Unsupported type: {type(kernel_size)}") if angle.shape != torch.Size([3, 2]): raise AssertionError( f"'angle' must be the shape of (3, 2). Got {angle.shape}.") angle = angle.to(device=device, dtype=dtype) yaw = _adapted_uniform((batch_size, ), angle[0][0], angle[0][1], same_on_batch) pitch = _adapted_uniform((batch_size, ), angle[1][0], angle[1][1], same_on_batch) roll = _adapted_uniform((batch_size, ), angle[2][0], angle[2][1], same_on_batch) angle_factor = torch.stack([yaw, pitch, roll], dim=1) direction = direction.to(device=device, dtype=dtype) direction_factor = _adapted_uniform((batch_size, ), direction[0], direction[1], same_on_batch) return dict( ksize_factor=ksize_factor.to(device=_device), angle_factor=angle_factor.to(device=_device, dtype=_dtype), direction_factor=direction_factor.to(device=_device, dtype=_dtype), )
def get_motion_kernel3d( kernel_size: int, angle: Union[torch.Tensor, Tuple[float, float, float]], direction: Union[torch.Tensor, float] = 0.0, mode: str = 'nearest', ) -> torch.Tensor: r"""Return 3D motion blur filter. Args: kernel_size (int): motion kernel width, height and depth. It should be odd and positive. angle (tensor or tuple): Range of yaw (x-axis), pitch (y-axis), roll (z-axis) to select from. If tensor, it must be :math:`(B, 3)`. If tuple, it must be (yaw, pitch, raw). direction (float): 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. mode (str): interpolation mode for rotating the kernel. ``'bilinear'`` or ``'nearest'``. Default: ``'nearest'``. Returns: torch.Tensor: the motion blur kernel. Shape: - Output: :math:`(B, kernel_size, kernel_size, kernel_size)` Examples: >>> get_motion_kernel3d(3, (0., 0., 0.), 0.) tensor([[[[0.0000, 0.0000, 0.0000], [0.0000, 0.0000, 0.0000], [0.0000, 0.0000, 0.0000]], <BLANKLINE> [[0.0000, 0.0000, 0.0000], [0.3333, 0.3333, 0.3333], [0.0000, 0.0000, 0.0000]], <BLANKLINE> [[0.0000, 0.0000, 0.0000], [0.0000, 0.0000, 0.0000], [0.0000, 0.0000, 0.0000]]]]) >>> get_motion_kernel3d(3, (90., 90., 0.), -0.5) tensor([[[[0.0000, 0.0000, 0.0000], [0.0000, 0.0000, 0.0000], [0.0000, 0.5000, 0.0000]], <BLANKLINE> [[0.0000, 0.0000, 0.0000], [0.0000, 0.3333, 0.0000], [0.0000, 0.0000, 0.0000]], <BLANKLINE> [[0.0000, 0.1667, 0.0000], [0.0000, 0.0000, 0.0000], [0.0000, 0.0000, 0.0000]]]]) """ if not isinstance(kernel_size, int) or kernel_size % 2 == 0 or kernel_size < 3: raise TypeError(f"ksize must be an odd integer >= than 3. Got {kernel_size}.") device, dtype = _extract_device_dtype( [angle if isinstance(angle, torch.Tensor) else None, direction if isinstance(direction, torch.Tensor) else None] ) if not isinstance(angle, torch.Tensor): angle = torch.tensor([angle], device=device, dtype=dtype) angle = cast(torch.Tensor, angle) if angle.dim() == 1: angle = angle.unsqueeze(0) if not (len(angle.shape) == 2 and angle.size(1) == 3): raise AssertionError(f"angle must be (B, 3). Got {angle}.") if not isinstance(direction, torch.Tensor): direction = torch.tensor([direction], device=device, dtype=dtype) direction = cast(torch.Tensor, direction) if direction.dim() == 0: direction = direction.unsqueeze(0) if direction.dim() != 1: raise AssertionError(f"direction must be a 1-dim tensor. Got {direction}.") if direction.size(0) != angle.size(0): raise AssertionError(f"direction and angle must have the same length. Got {direction} and {angle}.") kernel_tuple: Tuple[int, int, int] = (kernel_size, kernel_size, kernel_size) # direction from [-1, 1] to [0, 1] range direction = (torch.clamp(direction, -1.0, 1.0) + 1.0) / 2.0 kernel = torch.zeros((direction.size(0), *kernel_tuple), device=device, dtype=dtype) # Element-wise linspace # kernel[:, kernel_size // 2, kernel_size // 2, :] = torch.stack( # [(direction + ((1 - 2 * direction) / (kernel_size - 1)) * i) for i in range(kernel_size)], dim=-1) k = torch.stack([(direction + ((1 - 2 * direction) / (kernel_size - 1)) * i) for i in range(kernel_size)], dim=-1) kernel = torch.nn.functional.pad( k[:, None, None], [0, 0, kernel_size // 2, kernel_size // 2, kernel_size // 2, kernel_size // 2, 0, 0] ) if kernel.shape != torch.Size([direction.size(0), *kernel_tuple]): raise AssertionError kernel = kernel.unsqueeze(1) # rotate (counterclockwise) kernel by given angle kernel = rotate3d(kernel, angle[:, 0], angle[:, 1], angle[:, 2], mode=mode, align_corners=True) kernel = kernel[:, 0] kernel = kernel / kernel.sum(dim=(1, 2, 3), keepdim=True) return kernel
def get_motion_kernel2d( kernel_size: int, angle: Union[torch.Tensor, float], direction: Union[torch.Tensor, float] = 0.0, mode: str = 'nearest', ) -> torch.Tensor: r"""Return 2D motion blur filter. Args: kernel_size (int): motion kernel width and height. It should be odd and positive. angle (torch.Tensor, float): angle of the motion blur in degrees (anti-clockwise rotation). direction (float): 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. mode (str): interpolation mode for rotating the kernel. ``'bilinear'`` or ``'nearest'``. Default: ``'nearest'`` Returns: torch.Tensor: the motion blur kernel. Shape: - Output: :math:`(B, ksize, ksize)` Examples:: >>> get_motion_kernel2d(5, 0., 0.) tensor([[[0.0000, 0.0000, 0.0000, 0.0000, 0.0000], [0.0000, 0.0000, 0.0000, 0.0000, 0.0000], [0.2000, 0.2000, 0.2000, 0.2000, 0.2000], [0.0000, 0.0000, 0.0000, 0.0000, 0.0000], [0.0000, 0.0000, 0.0000, 0.0000, 0.0000]]]) >>> get_motion_kernel2d(3, 215., -0.5) tensor([[[0.0000, 0.0000, 0.1667], [0.0000, 0.3333, 0.0000], [0.5000, 0.0000, 0.0000]]]) """ device, dtype = _extract_device_dtype( [angle if isinstance(angle, torch.Tensor) else None, direction if isinstance(direction, torch.Tensor) else None] ) if not isinstance(kernel_size, int) or kernel_size % 2 == 0 or kernel_size < 3: raise TypeError("ksize must be an odd integer >= than 3") if not isinstance(angle, torch.Tensor): angle = torch.tensor([angle], device=device, dtype=dtype) angle = cast(torch.Tensor, angle) if angle.dim() == 0: angle = angle.unsqueeze(0) if angle.dim() != 1: raise AssertionError(f"angle must be a 1-dim tensor. Got {angle}.") if not isinstance(direction, torch.Tensor): direction = torch.tensor([direction], device=device, dtype=dtype) direction = cast(torch.Tensor, direction) if direction.dim() == 0: direction = direction.unsqueeze(0) if direction.dim() != 1: raise AssertionError(f"direction must be a 1-dim tensor. Got {direction}.") if direction.size(0) != angle.size(0): raise AssertionError(f"direction and angle must have the same length. Got {direction} and {angle}.") kernel_tuple: Tuple[int, int] = (kernel_size, kernel_size) # direction from [-1, 1] to [0, 1] range direction = (torch.clamp(direction, -1.0, 1.0) + 1.0) / 2.0 # kernel = torch.zeros((direction.size(0), *kernel_tuple), device=device, dtype=dtype) # Element-wise linspace # kernel[:, kernel_size // 2, :] = torch.stack( # [(direction + ((1 - 2 * direction) / (kernel_size - 1)) * i) for i in range(kernel_size)], dim=-1) # Alternatively # m = ((1 - 2 * direction)[:, None].repeat(1, kernel_size) / (kernel_size - 1)) # kernel[:, kernel_size // 2, :] = direction[:, None].repeat(1, kernel_size) + m * torch.arange(0, kernel_size) k = torch.stack([(direction + ((1 - 2 * direction) / (kernel_size - 1)) * i) for i in range(kernel_size)], dim=-1) kernel = torch.nn.functional.pad(k[:, None], [0, 0, kernel_size // 2, kernel_size // 2, 0, 0]) if kernel.shape != torch.Size([direction.size(0), *kernel_tuple]): raise AssertionError kernel = kernel.unsqueeze(1) # rotate (counterclockwise) kernel by given angle kernel = rotate(kernel, angle, mode=mode, align_corners=True) kernel = kernel[:, 0] kernel = kernel / kernel.sum(dim=(1, 2), keepdim=True) return kernel
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") assert isinstance(width, (int,)) and isinstance(height, (int,)) and width > 0 and height > 0, \ 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) assert len(scale.shape) == 1 and (len(scale) == 2 or len(scale) == 4), \ 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) _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) 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.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)
def random_crop_generator( batch_size: int, input_size: Tuple[int, int], size: Union[Tuple[int, int], torch.Tensor], resize_to: Optional[Tuple[int, int]] = 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 ```crop``` transformation for crop transform. Args: batch_size (int): the tensor batch size. input_size (tuple): Input image shape, like (h, w). size (tuple): Desired size of the crop operation, like (h, w). If tensor, it must be (B, 2). resize_to (tuple): Desired output size of the crop, like (h, w). If None, no resize will be performed. 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. - src (torch.Tensor): cropping bounding boxes with a shape of (B, 4, 2). - dst (torch.Tensor): output bounding boxes with a shape (B, 4, 2). Note: The generated random numbers are not reproducible across different devices and dtypes. Example: >>> _ = torch.manual_seed(0) >>> crop_size = torch.tensor([[25, 28], [27, 29], [26, 28]]) >>> random_crop_generator(3, (30, 30), size=crop_size, same_on_batch=False) {'src': tensor([[[ 1, 0], [28, 0], [28, 24], [ 1, 24]], <BLANKLINE> [[ 1, 1], [29, 1], [29, 27], [ 1, 27]], <BLANKLINE> [[ 0, 3], [27, 3], [27, 28], [ 0, 28]]]), 'dst': tensor([[[ 0, 0], [27, 0], [27, 24], [ 0, 24]], <BLANKLINE> [[ 0, 0], [28, 0], [28, 26], [ 0, 26]], <BLANKLINE> [[ 0, 0], [27, 0], [27, 25], [ 0, 25]]])} """ _common_param_check(batch_size, same_on_batch) _device, _dtype = _extract_device_dtype([size if isinstance(size, torch.Tensor) else None]) if not isinstance(size, torch.Tensor): size = torch.tensor(size, device=device, dtype=dtype).repeat(batch_size, 1) else: size = size.to(device=device, dtype=dtype) assert size.shape == torch.Size([batch_size, 2]), ( "If `size` is a tensor, it must be shaped as (B, 2). " f"Got {size.shape} while expecting {torch.Size([batch_size, 2])}.") size = size.long() x_diff = input_size[1] - size[:, 1] + 1 y_diff = input_size[0] - size[:, 0] + 1 if (x_diff < 0).any() or (y_diff < 0).any(): raise ValueError("input_size %s cannot be smaller than crop size %s in any dimension." % (str(input_size), str(size))) if batch_size == 0: return dict( src=torch.zeros([0, 4, 2], device=_device, dtype=torch.long), dst=torch.zeros([0, 4, 2], device=_device, dtype=torch.long), ) if same_on_batch: # If same_on_batch, select the first then repeat. x_start = _adapted_uniform((batch_size,), 0, x_diff[0].to(device=device, dtype=dtype), same_on_batch).long() y_start = _adapted_uniform((batch_size,), 0, y_diff[0].to(device=device, dtype=dtype), same_on_batch).long() else: x_start = _adapted_uniform((1,), 0, x_diff.to(device=device, dtype=dtype), same_on_batch).long() y_start = _adapted_uniform((1,), 0, y_diff.to(device=device, dtype=dtype), same_on_batch).long() crop_src = bbox_generator( x_start.view(-1), y_start.view(-1), size[:, 1], size[:, 0]).to(device=_device, dtype=torch.long) if resize_to is None: crop_dst = bbox_generator( torch.tensor([0] * batch_size, device=device, dtype=torch.long), torch.tensor([0] * batch_size, device=device, dtype=torch.long), size[:, 1], size[:, 0]).to(device=_device, dtype=torch.long) else: assert len(resize_to) == 2 and isinstance(resize_to[0], (int,)) and isinstance(resize_to[1], (int,)) \ and resize_to[0] > 0 and resize_to[1] > 0, \ f"`resize_to` must be a tuple of 2 positive integers. Got {resize_to}." crop_dst = torch.tensor([[ [0, 0], [resize_to[1] - 1, 0], [resize_to[1] - 1, resize_to[0] - 1], [0, resize_to[0] - 1], ]], device=_device, dtype=torch.long).repeat(batch_size, 1, 1) return dict(src=crop_src, dst=crop_dst)
def random_crop_size_generator( batch_size: int, size: Tuple[int, int], scale: torch.Tensor, ratio: torch.Tensor, same_on_batch: bool = False, device: torch.device = torch.device('cpu'), dtype: torch.dtype = torch.float32 ) -> Dict[str, torch.Tensor]: r"""Get cropping heights and widths for ```crop``` transformation for resized crop transform. Args: batch_size (int): the tensor batch size. size (Tuple[int, int]): expected output size of each edge. scale (torch.Tensor): range of size of the origin size cropped with (2,) shape. ratio (torch.Tensor): range of aspect ratio of the origin aspect ratio cropped with (2,) shape. 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. - size (torch.Tensor): element-wise cropping sizes with a shape of (B, 2). Note: The generated random numbers are not reproducible across different devices and dtypes. Examples: >>> _ = torch.manual_seed(42) >>> random_crop_size_generator(3, (30, 30), scale=torch.tensor([.7, 1.3]), ratio=torch.tensor([.9, 1.])) {'size': tensor([[29, 29], [27, 28], [26, 29]])} """ _common_param_check(batch_size, same_on_batch) _joint_range_check(scale, "scale") _joint_range_check(ratio, "ratio") assert len(size) == 2 and type(size[0]) == int and size[1] > 0 and type(size[1]) == int and size[1] > 0, \ f"'height' and 'width' must be integers. Got {size}." _device, _dtype = _extract_device_dtype([scale, ratio]) if batch_size == 0: return dict(size=torch.zeros([0, 2], device=_device, dtype=_dtype)) scale = scale.to(device=device, dtype=dtype) ratio = ratio.to(device=device, dtype=dtype) # 10 trails for each element area = _adapted_uniform( (batch_size, 10), scale[0] * size[0] * size[1], scale[1] * size[0] * size[1], same_on_batch) log_ratio = _adapted_uniform( (batch_size, 10), torch.log(ratio[0]), torch.log(ratio[1]), same_on_batch) aspect_ratio = torch.exp(log_ratio) w = torch.sqrt(area * aspect_ratio).round().int() h = torch.sqrt(area / aspect_ratio).round().int() # Element-wise w, h condition cond = ((0 < w) * (w < size[1]) * (0 < h) * (h < size[0])).int() # torch.argmax is not reproducible accross devices: https://github.com/pytorch/pytorch/issues/17738 # Here, we will select the first occurance of the duplicated elements. cond_bool, argmax_dim1 = ((cond.cumsum(1) == 1) & cond.bool()).max(1) h_out = w[torch.arange(0, batch_size, device=device, dtype=torch.long), argmax_dim1] w_out = h[torch.arange(0, batch_size, device=device, dtype=torch.long), argmax_dim1] if not cond_bool.all(): # Fallback to center crop in_ratio = float(size[0]) / float(size[1]) if (in_ratio < min(ratio)): h_ct = torch.tensor(size[0], device=device, dtype=dtype) w_ct = torch.round(h_ct / min(ratio)) elif (in_ratio > max(ratio)): w_ct = torch.tensor(size[1], device=device, dtype=dtype) h_ct = torch.round(w_ct * max(ratio)) else: # whole image h_ct = torch.tensor(size[0], device=device, dtype=dtype) w_ct = torch.tensor(size[1], device=device, dtype=dtype) h_ct = h_ct.int() w_ct = w_ct.int() h_out = h_out.where(cond_bool, h_ct) w_out = w_out.where(cond_bool, w_ct) return dict(size=torch.stack([h_out, w_out], dim=1).to(device=_device, dtype=torch.long))
def random_color_jitter_generator( batch_size: int, brightness: Optional[torch.Tensor] = None, contrast: Optional[torch.Tensor] = None, saturation: Optional[torch.Tensor] = None, hue: 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"""Generate random color jiter parameters for a batch of images. Args: batch_size (int): the number of images. brightness (torch.Tensor, optional): Brightness factor tensor of range (a, b). The provided range must follow 0 <= a <= b <= 2. Default value is [0., 0.]. contrast (torch.Tensor, optional): Contrast factor tensor of range (a, b). The provided range must follow 0 <= a <= b. Default value is [0., 0.]. saturation (torch.Tensor, optional): Saturation factor tensor of range (a, b). The provided range must follow 0 <= a <= b. Default value is [0., 0.]. hue (torch.Tensor, optional): Saturation factor tensor of range (a, b). The provided range must follow -0.5 <= a <= b < 0.5. Default value is [0., 0.]. 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. - brightness_factor (torch.Tensor): element-wise brightness factors with a shape of (B,). - contrast_factor (torch.Tensor): element-wise contrast factors with a shape of (B,). - hue_factor (torch.Tensor): element-wise hue factors with a shape of (B,). - saturation_factor (torch.Tensor): element-wise saturation factors with a shape of (B,). - order (torch.Tensor): applying orders of the color adjustments with a shape of (4). In which, 0 is brightness adjustment; 1 is contrast adjustment; 2 is saturation adjustment; 3 is hue adjustment. Note: The generated random numbers are not reproducible across different devices and dtypes. """ _common_param_check(batch_size, same_on_batch) _device, _dtype = _extract_device_dtype([brightness, contrast, hue, saturation]) brightness = torch.as_tensor([0., 0.] if brightness is None else brightness, device=device, dtype=dtype) contrast = torch.as_tensor([0., 0.] if contrast is None else contrast, device=device, dtype=dtype) hue = torch.as_tensor([0., 0.] if hue is None else hue, device=device, dtype=dtype) saturation = torch.as_tensor([0., 0.] if saturation is None else saturation, device=device, dtype=dtype) _joint_range_check(brightness, "brightness", (0, 2)) _joint_range_check(contrast, "contrast", (0, float('inf'))) _joint_range_check(hue, "hue", (-0.5, 0.5)) _joint_range_check(saturation, "saturation", (0, float('inf'))) brightness_factor = _adapted_uniform((batch_size,), brightness[0], brightness[1], same_on_batch) contrast_factor = _adapted_uniform((batch_size,), contrast[0], contrast[1], same_on_batch) hue_factor = _adapted_uniform((batch_size,), hue[0], hue[1], same_on_batch) saturation_factor = _adapted_uniform((batch_size,), saturation[0], saturation[1], same_on_batch) return dict(brightness_factor=brightness_factor.to(device=_device, dtype=_dtype), contrast_factor=contrast_factor.to(device=_device, dtype=_dtype), hue_factor=hue_factor.to(device=_device, dtype=_dtype), saturation_factor=saturation_factor.to(device=_device, dtype=_dtype), order=torch.randperm(4, device=_device, dtype=_dtype).long())
def random_rectangles_params_generator( batch_size: int, height: int, width: int, scale: torch.Tensor, ratio: torch.Tensor, value: float = 0., same_on_batch: bool = False, device: torch.device = torch.device('cpu'), dtype: torch.dtype = torch.float32 ) -> Dict[str, torch.Tensor]: r"""Get parameters for ```erasing``` transformation for erasing transform. Args: batch_size (int): the tensor batch size. height (int) : height of the image. width (int): width of the image. scale (torch.Tensor): range of size of the origin size cropped. Shape (2). ratio (torch.Tensor): range of aspect ratio of the origin aspect ratio cropped. Shape (2). value (float): value to be filled in the erased area. 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. - widths (torch.Tensor): element-wise erasing widths with a shape of (B,). - heights (torch.Tensor): element-wise erasing heights with a shape of (B,). - xs (torch.Tensor): element-wise erasing x coordinates with a shape of (B,). - ys (torch.Tensor): element-wise erasing y coordinates with a shape of (B,). - values (torch.Tensor): element-wise filling values 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) _device, _dtype = _extract_device_dtype([ratio, scale]) assert type(height) == int and height > 0 and type(width) == int and width > 0, \ f"'height' and 'width' must be integers. Got {height}, {width}." assert isinstance(value, (int, float)) and value >= 0 and value <= 1, \ f"'value' must be a number between 0 - 1. Got {value}." _joint_range_check(scale, 'scale', bounds=(0, float('inf'))) _joint_range_check(ratio, 'ratio', bounds=(0, float('inf'))) images_area = height * width target_areas = _adapted_uniform((batch_size,), scale[0].to(device=device, dtype=dtype), scale[1].to(device=device, dtype=dtype), same_on_batch) * images_area if ratio[0] < 1. and ratio[1] > 1.: aspect_ratios1 = _adapted_uniform((batch_size,), ratio[0].to(device=device, dtype=dtype), 1, same_on_batch) aspect_ratios2 = _adapted_uniform((batch_size,), 1, ratio[1].to(device=device, dtype=dtype), same_on_batch) if same_on_batch: rand_idxs = torch.round(_adapted_uniform( (1,), torch.tensor(0, device=device, dtype=dtype), torch.tensor(1, device=device, dtype=dtype), same_on_batch)).repeat(batch_size).bool() else: rand_idxs = torch.round(_adapted_uniform( (batch_size,), torch.tensor(0, device=device, dtype=dtype), torch.tensor(1, device=device, dtype=dtype), same_on_batch)).bool() aspect_ratios = torch.where(rand_idxs, aspect_ratios1, aspect_ratios2) else: aspect_ratios = _adapted_uniform( (batch_size,), ratio[0].to(device=device, dtype=dtype), ratio[1].to(device=device, dtype=dtype), same_on_batch) # based on target areas and aspect ratios, rectangle params are computed heights = torch.min( torch.max(torch.round((target_areas * aspect_ratios) ** (1 / 2)), torch.tensor(1., device=device, dtype=dtype)), torch.tensor(height, device=device, dtype=dtype) ) widths = torch.min( torch.max(torch.round((target_areas / aspect_ratios) ** (1 / 2)), torch.tensor(1., device=device, dtype=dtype)), torch.tensor(width, device=device, dtype=dtype) ) xs_ratio = _adapted_uniform((batch_size,), torch.tensor(0, device=device, dtype=dtype), torch.tensor(1, device=device, dtype=dtype), same_on_batch) ys_ratio = _adapted_uniform((batch_size,), torch.tensor(0, device=device, dtype=dtype), torch.tensor(1, device=device, dtype=dtype), same_on_batch) xs = xs_ratio * (torch.tensor(width, device=device, dtype=dtype) - widths + 1) ys = ys_ratio * (torch.tensor(height, device=device, dtype=dtype) - heights + 1) return dict(widths=widths.to(device=_device, dtype=torch.int32), heights=heights.to(device=_device, dtype=torch.int32), xs=xs.to(device=_device, dtype=torch.int32), ys=ys.to(device=_device, dtype=torch.int32), values=torch.tensor([value] * batch_size, device=_device, dtype=_dtype))
def random_motion_blur_generator( batch_size: int, kernel_size: Union[int, Tuple[int, int]], angle: torch.Tensor, direction: torch.Tensor, same_on_batch: bool = False, device: torch.device = torch.device('cpu'), dtype: torch.dtype = torch.float32 ) -> Dict[str, torch.Tensor]: r"""Get parameters for motion blur. Args: batch_size (int): the tensor batch size. kernel_size (int or (int, int)): motion kernel size (odd and positive) or range. angle (torch.Tensor): angle of the motion blur in degrees (anti-clockwise rotation). direction (torch.Tensor): 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. 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. - ksize_factor (torch.Tensor): element-wise kernel size factors with a shape of (B,). - angle_factor (torch.Tensor): element-wise angle factors with a shape of (B,). - direction_factor (torch.Tensor): element-wise direction factors 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(angle, 'angle') _joint_range_check(direction, 'direction', (-1, 1)) _device, _dtype = _extract_device_dtype([angle, direction]) if isinstance(kernel_size, int): assert kernel_size >= 3 and kernel_size % 2 == 1, \ f"`kernel_size` must be odd and greater than 3. Got {kernel_size}." ksize_factor = torch.tensor([kernel_size] * batch_size, device=device, dtype=dtype) elif isinstance(kernel_size, tuple): # kernel_size is fixed across the batch assert len(kernel_size) == 2, f"`kernel_size` must be (2,) if it is a tuple. Got {kernel_size}." ksize_factor = _adapted_uniform( (batch_size,), kernel_size[0] // 2, kernel_size[1] // 2, same_on_batch=True).int() * 2 + 1 else: raise TypeError(f"Unsupported type: {type(kernel_size)}") angle_factor = _adapted_uniform( (batch_size,), angle[0].to(device=device, dtype=dtype), angle[1].to(device=device, dtype=dtype), same_on_batch) direction_factor = _adapted_uniform( (batch_size,), direction[0].to(device=device, dtype=dtype), direction[1].to(device=device, dtype=dtype), same_on_batch) return dict(ksize_factor=ksize_factor.to(device=_device, dtype=torch.int32), angle_factor=angle_factor.to(device=_device, dtype=_dtype), direction_factor=direction_factor.to(device=_device, dtype=_dtype))