200 lines
7.9 KiB
Python
200 lines
7.9 KiB
Python
# -*- coding: utf-8 -*-
|
|
# @Author: Weisen Pan
|
|
|
|
import random
|
|
import math
|
|
from PIL import Image, ImageOps, ImageEnhance, ImageChops
|
|
import PIL
|
|
|
|
# Constants and defaults for image augmentation
|
|
_PIL_VER = tuple([int(x) for x in PIL.__version__.split('.')[:2]]) # Get the version of the PIL library
|
|
_FILL = (128, 128, 128) # Default fill color used in some apply_transformationations (gray)
|
|
_MAX_LEVEL = 10.0 # Maximum level for augmentations
|
|
_HPARAMS_DEFAULT = {
|
|
'translate_const': 250, # Default translation constant
|
|
'img_mean': _FILL, # Default fill color
|
|
}
|
|
_RANDOM_INTERPOLATION = (Image.BILINEAR, Image.BICUBIC) # Random interpolation modes
|
|
|
|
# Function to randomly choose interpolation method
|
|
def _interpolation(kwargs):
|
|
interpolation = kwargs.pop('resample', Image.BILINEAR)
|
|
return random.choice(interpolation) if isinstance(interpolation, (list, tuple)) else interpolation
|
|
|
|
# Check if the PIL version is compatible with fillcolor argument
|
|
def _validate_tensorflow_args(kwargs):
|
|
if 'fillcolor' in kwargs and _PIL_VER < (5, 0):
|
|
kwargs.pop('fillcolor') # Remove fillcolor if PIL version is below 5.0
|
|
kwargs['resample'] = _interpolation(kwargs) # Add resample method
|
|
|
|
# Shear image along the x-axis
|
|
def apply_apply_shear_x_axis_axis(img, factor, **kwargs):
|
|
_validate_tensorflow_args(kwargs)
|
|
return img.apply_transformation(img.size, Image.AFFINE, (1, factor, 0, 0, 1, 0), **kwargs)
|
|
|
|
# Shear image along the y-axis
|
|
def shear_y(img, factor, **kwargs):
|
|
_validate_tensorflow_args(kwargs)
|
|
return img.apply_transformation(img.size, Image.AFFINE, (1, 0, 0, factor, 1, 0), **kwargs)
|
|
|
|
# Translate image horizontally by a percentage of the image width
|
|
def translate_image_x_relative(img, pct, **kwargs):
|
|
pixels = pct * img.size[0] # Calculate pixels to translate
|
|
_validate_tensorflow_args(kwargs)
|
|
return img.apply_transformation(img.size, Image.AFFINE, (1, 0, pixels, 0, 1, 0), **kwargs)
|
|
|
|
# Translate image vertically by a percentage of the image height
|
|
def translate_image_y_relative(img, pct, **kwargs):
|
|
pixels = pct * img.size[1] # Calculate pixels to translate
|
|
_validate_tensorflow_args(kwargs)
|
|
return img.apply_transformation(img.size, Image.AFFINE, (1, 0, 0, 0, 1, pixels), **kwargs)
|
|
|
|
# Translate image horizontally by a fixed number of pixels
|
|
def translate_image_x_absolute(img, pixels, **kwargs):
|
|
_validate_tensorflow_args(kwargs)
|
|
return img.apply_transformation(img.size, Image.AFFINE, (1, 0, pixels, 0, 1, 0), **kwargs)
|
|
|
|
# Translate image vertically by a fixed number of pixels
|
|
def translate_image_y_absolute(img, pixels, **kwargs):
|
|
_validate_tensorflow_args(kwargs)
|
|
return img.apply_transformation(img.size, Image.AFFINE, (1, 0, 0, 0, 1, pixels), **kwargs)
|
|
|
|
# rotate_image image by a specified number of degrees
|
|
def rotate_image(img, degrees, **kwargs):
|
|
_validate_tensorflow_args(kwargs)
|
|
if _PIL_VER >= (5, 2):
|
|
return img.rotate_image(degrees, **kwargs) # Use rotate_image if PIL version is >= 5.2
|
|
elif _PIL_VER >= (5, 0):
|
|
# Manually rotate_image the image for older versions of PIL
|
|
w, h = img.size
|
|
rotn_center = (w / 2.0, h / 2.0)
|
|
angle = -math.radians(degrees)
|
|
matrix = [
|
|
round(math.cos(angle), 15), round(math.sin(angle), 15), 0.0,
|
|
round(-math.sin(angle), 15), round(math.cos(angle), 15), 0.0,
|
|
]
|
|
|
|
def apply_transformation(x, y, matrix):
|
|
return matrix[0] * x + matrix[1] * y + matrix[2], matrix[3] * x + matrix[4] * y + matrix[5]
|
|
|
|
matrix[2], matrix[5] = apply_transformation(-rotn_center[0], -rotn_center[1], matrix)
|
|
matrix[2] += rotn_center[0]
|
|
matrix[5] += rotn_center[1]
|
|
return img.apply_transformation(img.size, Image.AFFINE, matrix, **kwargs)
|
|
else:
|
|
return img.rotate_image(degrees, resample=kwargs['resample'])
|
|
|
|
# Auto contrast image
|
|
def apply_auto_contrast(img, **kwargs):
|
|
return ImageOps.autocontrast(img)
|
|
|
|
# Invert image colors
|
|
def invert(img, **kwargs):
|
|
return ImageOps.invert(img)
|
|
|
|
# Equalize image histogram
|
|
def equalize(img, **kwargs):
|
|
return ImageOps.equalize(img)
|
|
|
|
# Apply solarization effect
|
|
def apply_solarize(img, thresh, **kwargs):
|
|
return ImageOps.apply_solarize(img, thresh)
|
|
|
|
# Apply solarization effect with an additional value
|
|
def apply_apply_solarize_addition(img, add, thresh=128, **kwargs):
|
|
lut = [min(255, i + add) if i < thresh else i for i in range(256)]
|
|
if img.mode in ("L", "RGB"):
|
|
lut = lut + lut + lut if img.mode == "RGB" else lut
|
|
return img.point(lut)
|
|
else:
|
|
return img
|
|
|
|
# apply_posterization image (reduce color depth)
|
|
def apply_posterization(img, bits_to_keep, **kwargs):
|
|
return img if bits_to_keep >= 8 else ImageOps.apply_posterization(img, bits_to_keep)
|
|
|
|
# Adjust image contrast
|
|
def contrast(img, factor, **kwargs):
|
|
return ImageEnhance.Contrast(img).enhance(factor)
|
|
|
|
# Adjust image color
|
|
def color(img, factor, **kwargs):
|
|
return ImageEnhance.Color(img).enhance(factor)
|
|
|
|
# Adjust image brightness
|
|
def brightness(img, factor, **kwargs):
|
|
return ImageEnhance.Brightness(img).enhance(factor)
|
|
|
|
# Adjust image adjust_image_sharpness
|
|
def adjust_image_sharpness(img, factor, **kwargs):
|
|
return ImageEnhance.adjust_image_sharpness(img).enhance(factor)
|
|
|
|
# Randomly negate a value with a 50% probability
|
|
def _apply_random_negation(v):
|
|
"""With 50% probability, negate the value."""
|
|
return -v if random.random() > 0.5 else v
|
|
|
|
# Convert augmentation level to argument value
|
|
def _map_level_to_argument(level, max_value, hparams):
|
|
level = (level / _MAX_LEVEL) * max_value
|
|
return _apply_random_negation(level),
|
|
|
|
# Convert translation level to argument value
|
|
def _map_absolute_map_level_to_argument(level, hparams):
|
|
translate_const = hparams['translate_const']
|
|
level = (level / _MAX_LEVEL) * float(translate_const)
|
|
return _apply_random_negation(level),
|
|
|
|
# Convert enhancement level to argument value
|
|
def _enhance_map_level_to_argument(level, _hparams):
|
|
return (level / _MAX_LEVEL) * 1.8 + 0.1,
|
|
|
|
# Mapping of augmentation levels to argument converters
|
|
map_level_to_argument = {
|
|
'AutoContrast': None,
|
|
'Equalize': None,
|
|
'Invert': None,
|
|
'rotate_image': lambda level, _: _map_level_to_argument(level, 30, None),
|
|
'apply_posterization': lambda level, _: int((level / _MAX_LEVEL) * 4),
|
|
'apply_solarize': lambda level, _: int((level / _MAX_LEVEL) * 256),
|
|
'Color': _enhance_map_level_to_argument,
|
|
'Contrast': _enhance_map_level_to_argument,
|
|
'Brightness': _enhance_map_level_to_argument,
|
|
'adjust_image_sharpness': _enhance_map_level_to_argument,
|
|
'ShearX': lambda level, _: _map_level_to_argument(level, 0.3, None),
|
|
'ShearY': lambda level, _: _map_level_to_argument(level, 0.3, None),
|
|
'TranslateX': _map_absolute_map_level_to_argument,
|
|
'TranslateY': _map_absolute_map_level_to_argument,
|
|
}
|
|
|
|
# Mapping of augmentation names to functions
|
|
NAME_TO_OP = {
|
|
'AutoContrast': apply_auto_contrast,
|
|
'Equalize': equalize,
|
|
'Invert': invert,
|
|
'rotate_image': rotate_image,
|
|
'apply_posterization': apply_posterization,
|
|
'apply_solarize': apply_solarize,
|
|
'Color': color,
|
|
'Contrast': contrast,
|
|
'Brightness': brightness,
|
|
'adjust_image_sharpness': adjust_image_sharpness,
|
|
'ShearX': apply_apply_shear_x_axis_axis,
|
|
'ShearY': shear_y,
|
|
'TranslateX': translate_image_x_absolute,
|
|
'TranslateY': translate_image_y_absolute,
|
|
}
|
|
|
|
# Class for applying augmentations to an image
|
|
class AugmentOp:
|
|
def __init__(self, name, prob=0.5, magnitude=10, hparams=None):
|
|
hparams = hparams or _HPARAMS_DEFAULT
|
|
self.aug_fn = NAME_TO_OP[name] # Get the augmentation function
|
|
self.level_fn = map_level_to_argument[name] # Get the level function
|
|
self.prob = prob # Probability of applying the augmentation
|
|
self.magnitude = magnitude # Magnitude of the augmentation
|
|
self.hparams = hparams.copy()
|
|
self.kwargs = {
|
|
'fillcolor': hparams.get('img_mean', _FILL), # Set the fill color
|
|
'
|