|
- from functools import partial
- from typing import Optional, List, Union, Tuple
- from . import functional as F
- from .base import DualTransform, ImageOnlyTransform
-
-
- class HorizontalFlip(DualTransform):
- """Flip images horizontally (left->right)"""
-
- identity_param = False
-
- def __init__(self):
- super().__init__("apply", [False, True])
-
- def apply_aug_image(self, image, apply=False, **kwargs):
- if apply:
- image = F.hflip(image)
- return image
-
- def apply_deaug_mask(self, mask, apply=False, **kwargs):
- if apply:
- mask = F.hflip(mask)
- return mask
-
- def apply_deaug_label(self, label, apply=False, **kwargs):
- return label
-
- def apply_deaug_keypoints(self, keypoints, apply=False, **kwargs):
- if apply:
- keypoints = F.keypoints_hflip(keypoints)
- return keypoints
-
-
- class VerticalFlip(DualTransform):
- """Flip images vertically (up->down)"""
-
- identity_param = False
-
- def __init__(self):
- super().__init__("apply", [False, True])
-
- def apply_aug_image(self, image, apply=False, **kwargs):
- if apply:
- image = F.vflip(image)
- return image
-
- def apply_deaug_mask(self, mask, apply=False, **kwargs):
- if apply:
- mask = F.vflip(mask)
- return mask
-
- def apply_deaug_label(self, label, apply=False, **kwargs):
- return label
-
- def apply_deaug_keypoints(self, keypoints, apply=False, **kwargs):
- if apply:
- keypoints = F.keypoints_vflip(keypoints)
- return keypoints
-
-
- class HorizontalShift(DualTransform):
- """Shift images horizontally (left->right)"""
-
- identity_param = 0
-
- def __init__(self, shifts: List[float]):
- if self.identity_param not in shifts:
- shifts = [self.identity_param] + list(shifts)
- super().__init__("shifts", shifts)
-
- def apply_aug_image(self, image, shifts=0, **kwargs):
- image = F.hshift(image, shifts)
- return image
-
- def apply_deaug_mask(self, mask, shifts=0, **kwargs):
- return self.apply_aug_image(mask, -shifts)
-
- def apply_deaug_label(self, label, shifts=0, **kwargs):
- return label
-
- def apply_deaug_keypoints(self, keypoints, shifts=0, **kwargs):
- return F.keypoints_hshift(keypoints, -shifts)
-
-
- class VerticalShift(DualTransform):
- """Shift images vertically (up->down)"""
-
- identity_param = 0
-
- def __init__(self, shifts: List[float]):
- if self.identity_param not in shifts:
- shifts = [self.identity_param] + list(shifts)
- super().__init__("shifts", shifts)
-
- def apply_aug_image(self, image, shifts=0, **kwargs):
- image = F.vshift(image, shifts)
- return image
-
- def apply_deaug_mask(self, mask, shifts=0, **kwargs):
- return self.apply_aug_image(mask, -shifts)
-
- def apply_deaug_label(self, label, shifts=0, **kwargs):
- return label
-
- def apply_deaug_keypoints(self, keypoints, shifts=0, **kwargs):
- return F.keypoints_vshift(keypoints, -shifts)
-
-
- class Rotate90(DualTransform):
- """Rotate images 0/90/180/270 degrees
-
- Args:
- angles (list): angles to rotate images
- """
-
- identity_param = 0
-
- def __init__(self, angles: List[int]):
- if self.identity_param not in angles:
- angles = [self.identity_param] + list(angles)
-
- super().__init__("angle", angles)
-
- def apply_aug_image(self, image, angle=0, **kwargs):
- k = angle // 90 if angle >= 0 else (angle + 360) // 90
- return F.rot90(image, k)
-
- def apply_deaug_mask(self, mask, angle=0, **kwargs):
- return self.apply_aug_image(mask, -angle)
-
- def apply_deaug_label(self, label, angle=0, **kwargs):
- return label
-
- def apply_deaug_keypoints(self, keypoints, angle=0, **kwargs):
- angle *= -1
- k = angle // 90 if angle >= 0 else (angle + 360) // 90
- return F.keypoints_rot90(keypoints, k=k)
-
-
- class Scale(DualTransform):
- """Scale images
-
- Args:
- scales (List[Union[int, float]]): scale factors for spatial image dimensions
- interpolation (str): one of "nearest"/"lenear" (see more in paddle.nn.interpolate)
- align_corners (bool): see more in paddle.nn.interpolate
- """
-
- identity_param = 1
-
- def __init__(
- self,
- scales: List[Union[int, float]],
- interpolation: str = "nearest",
- align_corners: bool = False,
- ):
- if self.identity_param not in scales:
- scales = [self.identity_param] + list(scales)
- self.interpolation = interpolation
- self.align_corners = align_corners
-
- super().__init__("scale", scales)
-
- def apply_aug_image(self, image, scale=1, **kwargs):
- if scale != self.identity_param:
- image = F.scale(
- image,
- scale,
- interpolation=self.interpolation,
- align_corners=self.align_corners,
- )
- return image
-
- def apply_deaug_mask(self, mask, scale=1, **kwargs):
- if scale != self.identity_param:
- mask = F.scale(
- mask,
- 1 / scale,
- interpolation=self.interpolation,
- align_corners=self.align_corners,
- )
- return mask
-
- def apply_deaug_label(self, label, scale=1, **kwargs):
- return label
-
- def apply_deaug_keypoints(self, keypoints, scale=1, **kwargs):
- return keypoints
-
-
- class Resize(DualTransform):
- """Resize images
-
- Args:
- sizes (List[Tuple[int, int]): scale factors for spatial image dimensions
- original_size Tuple(int, int): optional, image original size for deaugmenting mask
- interpolation (str): one of "nearest"/"lenear" (see more in paddle.nn.interpolate)
- align_corners (bool): see more in paddle.nn.interpolate
- """
-
- def __init__(
- self,
- sizes: List[Tuple[int, int]],
- original_size: Tuple[int, int] = None,
- interpolation: str = "nearest",
- align_corners: bool = False,
- ):
- if original_size is not None and original_size not in sizes:
- sizes = [original_size] + list(sizes)
- self.interpolation = interpolation
- self.align_corners = align_corners
- self.original_size = original_size
-
- super().__init__("size", sizes)
-
- def apply_aug_image(self, image, size, **kwargs):
- if size != self.original_size:
- image = F.resize(
- image,
- size,
- interpolation=self.interpolation,
- align_corners=self.align_corners,
- )
- return image
-
- def apply_deaug_mask(self, mask, size, **kwargs):
- if self.original_size is None:
- raise ValueError(
- "Provide original image size to make mask backward transformation"
- )
- if size != self.original_size:
- mask = F.resize(
- mask,
- self.original_size,
- interpolation=self.interpolation,
- align_corners=self.align_corners,
- )
- return mask
-
- def apply_deaug_label(self, label, size=1, **kwargs):
- return label
-
- def apply_deaug_keypoints(self, keypoints, size=1, **kwargs):
- return keypoints
-
-
- class Add(ImageOnlyTransform):
- """Add value to images
-
- Args:
- values (List[float]): values to add to each pixel
- """
-
- identity_param = 0
-
- def __init__(self, values: List[float]):
-
- if self.identity_param not in values:
- values = [self.identity_param] + list(values)
- super().__init__("value", values)
-
- def apply_aug_image(self, image, value=0, **kwargs):
- if value != self.identity_param:
- image = F.add(image, value)
- return image
-
-
- class Multiply(ImageOnlyTransform):
- """Multiply images by factor
-
- Args:
- factors (List[float]): factor to multiply each pixel by
- """
-
- identity_param = 1
-
- def __init__(self, factors: List[float]):
- if self.identity_param not in factors:
- factors = [self.identity_param] + list(factors)
- super().__init__("factor", factors)
-
- def apply_aug_image(self, image, factor=1, **kwargs):
- if factor != self.identity_param:
- image = F.multiply(image, factor)
- return image
-
-
- class FiveCrops(ImageOnlyTransform):
- """Makes 4 crops for each corner + center crop
-
- Args:
- crop_height (int): crop height in pixels
- crop_width (int): crop width in pixels
- """
-
- def __init__(self, crop_height, crop_width):
- crop_functions = (
- partial(F.crop_lt, crop_h=crop_height, crop_w=crop_width),
- partial(F.crop_lb, crop_h=crop_height, crop_w=crop_width),
- partial(F.crop_rb, crop_h=crop_height, crop_w=crop_width),
- partial(F.crop_rt, crop_h=crop_height, crop_w=crop_width),
- partial(F.center_crop, crop_h=crop_height, crop_w=crop_width),
- )
- super().__init__("crop_fn", crop_functions)
-
- def apply_aug_image(self, image, crop_fn=None, **kwargs):
- return crop_fn(image)
-
- def apply_deaug_mask(self, mask, **kwargs):
- raise ValueError("`FiveCrop` augmentation is not suitable for mask!")
-
- def apply_deaug_keypoints(self, keypoints, **kwargs):
- raise ValueError("`FiveCrop` augmentation is not suitable for keypoints!")
-
-
- class AdjustContrast(ImageOnlyTransform):
- """Adjusts contrast of an Image.
-
- Args:
- factors (List[float]): How much to adjust the contrast.
- """
-
- identity_param = 1.
-
- def __init__(self, factors: List[float]):
-
- if self.identity_param not in factors:
- factors = [self.identity_param] + list(factors)
- super().__init__("factor", factors)
-
- def apply_aug_image(self, image, factor=1., **kwargs):
- if factor != self.identity_param:
- return F.adjust_contrast(image, factor)
- return image
-
-
- class AdjustBrightness(ImageOnlyTransform):
- """Adjusts brightness of an Image.
-
- Args:
- factors (List[float]): How much to adjust the brightness.
- """
-
- identity_param = 1.
-
- def __init__(self, factors: List[float]):
-
- if self.identity_param not in factors:
- factors = [self.identity_param] + list(factors)
- super().__init__("factor", factors)
-
- def apply_aug_image(self, image, factor=1., **kwargs):
- if factor != self.identity_param:
- return F.adjust_brightness(image, factor)
- return image
-
-
- class AverageBlur(ImageOnlyTransform):
- """Apply average blur to the input image
-
- Args:
- kernel_sizes (List[Union[Tuple[int, int], int]]): kernel size of average blur
- """
-
- identity_param = 1
-
- def __init__(
- self,
- kernel_sizes: List[Union[Tuple[int, int], int]],
- ):
- if self.identity_param not in kernel_sizes:
- kernel_sizes = [self.identity_param] + list(kernel_sizes)
- super().__init__("kernel_size", kernel_sizes)
-
- def apply_aug_image(self, image, kernel_size, **kwargs):
- if kernel_size == self.identity_param:
- return image
- return F.average_blur(image, kernel_size)
-
-
- class GaussianBlur(ImageOnlyTransform):
- """Apply gaussian blur to the input image
-
- Args:
- kernel_sizes (List[Union[Tuple[int, int], int]]): kernel size of gaussian blur
- sigma (Optional[Union[Tuple[float, float], float]]): gaussian kernel standard deviation
- """
-
- identity_param = 1
-
- def __init__(
- self,
- kernel_sizes: List[Union[Tuple[int, int], int]],
- sigma: Optional[Union[Tuple[float, float], float]] = None
- ):
- if self.identity_param not in kernel_sizes:
- kernel_sizes = [self.identity_param] + list(kernel_sizes)
- self.sigma = sigma
- super().__init__("kernel_size", kernel_sizes)
-
- def apply_aug_image(self, image, kernel_size, **kwargs):
- if kernel_size == self.identity_param:
- return image
- return F.gaussian_blur(image, kernel_size, self.sigma)
-
-
- class Sharpen(ImageOnlyTransform):
- """Apply sharpen to the input image
-
- Args:
- kernel_sizes (List[int]): kernel size of sharpen
- """
-
- identity_param = 1
-
- def __init__(self, kernel_sizes: List[int]):
- if self.identity_param not in kernel_sizes:
- kernel_sizes = [self.identity_param] + list(kernel_sizes)
- super().__init__("kernel_size", kernel_sizes)
-
- def apply_aug_image(self, image, kernel_size, **kwargs):
- if kernel_size == self.identity_param:
- return image
- return F.sharpen(image, kernel_size)
|