[PaddlePaddle Hackathon] add image augment algorithms

README.md

@@ -125,12 +125,19 @@ mask = mean(masks)
 |----------------|:-------------------------:|:---------------------------------:|
 | HorizontalFlip | -                         | -                                 |
 | VerticalFlip   | -                         | -                                 |
+| HorizontalShift| shifts                    | List\[float]                      |
+| VerticalShift  | shifts                    | List\[float]                      |
 | Rotate90       | angles                    | List\[0, 90, 180, 270]            |
 | Scale          | scales<br>interpolation   | List\[float]<br>"nearest"/"linear"|
 | Resize         | sizes<br>original_size<br>interpolation   | List\[Tuple\[int, int]]<br>Tuple\[int,int]<br>"nearest"/"linear"|
 | Add            | values                    | List\[float]                      |
 | Multiply       | factors                   | List\[float]                      |
 | FiveCrops      | crop_height<br>crop_width | int<br>int                        |
+| AdjustContrast | factors                   | List\[float]                      |
+| AdjustBrightness|factors                   | List\[float]                      |
+| AverageBlur    | kernel_sizes              | List\[Union\[Tuple\[int, int], int]] |
+| GaussianBlur   | kernel_sizes<br>sigma     | List\[Union\[Tuple\[int, int], int]]<br>Optional\[Union\[Tuple\[float, float], float]]|
+| Sharpen        | kernel_sizes              | List[int]                         |
 
 ## Aliases (Combos)
 

patta/__init__.py

@@ -1,12 +1,26 @@
 from .wrappers import (
     SegmentationTTAWrapper,
     ClassificationTTAWrapper,
-    KeypointsTTAWrapper
+    KeypointsTTAWrapper,
 )
 from .base import Compose
 
 from .transforms import (
-    HorizontalFlip, VerticalFlip, Rotate90, Scale, Add, Multiply, FiveCrops, Resize
+    HorizontalFlip,
+    VerticalFlip,
+    HorizontalShift,
+    VerticalShift,
+    Rotate90,
+    Scale,
+    Add,
+    Multiply,
+    FiveCrops,
+    Resize,
+    AdjustContrast,
+    AdjustBrightness,
+    AverageBlur,
+    GaussianBlur,
+    Sharpen,
 )
 
 from . import aliases

patta/functional.py

@@ -1,6 +1,9 @@
+from typing import Optional, Tuple, Union
+
+import cv2
+import numpy as np
 import paddle
 import paddle.nn.functional as F
-import numpy as np
 
 
 def rot90(x, k=1):
@@ -23,6 +26,16 @@ def vflip(x):
     return x.flip([2])
 
 
+def hshift(x, shifts=0):
+    """shift batch of images horizontally"""
+    return paddle.roll(x, int(shifts*x.shape[3]), axis=3)
+
+
+def vshift(x, shifts=0):
+    """shift batch of images vertically"""
+    return paddle.roll(x, int(shifts*x.shape[2]), axis=2)
+
+
 def sum(x1, x2):
     """sum of two tensors"""
     return x1 + x2
@@ -53,9 +66,7 @@ def scale(x, scale_factor, interpolation="nearest", align_corners=False):
     h, w = x.shape[2:]
     new_h = int(h * scale_factor)
     new_w = int(w * scale_factor)
-    return F.interpolate(
-        x, size=(new_h, new_w), mode=interpolation, align_corners=align_corners
-    )
+    return F.interpolate(x, size=(new_h, new_w), mode=interpolation, align_corners=align_corners)
 
 
 def resize(x, size, interpolation="nearest", align_corners=False):
@@ -104,6 +115,111 @@ def center_crop(x, crop_h, crop_w):
     return x[:, :, y_min:y_max, x_min:x_max]
 
 
