mirror of
https://github.com/open-mmlab/mmpretrain.git
synced 2025-06-03 14:59:18 +08:00
8c11c01fdb
* half-done auto_augmentation * remove auto_augcode and support shear pipeline * fix typo * fix typo * use a non-square toy data instead
87 lines
3.6 KiB
Python
87 lines
3.6 KiB
Python
import mmcv
|
|
import numpy as np
|
|
|
|
from ..builder import PIPELINES
|
|
|
|
|
|
def random_negative(value, random_negative_prob):
|
|
"""Randomly negate value based on random_negative_prob."""
|
|
return -value if np.random.rand() < random_negative_prob else value
|
|
|
|
|
|
@PIPELINES.register_module()
|
|
class Shear(object):
|
|
"""Shear images.
|
|
|
|
Args:
|
|
magnitude (int | float): The magnitude used for shear.
|
|
pad_val (int, tuple[int]): Pixel pad_val value for constant fill. If a
|
|
tuple of length 3, it is used to pad_val R, G, B channels
|
|
respectively. Defaults to 128.
|
|
prob (float): The probability for performing Shear therefore should be
|
|
in range [0, 1]. Defaults to 0.5.
|
|
direction (str): The shearing direction. Options are 'horizontal' and
|
|
'vertical'. Defaults to 'horizontal'.
|
|
random_negative_prob (float): The probability that turns the magnitude
|
|
negative, which should be in range [0,1]. Defaults to 0.5.
|
|
interpolation (str): Interpolation method. Options are 'nearest',
|
|
'bilinear', 'bicubic', 'area', 'lanczos'. Defaults to 'bicubic'.
|
|
"""
|
|
|
|
def __init__(self,
|
|
magnitude,
|
|
pad_val=128,
|
|
prob=0.5,
|
|
direction='horizontal',
|
|
random_negative_prob=0.5,
|
|
interpolation='bicubic'):
|
|
assert isinstance(magnitude, (int, float)), 'The magnitude type must '\
|
|
f'be int or float, but got {type(magnitude)} instead.'
|
|
if isinstance(pad_val, int):
|
|
pad_val = tuple([pad_val] * 3)
|
|
elif isinstance(pad_val, tuple):
|
|
assert len(pad_val) == 3, 'pad_val as a tuple must have 3 ' \
|
|
f'elements, got {len(pad_val)} instead.'
|
|
assert all(isinstance(i, int) for i in pad_val), 'pad_val as a '\
|
|
'tuple must got elements of int type.'
|
|
else:
|
|
raise TypeError('pad_val must be int or tuple with 3 elements.')
|
|
assert 0 <= prob <= 1.0, 'The prob should be in range [0,1], ' \
|
|
f'got {prob} instead.'
|
|
assert direction in ('horizontal', 'vertical'), 'direction must be ' \
|
|
f'either "horizontal" or "vertical", got {direction} instead.'
|
|
assert 0 <= random_negative_prob <= 1.0, 'The random_negative_prob ' \
|
|
f'should be in range [0,1], got {random_negative_prob} instead.'
|
|
|
|
self.magnitude = magnitude
|
|
self.pad_val = pad_val
|
|
self.prob = prob
|
|
self.direction = direction
|
|
self.random_negative_prob = random_negative_prob
|
|
self.interpolation = interpolation
|
|
|
|
def __call__(self, results):
|
|
if np.random.rand() > self.prob:
|
|
return results
|
|
magnitude = random_negative(self.magnitude, self.random_negative_prob)
|
|
for key in results.get('img_fields', ['img']):
|
|
img = results[key]
|
|
img_sheared = mmcv.imshear(
|
|
img,
|
|
magnitude,
|
|
direction=self.direction,
|
|
border_value=self.pad_val,
|
|
interpolation=self.interpolation)
|
|
results[key] = img_sheared.astype(img.dtype)
|
|
return results
|
|
|
|
def __repr__(self):
|
|
repr_str = self.__class__.__name__
|
|
repr_str += f'(magnitude={self.magnitude}, '
|
|
repr_str += f'pad_val={self.pad_val}, '
|
|
repr_str += f'prob={self.prob}, '
|
|
repr_str += f'direction={self.direction}, '
|
|
repr_str += f'random_negative_prob={self.random_negative_prob}, '
|
|
repr_str += f'interpolation={self.interpolation})'
|
|
return repr_str
|