{"cells":[{"metadata":{},"cell_type":"markdown","source":"# An Implementation of Augmix Based on albumentations\n\n\nHi, here is my implementation of Augmix augmentation based on albumentations. If you find it helpful please upvote me.\nThanks!\n\n![](https://www.googleapis.com/download/storage/v1/b/kaggle-user-content/o/inbox%2F448347%2Fe9d3259d5b0ef3dadba0238bf3901f1e%2F2020-02-10%2013.50.53.png?generation=1581310282118474&alt=media)\n\nAugmix: https://arxiv.org/abs/1912.02781\n\nOfficial implementation: https://github.com/google-research/augmix\n\nalbumentations： https://github.com/albumentations-team/albumentations\n\n**Notice: I didn't normalize images in my implementation** : https://github.com/google-research/augmix/blob/master/augment_and_mix.py#L25-L30"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nimport cv2\nimport numpy as np\nimport pandas as pd\nimport albumentations\nfrom PIL import Image, ImageOps, ImageEnhance\nfrom albumentations.core.transforms_interface import ImageOnlyTransform\nfrom albumentations.augmentations import functional as F\nimport matplotlib.pyplot as plt\nimport torch\nfrom torch.utils.data import TensorDataset, DataLoader, Dataset\n\ndata_dir = '../input/bengaliai-cv19'\nfiles_train = [f'train_image_data_{fid}.parquet' for fid in range(1)]\n\nHEIGHT = 137\nWIDTH = 236","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"df_train = pd.read_csv(os.path.join(data_dir, f'train.csv'))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def read_data(files):\n    tmp = []\n    for f in files:\n        F = os.path.join(data_dir, f)\n        data = pd.read_parquet(F)\n        tmp.append(data)\n    tmp = pd.concat(tmp)\n\n    data = tmp.iloc[:, 1:].values\n    return data\n\n# train data\ndata_train = read_data(files_train)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def int_parameter(level, maxval):\n    \"\"\"Helper function to scale `val` between 0 and maxval .\n    Args:\n    level: Level of the operation that will be between [0, `PARAMETER_MAX`].\n    maxval: Maximum value that the operation can have. This will be scaled to\n      level/PARAMETER_MAX.\n    Returns:\n    An int that results from scaling `maxval` according to `level`.\n    \"\"\"\n    return int(level * maxval / 10)\n\n\ndef float_parameter(level, maxval):\n    \"\"\"Helper function to scale `val` between 0 and maxval.\n    Args:\n    level: Level of the operation that will be between [0, `PARAMETER_MAX`].\n    maxval: Maximum value that the operation can have. This will be scaled to\n      level/PARAMETER_MAX.\n    Returns:\n    A float that results from scaling `maxval` according to `level`.\n    \"\"\"\n    return float(level) * maxval / 10.\n\n\ndef sample_level(n):\n    return np.random.uniform(low=0.1, high=n)\n\n\ndef autocontrast(pil_img, _):\n    return ImageOps.autocontrast(pil_img)\n\n\ndef equalize(pil_img, _):\n    return ImageOps.equalize(pil_img)\n\n\ndef posterize(pil_img, level):\n    level = int_parameter(sample_level(level), 4)\n    return ImageOps.posterize(pil_img, 4 - level)\n\n\ndef rotate(pil_img, level):\n    degrees = int_parameter(sample_level(level), 30)\n    if np.random.uniform() > 0.5:\n        degrees = -degrees\n    return pil_img.rotate(degrees, resample=Image.BILINEAR)\n\n\ndef solarize(pil_img, level):\n    level = int_parameter(sample_level(level), 256)\n    return ImageOps.solarize(pil_img, 256 - level)\n\n\ndef shear_x(pil_img, level):\n    level = float_parameter(sample_level(level), 0.3)\n    if np.random.uniform() > 0.5:\n        level = -level\n    return pil_img.transform(pil_img.size,\n                           Image.AFFINE, (1, level, 0, 0, 1, 0),\n                           resample=Image.BILINEAR)\n\n\ndef shear_y(pil_img, level):\n    level = float_parameter(sample_level(level), 0.3)\n    if np.random.uniform() > 0.5:\n        level = -level\n    return pil_img.transform(pil_img.size,\n                           Image.AFFINE, (1, 0, 0, level, 1, 0),\n                           resample=Image.BILINEAR)\n\n\ndef translate_x(pil_img, level):\n    level = int_parameter(sample_level(level), pil_img.size[0] / 3)\n    if np.random.random() > 0.5:\n        level = -level\n    return pil_img.transform(pil_img.size,\n                           Image.AFFINE, (1, 0, level, 0, 1, 0),\n                           resample=Image.BILINEAR)\n\n\ndef translate_y(pil_img, level):\n    level = int_parameter(sample_level(level), pil_img.size[0] / 3)\n    if np.random.random() > 0.5:\n        level = -level\n    return pil_img.transform(pil_img.size,\n                           Image.AFFINE, (1, 0, 0, 0, 1, level),\n                           resample=Image.BILINEAR)\n\n\n# operation that overlaps with ImageNet-C's test set\ndef color(pil_img, level):\n    level = float_parameter(sample_level(level), 1.8) + 0.1\n    return ImageEnhance.Color(pil_img).enhance(level)\n\n\n# operation that overlaps with ImageNet-C's test set\ndef contrast(pil_img, level):\n    level = float_parameter(sample_level(level), 1.8) + 0.1\n    return ImageEnhance.Contrast(pil_img).enhance(level)\n\n\n# operation that overlaps with ImageNet-C's test set\ndef brightness(pil_img, level):\n    level = float_parameter(sample_level(level), 1.8) + 0.1\n    return ImageEnhance.Brightness(pil_img).enhance(level)\n\n\n# operation that overlaps with ImageNet-C's test set\ndef sharpness(pil_img, level):\n    level = float_parameter(sample_level(level), 1.8) + 0.1\n    return ImageEnhance.Sharpness(pil_img).enhance(level)\n\n\naugmentations = [\n    autocontrast, equalize, posterize, rotate, solarize, shear_x, shear_y,\n    translate_x, translate_y\n]\n\naugmentations_all = [\n    autocontrast, equalize, posterize, rotate, solarize, shear_x, shear_y,\n    translate_x, translate_y, color, contrast, brightness, sharpness\n]\n\ndef normalize(image):\n    \"\"\"Normalize input image channel-wise to zero mean and unit variance.\"\"\"\n    return image - 127\n\ndef apply_op(image, op, severity):\n    #   image = np.clip(image, 0, 255)\n    pil_img = Image.fromarray(image)  # Convert to PIL.Image\n    pil_img = op(pil_img, severity)\n    return np.asarray(pil_img)\n\ndef augment_and_mix(image, severity=3, width=3, depth=-1, alpha=1.):\n    \"\"\"Perform AugMix augmentations and compute mixture.\n    Args:\n    image: Raw input image as float32 np.ndarray of shape (h, w, c)\n    severity: Severity of underlying augmentation operators (between 1 to 10).\n    width: Width of augmentation chain\n    depth: Depth of augmentation chain. -1 enables stochastic depth uniformly\n      from [1, 3]\n    alpha: Probability coefficient for Beta and Dirichlet distributions.\n    Returns:\n    mixed: Augmented and mixed image.\n    \"\"\"\n    ws = np.float32(\n      np.random.dirichlet([alpha] * width))\n    m = np.float32(np.random.beta(alpha, alpha))\n\n    mix = np.zeros_like(image).astype(np.float32)\n    for i in range(width):\n        image_aug = image.copy()\n        depth = depth if depth > 0 else np.random.randint(1, 4)\n        for _ in range(depth):\n            op = np.random.choice(augmentations)\n            image_aug = apply_op(image_aug, op, severity)\n        # Preprocessing commutes since all coefficients are convex\n        mix += ws[i] * image_aug\n#         mix += ws[i] * normalize(image_aug)\n\n    mixed = (1 - m) * image + m * mix\n#     mixed = (1 - m) * normalize(image) + m * mix\n    return mixed\n\n\nclass RandomAugMix(ImageOnlyTransform):\n\n    def __init__(self, severity=3, width=3, depth=-1, alpha=1., always_apply=False, p=0.5):\n        super().__init__(always_apply, p)\n        self.severity = severity\n        self.width = width\n        self.depth = depth\n        self.alpha = alpha\n\n    def apply(self, image, **params):\n        image = augment_and_mix(\n            image,\n            self.severity,\n            self.width,\n            self.depth,\n            self.alpha\n        )\n        return image\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class BengaliDataset(Dataset):\n    def __init__(self, csv, data, idx, split, mode, image_size, transform=None):\n\n        self.csv = csv.reset_index()\n        self.data = data\n        self.idx = np.asarray(idx)\n        self.split = split\n        self.mode = mode\n        self.image_size = image_size\n        self.transform = transform\n\n    def __len__(self):\n        return self.idx.shape[0]\n\n    def __getitem__(self, index):\n        index = self.idx[index]\n        this_img_id = self.csv.iloc[index].image_id\n        \n        image = self.data[index].reshape(HEIGHT, WIDTH)\n        image = cv2.resize(image, (self.image_size, self.image_size))\n\n        if self.transform is not None:\n            res = self.transform(image=image)\n            image = res['image'].astype(np.float32)\n        else:\n            image = image.astype(np.float32)\n\n        image /= 255\n        image = image[np.newaxis, :, :]\n        image = 1 - image\n        image = np.repeat(image, 3, 0)  # 1ch to 3ch\n\n        if self.mode == 'test':\n            return torch.tensor(image)\n        else:\n            label_1 = self.csv.iloc[index].grapheme_root\n            label_2 = self.csv.iloc[index].vowel_diacritic\n            label_3 = self.csv.iloc[index].consonant_diacritic\n            label = [label_1, label_2, label_3]\n            return torch.tensor(image), torch.tensor(label)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def plot_imgs(dataset_show):\n    from pylab import rcParams\n    rcParams['figure.figsize'] = 20,10\n    for i in range(2):\n        f, axarr = plt.subplots(1,5)\n        for p in range(5):\n            idx = np.random.randint(0, len(dataset_show))\n            img, label = dataset_show[idx]\n            axarr[p].imshow(img.transpose(0, 1).transpose(1,2).squeeze())\n            axarr[p].set_title(idx)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"markdown","source":"# Usage Example\n\n### Official default setting -> Looks dirty"},{"metadata":{"trusted":true},"cell_type":"code","source":"transforms_train = albumentations.Compose([\n    RandomAugMix(severity=3, width=3, alpha=1., p=1.),\n])\n\ndf_show = df_train.iloc[:1000]\ndataset_show = BengaliDataset(df_show, data_train, list(range(df_show.shape[0])), 'train', 'train', 128, transform=transforms_train)\nplot_imgs(dataset_show)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### severity=1, width=1 -> Not that dirty?"},{"metadata":{"trusted":true},"cell_type":"code","source":"transforms_train = albumentations.Compose([\n    RandomAugMix(severity=1, width=1, p=1.),\n])\n\ndf_show = df_train.iloc[:1000]\ndataset_show = BengaliDataset(df_show, data_train, list(range(df_show.shape[0])), 'train', 'train', 128, transform=transforms_train)\nplot_imgs(dataset_show)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### severity=7, width=7, alpha=5 -> Very aggressive transform"},{"metadata":{"trusted":true},"cell_type":"code","source":"transforms_train = albumentations.Compose([\n    RandomAugMix(severity=7, width=7, alpha=5, p=1.),\n])\n\ndf_show = df_train.iloc[:1000]\ndataset_show = BengaliDataset(df_show, data_train, list(range(df_show.shape[0])), 'train', 'train', 128, transform=transforms_train)\nplot_imgs(dataset_show)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":1}