+def adjust_contrast(x, contrast_factor: float=1.):
+    """adjusts contrast for batch of images"""
+    table = np.array([
+        (i - 74) * contrast_factor + 74
+        for i in range(0, 256)
+    ]).clip(0, 255).astype(np.uint8)
+    try:
+        x = x.paddle.to_tensor(x).numpy()
+    except:
+        x = x.numpy()
+    x = x.clip(0,255).astype(np.uint8)
+    x = cv2.LUT(x, table)
+    x = x.astype(np.float32)
+    return paddle.to_tensor(x)
+
+
+def adjust_brightness(x, brightness_factor: float=1.):
+    """adjusts brightness for batch of images"""
+    table = np.array([
+        i * brightness_factor
+        for i in range(0, 256)
+    ]).clip(0, 255).astype(np.uint8)
+    try:
+        x = x.paddle.to_tensor(x).numpy()
+    except:
+        x = x.numpy()
+    x = x.clip(0,255).astype(np.uint8)
+    x = cv2.LUT(x, table)
+    x = x.astype(np.float32)
+    return paddle.to_tensor(x)
+
+def filter2d(x, kernel):
+    """applies an arbitrary linear filter to an image"""
+    C = x.shape[1]
+    kernel = kernel.reshape((1, 1, *kernel.shape))
+    kernel = paddle.concat([kernel for _ in range(C)], axis=0)
+    return F.conv2d(x, kernel, stride=1, padding="same", groups=C)
+
+
+def average_blur(x, kernel_size: Union[Tuple[int, int], int] = 3):
+    """smooths the input image"""
+    cv2.blur
+    if isinstance(kernel_size, int):
+        kernel_size = (kernel_size, kernel_size)
+    if kernel_size == (1, 1):
+        return x
+    assert (
+        kernel_size[0] > 0 and kernel_size[1] > 0 and kernel_size[0] % 2 == 1 and kernel_size[1] % 2 == 1
+    ), "kernel_size both must be positive and odd"
+    kernel = np.ones(kernel_size[0] * kernel_size[1], dtype=np.float32)
+    kernel = kernel / (kernel_size[0] * kernel_size[1])
+    kernel = kernel.reshape((kernel_size[1], kernel_size[0]))
+    assert kernel.shape == (kernel_size[1], kernel_size[0])
+    kernel = paddle.to_tensor(kernel)
+    return filter2d(x, kernel)
+
+
+def gaussian_blur(
+    x,
+    kernel_size: Union[Tuple[int, int], int] = 3,
+    sigma: Optional[Union[Tuple[float, float], float]] = None
+):
+    """smooths the input image with the specified Gaussian kernel"""
+    cv2.GaussianBlur
+    if isinstance(kernel_size, int):
+        kernel_size = (kernel_size, kernel_size)
+    if kernel_size == (1, 1):
+        return x
+    if sigma is None:
+        sigma = (kernel_size[0] - 1) / 4.0
+    if isinstance(sigma, (int, float)):
+        sigma = (sigma, sigma)
+    assert kernel_size[0] > 0 and kernel_size[1] > 0 and kernel_size[0] % 2 == 1 and \
+           kernel_size[1] % 2 == 1, "kernel_size both must be positive and odd"
+    kernel_x = cv2.getGaussianKernel(kernel_size[0], sigma[0])
+    kernel_y = cv2.getGaussianKernel(kernel_size[1], sigma[1])
+    kernel = kernel_y @ kernel_x.transpose()
+    kernel = kernel.astype(np.float32)
+    assert kernel.shape == (kernel_size[1], kernel_size[0])
+    kernel = paddle.to_tensor(kernel)
+    return filter2d(x, kernel)
+
+
+def sharpen(x, kernel_size: int = 3):
+    """sharpen the input image"""
+    if kernel_size == 1:
+        return x
+    assert kernel_size > 0 and kernel_size % 2 == 1, "kernel_size both must be positive and odd"
+    kernel = get_laplacian_kernel(kernel_size).astype(np.float32)
+    kernel = kernel.astype(np.float32)
+    kernel = paddle.to_tensor(kernel)
+    x_laplacian = filter2d(x, kernel)
+    x = x - x_laplacian
+    x = x.clip(0, 255)
+    return x
+
+
+def get_laplacian_kernel(kernel_size):
+    assert kernel_size > 0 and kernel_size % 2 == 1, "kernel_size both must be positive and odd"
+    kernel = np.ones((kernel_size, kernel_size))
+    mid = kernel_size // 2
+    kernel[mid, mid] = 1 - kernel_size ** 2
+    return kernel
+
+
 def _disassemble_keypoints(keypoints):
     x = keypoints[:, 0]
     y = keypoints[:, 1]
@@ -124,6 +240,16 @@ def keypoints_vflip(keypoints):
     return _assemble_keypoints(x, 1. - y)
 
 
+def keypoints_hshift(keypoints, shifts):
+    x, y = _disassemble_keypoints(keypoints)
+    return _assemble_keypoints((x + shifts) % 1, y)
+
+
+def keypoints_vshift(keypoints, shifts):
+    x, y = _disassemble_keypoints(keypoints)
+    return _assemble_keypoints(x, (y + shifts) % 1)
+
+
 def keypoints_rot90(keypoints, k=1):
 
     if k not in {0, 1, 2, 3}:

patta/transforms.py

@@ -58,6 +58,54 @@ def apply_deaug_keypoints(self, keypoints, apply=False, **kwargs):
         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
 
@@ -263,3 +311,114 @@ def apply_deaug_mask(self, mask, **kwargs):
 
     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)