mirror of https://github.com/alibaba/EasyCV.git
137 lines
6.1 KiB
Python
137 lines
6.1 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
from numbers import Number
|
|
from typing import Sequence
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
from easycv.datasets.registry import PIPELINES
|
|
|
|
|
|
@PIPELINES.register_module()
|
|
class MMRandomErasing(object):
|
|
"""Randomly selects a rectangle region in an image and erase pixels.
|
|
Args:
|
|
erase_prob (float): Probability that image will be randomly erased.
|
|
Default: 0.5
|
|
min_area_ratio (float): Minimum erased area / input image area
|
|
Default: 0.02
|
|
max_area_ratio (float): Maximum erased area / input image area
|
|
Default: 0.4
|
|
aspect_range (sequence | float): Aspect ratio range of erased area.
|
|
if float, it will be converted to (aspect_ratio, 1/aspect_ratio)
|
|
Default: (3/10, 10/3)
|
|
mode (str): Fill method in erased area, can be:
|
|
- const (default): All pixels are assign with the same value.
|
|
- rand: each pixel is assigned with a random value in [0, 255]
|
|
fill_color (sequence | Number): Base color filled in erased area.
|
|
Defaults to (128, 128, 128).
|
|
fill_std (sequence | Number, optional): If set and ``mode`` is 'rand',
|
|
fill erased area with random color from normal distribution
|
|
(mean=fill_color, std=fill_std); If not set, fill erased area with
|
|
random color from uniform distribution (0~255). Defaults to None.
|
|
Note:
|
|
See `Random Erasing Data Augmentation
|
|
<https://arxiv.org/pdf/1708.04896.pdf>`_
|
|
This paper provided 4 modes: RE-R, RE-M, RE-0, RE-255, and use RE-M as
|
|
default. The config of these 4 modes are:
|
|
- RE-R: RandomErasing(mode='rand')
|
|
- RE-M: RandomErasing(mode='const', fill_color=(123.67, 116.3, 103.5))
|
|
- RE-0: RandomErasing(mode='const', fill_color=0)
|
|
- RE-255: RandomErasing(mode='const', fill_color=255)
|
|
"""
|
|
|
|
def __init__(self,
|
|
erase_prob=0.5,
|
|
min_area_ratio=0.02,
|
|
max_area_ratio=0.4,
|
|
aspect_range=(3 / 10, 10 / 3),
|
|
mode='const',
|
|
fill_color=(128, 128, 128),
|
|
fill_std=None):
|
|
assert isinstance(erase_prob, float) and 0. <= erase_prob <= 1.
|
|
assert isinstance(min_area_ratio, float) and 0. <= min_area_ratio <= 1.
|
|
assert isinstance(max_area_ratio, float) and 0. <= max_area_ratio <= 1.
|
|
assert min_area_ratio <= max_area_ratio, \
|
|
'min_area_ratio should be smaller than max_area_ratio'
|
|
if isinstance(aspect_range, float):
|
|
aspect_range = min(aspect_range, 1 / aspect_range)
|
|
aspect_range = (aspect_range, 1 / aspect_range)
|
|
assert isinstance(aspect_range, Sequence) and len(aspect_range) == 2 \
|
|
and all(isinstance(x, float) for x in aspect_range), \
|
|
'aspect_range should be a float or Sequence with two float.'
|
|
assert all(x > 0 for x in aspect_range), \
|
|
'aspect_range should be positive.'
|
|
assert aspect_range[0] <= aspect_range[1], \
|
|
'In aspect_range (min, max), min should be smaller than max.'
|
|
assert mode in ['const', 'rand']
|
|
if isinstance(fill_color, Number):
|
|
fill_color = [fill_color] * 3
|
|
assert isinstance(fill_color, Sequence) and len(fill_color) == 3 \
|
|
and all(isinstance(x, Number) for x in fill_color), \
|
|
'fill_color should be a float or Sequence with three int.'
|
|
if fill_std is not None:
|
|
if isinstance(fill_std, Number):
|
|
fill_std = [fill_std] * 3
|
|
assert isinstance(fill_std, Sequence) and len(fill_std) == 3 \
|
|
and all(isinstance(x, Number) for x in fill_std), \
|
|
'fill_std should be a float or Sequence with three int.'
|
|
|
|
self.erase_prob = erase_prob
|
|
self.min_area_ratio = min_area_ratio
|
|
self.max_area_ratio = max_area_ratio
|
|
self.aspect_range = aspect_range
|
|
self.mode = mode
|
|
self.fill_color = fill_color
|
|
self.fill_std = fill_std
|
|
|
|
def _fill_pixels(self, img, top, left, h, w):
|
|
if self.mode == 'const':
|
|
patch = np.empty((h, w, 3), dtype=np.uint8)
|
|
patch[:, :] = np.array(self.fill_color, dtype=np.uint8)
|
|
elif self.fill_std is None:
|
|
# Uniform distribution
|
|
patch = np.random.uniform(0, 256, (h, w, 3)).astype(np.uint8)
|
|
else:
|
|
# Normal distribution
|
|
patch = np.random.normal(self.fill_color, self.fill_std, (h, w, 3))
|
|
patch = np.clip(patch.astype(np.int32), 0, 255).astype(np.uint8)
|
|
|
|
img[top:top + h, left:left + w] = patch
|
|
return img
|
|
|
|
def __call__(self, results):
|
|
if np.random.rand() > self.erase_prob:
|
|
return results
|
|
|
|
for key in results.get('img_fields', ['img']):
|
|
img = np.array(results[key])
|
|
img_h, img_w = img.shape[:2]
|
|
|
|
# convert to log aspect to ensure equal probability of aspect ratio
|
|
log_aspect_range = np.log(
|
|
np.array(self.aspect_range, dtype=np.float32))
|
|
aspect_ratio = np.exp(np.random.uniform(*log_aspect_range))
|
|
area = img_h * img_w
|
|
area *= np.random.uniform(self.min_area_ratio, self.max_area_ratio)
|
|
|
|
h = min(int(round(np.sqrt(area * aspect_ratio))), img_h)
|
|
w = min(int(round(np.sqrt(area / aspect_ratio))), img_w)
|
|
top = np.random.randint(0, img_h - h) if img_h > h else 0
|
|
left = np.random.randint(0, img_w - w) if img_w > w else 0
|
|
img = self._fill_pixels(img, top, left, h, w)
|
|
results[key] = Image.fromarray(img.astype(np.uint8))
|
|
|
|
return results
|
|
|
|
def __repr__(self):
|
|
repr_str = self.__class__.__name__
|
|
repr_str += f'(erase_prob={self.erase_prob}, '
|
|
repr_str += f'min_area_ratio={self.min_area_ratio}, '
|
|
repr_str += f'max_area_ratio={self.max_area_ratio}, '
|
|
repr_str += f'aspect_range={self.aspect_range}, '
|
|
repr_str += f'mode={self.mode}, '
|
|
repr_str += f'fill_color={self.fill_color}, '
|
|
repr_str += f'fill_std={self.fill_std})'
|
|
return repr_str
|