|
""" |
|
This implementation is based on |
|
https://github.com/rwightman/pytorch-image-models/blob/master/timm/data/random_erasing.py |
|
pulished under an Apache License 2.0. |
|
""" |
|
import math |
|
import random |
|
import torch |
|
|
|
|
|
def _get_pixels( |
|
per_pixel, rand_color, patch_size, dtype=torch.float32, device="cuda" |
|
): |
|
|
|
|
|
|
|
if per_pixel: |
|
return torch.empty(patch_size, dtype=dtype, device=device).normal_() |
|
elif rand_color: |
|
return torch.empty( |
|
(patch_size[0], 1, 1), dtype=dtype, device=device |
|
).normal_() |
|
else: |
|
return torch.zeros((patch_size[0], 1, 1), dtype=dtype, device=device) |
|
|
|
|
|
class RandomErasing: |
|
"""Randomly selects a rectangle region in an image and erases its pixels. |
|
'Random Erasing Data Augmentation' by Zhong et al. |
|
See https://arxiv.org/pdf/1708.04896.pdf |
|
This variant of RandomErasing is intended to be applied to either a batch |
|
or single image tensor after it has been normalized by dataset mean and std. |
|
Args: |
|
probability: Probability that the Random Erasing operation will be performed. |
|
min_area: Minimum percentage of erased area wrt input image area. |
|
max_area: Maximum percentage of erased area wrt input image area. |
|
min_aspect: Minimum aspect ratio of erased area. |
|
mode: pixel color mode, one of 'const', 'rand', or 'pixel' |
|
'const' - erase block is constant color of 0 for all channels |
|
'rand' - erase block is same per-channel random (normal) color |
|
'pixel' - erase block is per-pixel random (normal) color |
|
max_count: maximum number of erasing blocks per image, area per box is scaled by count. |
|
per-image count is randomly chosen between 1 and this value. |
|
""" |
|
|
|
def __init__( |
|
self, |
|
probability=0.5, |
|
min_area=0.02, |
|
max_area=1 / 3, |
|
min_aspect=0.3, |
|
max_aspect=None, |
|
mode="const", |
|
min_count=1, |
|
max_count=None, |
|
num_splits=0, |
|
device="cuda", |
|
cube=True, |
|
): |
|
self.probability = probability |
|
self.min_area = min_area |
|
self.max_area = max_area |
|
max_aspect = max_aspect or 1 / min_aspect |
|
self.log_aspect_ratio = (math.log(min_aspect), math.log(max_aspect)) |
|
self.min_count = min_count |
|
self.max_count = max_count or min_count |
|
self.num_splits = num_splits |
|
mode = mode.lower() |
|
self.rand_color = False |
|
self.per_pixel = False |
|
self.cube = cube |
|
if mode == "rand": |
|
self.rand_color = True |
|
elif mode == "pixel": |
|
self.per_pixel = True |
|
else: |
|
assert not mode or mode == "const" |
|
self.device = device |
|
|
|
def _erase(self, img, chan, img_h, img_w, dtype): |
|
if random.random() > self.probability: |
|
return |
|
area = img_h * img_w |
|
count = ( |
|
self.min_count |
|
if self.min_count == self.max_count |
|
else random.randint(self.min_count, self.max_count) |
|
) |
|
for _ in range(count): |
|
for _ in range(10): |
|
target_area = ( |
|
random.uniform(self.min_area, self.max_area) * area / count |
|
) |
|
aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio)) |
|
h = int(round(math.sqrt(target_area * aspect_ratio))) |
|
w = int(round(math.sqrt(target_area / aspect_ratio))) |
|
if w < img_w and h < img_h: |
|
top = random.randint(0, img_h - h) |
|
left = random.randint(0, img_w - w) |
|
img[:, top : top + h, left : left + w] = _get_pixels( |
|
self.per_pixel, |
|
self.rand_color, |
|
(chan, h, w), |
|
dtype=dtype, |
|
device=self.device, |
|
) |
|
break |
|
|
|
def _erase_cube( |
|
self, |
|
img, |
|
batch_start, |
|
batch_size, |
|
chan, |
|
img_h, |
|
img_w, |
|
dtype, |
|
): |
|
if random.random() > self.probability: |
|
return |
|
area = img_h * img_w |
|
count = ( |
|
self.min_count |
|
if self.min_count == self.max_count |
|
else random.randint(self.min_count, self.max_count) |
|
) |
|
for _ in range(count): |
|
for _ in range(100): |
|
target_area = ( |
|
random.uniform(self.min_area, self.max_area) * area / count |
|
) |
|
aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio)) |
|
h = int(round(math.sqrt(target_area * aspect_ratio))) |
|
w = int(round(math.sqrt(target_area / aspect_ratio))) |
|
if w < img_w and h < img_h: |
|
top = random.randint(0, img_h - h) |
|
left = random.randint(0, img_w - w) |
|
for i in range(batch_start, batch_size): |
|
img_instance = img[i] |
|
img_instance[ |
|
:, top : top + h, left : left + w |
|
] = _get_pixels( |
|
self.per_pixel, |
|
self.rand_color, |
|
(chan, h, w), |
|
dtype=dtype, |
|
device=self.device, |
|
) |
|
break |
|
|
|
def __call__(self, input): |
|
if len(input.size()) == 3: |
|
self._erase(input, *input.size(), input.dtype) |
|
else: |
|
batch_size, chan, img_h, img_w = input.size() |
|
|
|
batch_start = ( |
|
batch_size // self.num_splits if self.num_splits > 1 else 0 |
|
) |
|
if self.cube: |
|
self._erase_cube( |
|
input, |
|
batch_start, |
|
batch_size, |
|
chan, |
|
img_h, |
|
img_w, |
|
input.dtype, |
|
) |
|
else: |
|
for i in range(batch_start, batch_size): |
|
self._erase(input[i], chan, img_h, img_w, input.dtype) |
|
return input |
|
|