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