def apply_erase_rectangles(input: torch.Tensor, params: Dict[str, torch.Tensor]) -> torch.Tensor:
r"""Apply rectangle erase by params.
Generate a {0, 1} mask with drawed rectangle having parameters defined by params and size by input.size().
Args:
input (torch.Tensor): Tensor to be transformed with shape :math:`(*, C, H, W)`.
params (Dict[str, torch.Tensor]):
- params['widths']: widths tensor
- params['heights']: heights tensor
- params['xs']: x positions tensor
- params['ys']: y positions tensor
- params['values']: the value to fill in
Returns:
torch.Tensor: Erased image with shape :math:`(B, C, H, W)`.
"""
if not (params['widths'].size() == params['heights'].size() == params['xs'].size() == params['ys'].size()):
raise TypeError(
"rectangle params components must have same shape. "
f"Got ({params['widths'].size()}, {params['heights'].size()}) "
f"and ({params['xs'].size()}, {params['ys'].size()})"
)
values = params['values'].unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).repeat(1, *input.shape[1:]).to(input)
_, c, h, w = input.size()
bboxes = bbox_generator(params['xs'], params['ys'], params['widths'], params['heights'])
mask = bbox_to_mask(bboxes, w, h) # Returns B, H, W
mask = mask.unsqueeze(1).repeat(1, c, 1, 1).to(input) # Transform to B, c, H, W
transformed = torch.where(mask == 1., values, input)
return transformed
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_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_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) -> 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.
Returns:
params Dict[str, torch.Tensor]: parameters to be passed for transformation.
Example:
>>> _ = torch.manual_seed(0)
>>> crop_size = random_crop_size_generator(
... 3, (30, 30), scale=torch.tensor([.7, 1.3]), ratio=torch.tensor([.9, 1.]))['size']
>>> crop_size
tensor([[26, 29],
[27, 28],
[25, 28]], dtype=torch.int32)
>>> random_crop_generator(3, (30, 30), size=crop_size, same_on_batch=False)
{'src': tensor([[[ 1, 3],
[29, 3],
[29, 28],
[ 1, 28]],
<BLANKLINE>
[[ 2, 3],
[29, 3],
[29, 29],
[ 2, 29]],
<BLANKLINE>
[[ 0, 2],
[27, 2],
[27, 26],
[ 0, 26]]]), 'dst': tensor([[[ 0, 0],
[28, 0],
[28, 25],
[ 0, 25]],
<BLANKLINE>
[[ 0, 0],
[27, 0],
[27, 26],
[ 0, 26]],
<BLANKLINE>
[[ 0, 0],
[27, 0],
[27, 24],
[ 0, 24]]])}
"""
_common_param_check(batch_size, same_on_batch)
if not isinstance(size, torch.Tensor):
size = torch.tensor(size).repeat(batch_size, 1)
assert size.shape == torch.Size([batch_size, 2]), \
f"If `size` is a tensor, it must be shaped as (B, 2). Got {size.shape}."
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 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).long()
y_start = _adapted_uniform((batch_size, ), 0, y_diff[0],
same_on_batch).long()
else:
x_start = _adapted_uniform((1, ), 0, x_diff, same_on_batch).long()
y_start = _adapted_uniform((1, ), 0, y_diff, same_on_batch).long()
crop_src = bbox_generator(x_start.view(-1), y_start.view(-1),
size[:, 1] - 1, size[:, 0] - 1)
if resize_to is None:
crop_dst = bbox_generator(
torch.tensor([0] * batch_size,
device=x_start.device,
dtype=x_start.dtype),
torch.tensor([0] * batch_size,
device=x_start.device,
dtype=x_start.dtype), size[:, 1] - 1, size[:, 0] - 1)
else:
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=x_start.device,
dtype=x_start.dtype).repeat(batch_size, 1, 1)
return dict(src=crop_src, dst=crop_dst)
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
) -> 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 (float or torch.Tensor, optional): hyperparameter for generating cut size from beta distribution.
If None, it will be set to 1.
cut_size ((float, float) or 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.
Returns:
params Dict[str, torch.Tensor]: parameters to be passed for transformation.
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([[[[ 36, 25],
[209, 25],
[209, 198],
[ 36, 198]],
<BLANKLINE>
[[157, 137],
[156, 137],
[156, 136],
[157, 136]],
<BLANKLINE>
[[ 3, 12],
[210, 12],
[210, 219],
[ 3, 219]]],
<BLANKLINE>
<BLANKLINE>
[[[ 83, 126],
[177, 126],
[177, 220],
[ 83, 220]],
<BLANKLINE>
[[ 55, 8],
[206, 8],
[206, 159],
[ 55, 159]],
<BLANKLINE>
[[ 97, 70],
[ 96, 70],
[ 96, 69],
[ 97, 69]]]])}
"""
if beta is None:
beta = torch.tensor(1.)
if cut_size is None:
cut_size = torch.tensor([0., 1.])
assert num_mix >= 1 and isinstance(num_mix, (int,)), \
f"`num_mix` must be an integer greater than 1. Got {num_mix}."
_joint_range_check(cut_size, 'cut_size', bounds=(0, 1))
batch_probs: torch.Tensor = random_prob_generator(batch_size * num_mix, p, same_on_batch)
mix_pairs: torch.Tensor = torch.rand(num_mix, batch_size).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() - 1
cut_width = (cutmix_rate * width).long() - 1
_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, dtype=torch.float32), width - cut_width - 1, same_on_batch).long()
y_start = _adapted_uniform(
_gen_shape, torch.zeros_like(cut_height, dtype=torch.float32), height - cut_height - 1, same_on_batch).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,
crop_src=crop_src
)