<commit_before>from collections import Iterable
from numbers import Number
from types import FunctionType
import random

import numpy as np
from skimage.transform import warp, resize, AffineTransform


class ImageAugmentation:
    def __init__(self, augmentations):
        self.augmentations = augmentations

    def __call__(self, batch):
        xb, yb = batch
        assert (len(xb) == len(yb))
        for idx in range(len(xb)):
            x, y = xb[idx], yb[idx]
            for aug in self.augmentations:
                x, y = aug(x, y)
            xb[idx], yb[idx] = x, y
        return (xb, yb)


class RandomWarpAugmentation:
    """
    Perform random warping transformation on the input data.

    Parameters can be either constant values, a list/tuple containing the lower and upper bounds
    for a uniform distribution or a value generating function:

    Examples:
        * RandomWarpAugmentation(rotation=0.5 * np.pi)
        * RandomWarpAugmentation(rotation=(-0.25 * np.pi, 0.25 * np.pi))
        * RandomWarpAugmentation(rotation=lambda: np.random.normal(0, np.pi))

    Args:
        fliph_probability: probability of random flips on horizontal (first) axis
        flipv_probability: probability of random flips on vertial (second) axis
        translation: translation of image data among all axis
        rotation: rotation angle in counter-clockwise direction as radians
        scale: scale as proportion of input size
        shear: shear angle in counter-clockwise direction as radians.
        use_diff: whether to use diffeomorphism augmentation (also known as elastic distortion)
        diff_scale: scale parameter of diffeomorphism augmentation
        diff_alpha: alpha parameter of diffeomorphism augmentation
        diff_fix_border: fix_border parameter of diffeomorphism augmentation
        augment_x: whether to augment input data
        augment_y: whether to augment label data
    """
    def __init__(self,
                 fliph_probability=0.5,
                 flipv_probability=0.5,
                 translation=(-5, 5),
                 rotation=(-np.pi / 8, np.pi / 8),
                 scale=(0.9, 1.1),
                 shear=(-0.1 * np.pi, 0.1 * np.pi),
                 use_diff=True,
                 diff_scale=8,
                 diff_alpha=.75,
                 diff_fix_border=False,
                 augment_x=True,
                 augment_y=False):
        self.fliph_probability = self.parse_parameter(fliph_probability)
        self.flipv_probability = self.parse_parameter(flipv_probability)
        self.translation = self.parse_parameter(translation)
        self.rotation = self.parse_parameter(rotation)
        self.scale = self.parse_parameter(scale)
        self.shear = self.parse_parameter(shear)
        self.use_diff = self.parse_parameter(use_diff)
        self.diff_scale = self.parse_parameter(diff_scale)
        self.diff_alpha = self.parse_parameter(diff_alpha)
        self.diff_fix_border = self.parse_parameter(diff_fix_border)
        self.augment_x = augment_x
        self.augment_y = augment_y

    @staticmethod
    def parse_parameter(param):
        if isinstance(param, Iterable):
            if len(param) != 2:
                raise ValueError('lower and upper bound required')
            lower, upper = param
            return lambda: np.random.uniform(lower, upper)
        elif isinstance(param, Number):
            return lambda: param
        elif isinstance(param, FunctionType):
            return param
        else:<commit_msg>Don't use augementations by default<commit_after>from collections import Iterable
from numbers import Number
from types import FunctionType
import random

import numpy as np
from skimage.transform import warp, resize, AffineTransform


class ImageAugmentation:
    def __init__(self, augmentations):
        self.augmentations = augmentations

    def __call__(self, batch):
        xb, yb = batch
        assert (len(xb) == len(yb))
        for idx in range(len(xb)):
            x, y = xb[idx], yb[idx]
            for aug in self.augmentations:
                x, y = aug(x, y)
            xb[idx], yb[idx] = x, y
        return (xb, yb)


class RandomWarpAugmentation:
    """
    Perform random warping transformation on the input data.

    Parameters can be either constant values, a list/tuple containing the lower and upper bounds
    for a uniform distribution or a value generating function:

    Examples:
        * RandomWarpAugmentation(rotation=0.5 * np.pi)
        * RandomWarpAugmentation(rotation=(-0.25 * np.pi, 0.25 * np.pi))
        * RandomWarpAugmentation(rotation=lambda: np.random.normal(0, np.pi))
        
    Sensible starting values for parameter tuning:
        * fliph_probability = 0.5
        * flipv_probability = 0.5
        * translation = (-5, 5)
        * rotation = (-np.pi / 8, np.pi / 8)
        * scale= (0.9, 1.1)
        * shear = (-0.1 * np.pi, 0.1 * np.pi)
        * diff_scale = 8
        * diff_alpha = .75

    Args:
        fliph_probability: probability of random flips on horizontal (first) axis
        flipv_probability: probability of random flips on vertial (second) axis
        translation: translation of image data among all axis
        rotation: rotation angle in counter-clockwise direction as radians
        scale: scale as proportion of input size
        shear: shear angle in counter-clockwise direction as radians.
        use_diff: whether to use diffeomorphism augmentation (also known as elastic distortion)
        diff_scale: scale parameter of diffeomorphism augmentation
        diff_alpha: alpha parameter of diffeomorphism augmentation
        diff_fix_border: fix_border parameter of diffeomorphism augmentation
        augment_x: whether to augment input data
        augment_y: whether to augment label data
    """
    def __init__(self,
                 fliph_probability=0.,
                 flipv_probability=0.,
                 translation=(0, 0),
                 rotation=(0, 0),
                 scale=(1, 1),
                 shear=(0, 0),
                 use_diff=False,
                 diff_scale=1,
                 diff_alpha=1,
                 diff_fix_border=False,
                 augment_x=True,
                 augment_y=False):
        self.fliph_probability = self.parse_parameter(fliph_probability)
        self.flipv_probability = self.parse_parameter(flipv_probability)
        self.translation = self.parse_parameter(translation)
        self.rotation = self.parse_parameter(rotation)
        self.scale = self.parse_parameter(scale)
        self.shear = self.parse_parameter(shear)
        self.use_diff = self.parse_parameter(use_diff)
        self.diff_scale = self.parse_parameter(diff_scale)
        self.diff_alpha = self.parse_parameter(diff_alpha)
        self.diff_fix_border = self.parse_parameter(diff_fix_border)
        self.augment_x = augment_x
        self.augment_y = augment_y

    @staticmethod
    def parse_parameter(param):
        if isinstance(param, Iterable):
            if len(param) != 2:
                raise ValueError('lower and upper bound required')
            lower, upper = param
            return lambda: np.random.uniform(lower, upper)
        elif isinstance(param, Number):
            return lambda: param
        elif isinstance(param, FunctionType):
            return param
        else:
{'fraction_non_alphanumeric': 0.07693349586888798, 'fraction_numerical': 0.009481240688067182, 'mean_word_length': 2.3366470854044286, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3686309', 'n_tokens_mistral': 2058, 'n_tokens_neox': 1977, 'n_words': 583}