{"cells":[{"metadata":{},"cell_type":"markdown","source":"# ReadMe\n\nThis is my attempt to implement [AugMix](https://arxiv.org/pdf/1912.02781.pdf) on TPU. In this notebook I implemented data augmentation part whichs seems to be working well. \nHowever, AugMix performs better when used with special loss function (Jensen-Shannon Divergence Consistency Loss). While experimenting with custom implementation of this loss using optimized training loop from [this notebook](https://www.kaggle.com/mgornergoogle/custom-training-loop-with-100-flowers-on-tpu#Optimized-custom-training-loop), I encountered significant memory issues what made it pretty useless for the competetition and thus I did not include this loss function. \n\n\n\nAugMix utilizes simple augmentation operations which are stochastically sampled and layered to produce a high diversity of augmented images. \n\n![visualization of augmix](https://i.ibb.co/YNfsHPF/Capture.png)\nAbove image is from original paper. https://arxiv.org/pdf/1912.02781.pdf\n\nThis is also my first contact with tensorflow (micro project), so if you spot any errors and mistakes, please report in the comments section. "},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import math, re, os\nimport tensorflow as tf, tensorflow.keras.backend as K\nimport tensorflow_probability as tfp\nimport numpy as np\nfrom matplotlib import pyplot as plt\nfrom kaggle_datasets import KaggleDatasets\nprint(\"Tensorflow version \" + tf.__version__)\nAUTO = tf.data.experimental.AUTOTUNE","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"# Detect hardware, return appropriate distribution strategy\ntry:\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()  # TPU detection. No parameters necessary if TPU_NAME environment variable is set. On Kaggle this is always the case.\n    print('Running on TPU ', tpu.master())\nexcept ValueError:\n    tpu = None\n\nif tpu:\n    tf.config.experimental_connect_to_cluster(tpu)\n    tf.tpu.experimental.initialize_tpu_system(tpu)\n    strategy = tf.distribute.experimental.TPUStrategy(tpu)\nelse:\n    strategy = tf.distribute.get_strategy() # default distribution strategy in Tensorflow. Works on CPU and single GPU.\n\nprint(\"REPLICAS: \", strategy.num_replicas_in_sync)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"GCS_DS_PATH = KaggleDatasets().get_gcs_path('flower-classification-with-tpus') # you can list the bucket with \"!gsutil ls $GCS_DS_PATH\"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Configuration"},{"metadata":{"trusted":true},"cell_type":"code","source":"IMAGE_SIZE = [512, 512]\n\nGCS_PATH_SELECT = { # available image sizes\n    192: GCS_DS_PATH + '/tfrecords-jpeg-192x192',\n    224: GCS_DS_PATH + '/tfrecords-jpeg-224x224',\n    331: GCS_DS_PATH + '/tfrecords-jpeg-331x331',\n    512: GCS_DS_PATH + '/tfrecords-jpeg-512x512'\n}\nGCS_PATH = GCS_PATH_SELECT[IMAGE_SIZE[0]]\n\nTRAINING_FILENAMES = tf.io.gfile.glob(GCS_PATH + '/train/*.tfrec')\nVALIDATION_FILENAMES = tf.io.gfile.glob(GCS_PATH + '/val/*.tfrec')\nTEST_FILENAMES = tf.io.gfile.glob(GCS_PATH + '/test/*.tfrec') # predictions on this dataset should be submitted for the competition\n\nCLASSES = ['pink primrose',    'hard-leaved pocket orchid', 'canterbury bells', 'sweet pea',     'wild geranium',     'tiger lily',           'moon orchid',              'bird of paradise', 'monkshood',        'globe thistle',         # 00 - 09\n           'snapdragon',       \"colt's foot\",               'king protea',      'spear thistle', 'yellow iris',       'globe-flower',         'purple coneflower',        'peruvian lily',    'balloon flower',   'giant white arum lily', # 10 - 19\n           'fire lily',        'pincushion flower',         'fritillary',       'red ginger',    'grape hyacinth',    'corn poppy',           'prince of wales feathers', 'stemless gentian', 'artichoke',        'sweet william',         # 20 - 29\n           'carnation',        'garden phlox',              'love in the mist', 'cosmos',        'alpine sea holly',  'ruby-lipped cattleya', 'cape flower',              'great masterwort', 'siam tulip',       'lenten rose',           # 30 - 39\n           'barberton daisy',  'daffodil',                  'sword lily',       'poinsettia',    'bolero deep blue',  'wallflower',           'marigold',                 'buttercup',        'daisy',            'common dandelion',      # 40 - 49\n           'petunia',          'wild pansy',                'primula',          'sunflower',     'lilac hibiscus',    'bishop of llandaff',   'gaura',                    'geranium',         'orange dahlia',    'pink-yellow dahlia',    # 50 - 59\n           'cautleya spicata', 'japanese anemone',          'black-eyed susan', 'silverbush',    'californian poppy', 'osteospermum',         'spring crocus',            'iris',             'windflower',       'tree poppy',            # 60 - 69\n           'gazania',          'azalea',                    'water lily',       'rose',          'thorn apple',       'morning glory',        'passion flower',           'lotus',            'toad lily',        'anthurium',             # 70 - 79\n           'frangipani',       'clematis',                  'hibiscus',         'columbine',     'desert-rose',       'tree mallow',          'magnolia',                 'cyclamen ',        'watercress',       'canna lily',            # 80 - 89\n           'hippeastrum ',     'bee balm',                  'pink quill',       'foxglove',      'bougainvillea',     'camellia',             'mallow',                   'mexican petunia',  'bromelia',         'blanket flower',        # 90 - 99\n           'trumpet creeper',  'blackberry lily',           'common tulip',     'wild rose']                                                                                                                                               # 100 - 102","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataset functions"},{"metadata":{"trusted":true},"cell_type":"code","source":"def decode_image(image_data):\n    image = tf.image.decode_jpeg(image_data, channels=3)\n    image = tf.cast(image, tf.float32) / 255.0  # convert image to floats in [0, 1] range\n    image = tf.reshape(image, [*IMAGE_SIZE, 3]) # explicit size needed for TPU\n    return image\n\ndef read_labeled_tfrecord(example):\n    LABELED_TFREC_FORMAT = {\n        \"image\": tf.io.FixedLenFeature([], tf.string), # tf.string means bytestring\n        \"class\": tf.io.FixedLenFeature([], tf.int64),  # shape [] means single element\n    }\n    example = tf.io.parse_single_example(example, LABELED_TFREC_FORMAT)\n    image = decode_image(example['image'])\n    label = tf.cast(example['class'], tf.int32)\n    return image, label # returns a dataset of (image, label) pairs\n\ndef read_unlabeled_tfrecord(example):\n    UNLABELED_TFREC_FORMAT = {\n        \"image\": tf.io.FixedLenFeature([], tf.string), # tf.string means bytestring\n        \"id\": tf.io.FixedLenFeature([], tf.string),  # shape [] means single element\n        # class is missing, this competitions's challenge is to predict flower classes for the test dataset\n    }\n    example = tf.io.parse_single_example(example, UNLABELED_TFREC_FORMAT)\n    image = decode_image(example['image'])\n    idnum = example['id']\n    return image, idnum # returns a dataset of image(s)\n\ndef load_dataset(filenames, labeled = True, ordered = False):\n    # Read from TFRecords. For optimal performance, reading from multiple files at once and\n    # Diregarding data order. Order does not matter since we will be shuffling the data anyway\n    \n    ignore_order = tf.data.Options()\n    if not ordered:\n        ignore_order.experimental_deterministic = False # disable order, increase speed\n        \n    dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads = AUTO) # automatically interleaves reads from multiple files\n    dataset = dataset.with_options(ignore_order) # use data as soon as it streams in, rather than in its original order\n    dataset = dataset.map(read_labeled_tfrecord if labeled else read_unlabeled_tfrecord, num_parallel_calls = AUTO) # returns a dataset of (image, label) pairs if labeled = True or (image, id) pair if labeld = False\n    return dataset\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Data visualization functions"},{"metadata":{"trusted":true},"cell_type":"code","source":"# numpy and matplotlib defaults\nnp.set_printoptions(threshold=15, linewidth=80)\n\ndef batch_to_numpy_images_and_labels(data):\n    try:\n        images, labels = data\n        numpy_labels = labels.numpy()\n        if numpy_labels.dtype == object: # binary string in this case, these are image ID strings\n            numpy_labels = [None for _ in enumerate(numpy_images)]\n    except:\n        images = data\n        numpy_labels = None\n    numpy_images = images.numpy()\n    return numpy_images, numpy_labels\n\ndef title_from_label_and_target(label, correct_label):\n    if correct_label is None:\n        return CLASSES[label], True\n    correct = (label == correct_label)\n    return \"{} [{}{}{}]\".format(CLASSES[label], 'OK' if correct else 'NO', u\"\\u2192\" if not correct else '',\n                                CLASSES[correct_label] if not correct else ''), correct\n\ndef display_one_flower(image, title, subplot, red=False, titlesize=16):\n    plt.subplot(*subplot)\n    plt.axis('off')\n    plt.imshow(image)\n    if len(title) > 0:\n        plt.title(title, fontsize=int(titlesize) if not red else int(titlesize/1.2), color='red' if red else 'black', fontdict={'verticalalignment':'center'}, pad=int(titlesize/1.5))\n    return (subplot[0], subplot[1], subplot[2]+1)\n    \ndef display_batch_of_images(databatch, predictions=None, figsize  = 13.0):\n    \"\"\"This will work with:\n    display_batch_of_images(images)\n    display_batch_of_images(images, predictions)\n    display_batch_of_images((images, labels))\n    display_batch_of_images((images, labels), predictions)\n    \"\"\"\n    # data\n    images, labels = batch_to_numpy_images_and_labels(databatch)\n    if labels is None:\n        labels = [None for _ in enumerate(images)]\n        \n    # auto-squaring: this will drop data that does not fit into square or square-ish rectangle\n    rows = int(math.sqrt(len(images)))\n    cols = len(images)//rows\n        \n    # size and spacing\n    FIGSIZE =  figsize\n    SPACING = 0.1\n    subplot=(rows,cols,1)\n    if rows < cols:\n        plt.figure(figsize=(FIGSIZE,FIGSIZE/cols*rows))\n    else:\n        plt.figure(figsize=(FIGSIZE/rows*cols,FIGSIZE))\n    \n    # display\n    for i, (image, label) in enumerate(zip(images[:rows*cols], labels[:rows*cols])):\n        title = '' if label is None else CLASSES[label]\n        correct = True\n        if predictions is not None:\n            title, correct = title_from_label_and_target(predictions[i], label)\n        dynamic_titlesize = FIGSIZE*SPACING/max(rows,cols)*40+3 # magic formula tested to work from 1x1 to 10x10 images\n        subplot = display_one_flower(image, title, subplot, not correct, titlesize=dynamic_titlesize)\n    \n    #layout\n    plt.tight_layout()\n    if label is None and predictions is None:\n        plt.subplots_adjust(wspace=0, hspace=0)\n    else:\n        plt.subplots_adjust(wspace=SPACING, hspace=SPACING)\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Data augmentation"},{"metadata":{},"cell_type":"markdown","source":"## Helper functions"},{"metadata":{"trusted":true},"cell_type":"code","source":"def int_parameter(level, maxval):\n    return tf.cast(level * maxval / 10, tf.int32)\n\ndef float_parameter(level, maxval):\n    return tf.cast((level) * maxval / 10., tf.float32)\n\ndef sample_level(n):\n    return tf.random.uniform(shape=[1], minval=0.1, maxval=n, dtype=tf.float32)\n    \ndef affine_transform(image, transform_matrix):\n    # input image - is one image of size [dim,dim,3] not a batch of [b,dim,dim,3]\n    DIM = IMAGE_SIZE[0]\n    XDIM = DIM%2 #fix for size 331\n    \n    x = tf.repeat(tf.range(DIM//2,-DIM//2,-1), DIM)\n    y = tf.tile(tf.range(-DIM//2,DIM//2), [DIM])\n    z = tf.ones([DIM*DIM], dtype='int32')\n    idx = tf.stack([x, y, z])\n    \n    # ROTATE DESTINATION PIXELS ONTO ORIGIN PIXELS\n    idx2 = K.dot(transform_matrix, tf.cast(idx, dtype='float32'))\n    idx2 = K.cast(idx2, dtype='int32')\n    idx2 = K.clip(idx2, -DIM//2+XDIM+1, DIM//2)\n    \n    # FIND ORIGIN PIXEL VALUES           \n    idx3 = tf.stack([DIM//2-idx2[0,], DIM//2-1+idx2[1,]])\n    d = tf.gather_nd(image, tf.transpose(idx3))\n    return tf.reshape(d,[DIM,DIM,3])\n\ndef blend(image1, image2, factor):\n    if factor == 0.0:\n        return tf.convert_to_tensor(image1)\n    if factor == 1.0:\n        return tf.convert_to_tensor(image2)\n\n    image1 = tf.cast(image1, tf.float32)\n    image2 = tf.cast(image2, tf.float32)\n\n    difference = image2 - image1\n    scaled = factor * difference\n\n    # Do addition in float.\n    temp = tf.cast(image1, tf.float32) + scaled\n\n    # Interpolate\n    if factor > 0.0 and factor < 1.0:\n        # Interpolation means we always stay within 0 and 255.\n        return tf.cast(temp, tf.uint8)\n\n    # Extrapolate:\n    #\n    # We need to clip and then cast.\n    return tf.cast(tf.clip_by_value(temp, 0.0, 255.0), tf.uint8)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Transformations\nThese are simple augmentations used by AugMix. Every function takes ```image``` and ```level``` (integer from 1 to 10) as arguments. The second one indicates how much variation will particular transformation yield, in other words, how strong it will be.\n\nTranslate, shear and rotate augmentations are based on [this notebook](https://www.kaggle.com/cdeotte/rotation-augmentation-gpu-tpu-0-96)."},{"metadata":{"trusted":true},"cell_type":"code","source":"def rotate(image, level):\n    degrees = float_parameter(sample_level(level), 30)\n    rand_var = tf.random.uniform(shape=[], dtype=tf.float32)\n    degrees = tf.cond(rand_var > 0.5, lambda: degrees, lambda: -degrees)\n\n    angle = math.pi*degrees/180 # convert degrees to radians\n    angle = tf.cast(angle, tf.float32)\n    # define rotation matrix\n    c1 = tf.math.cos(angle)\n    s1 = tf.math.sin(angle)\n    one = tf.constant([1],dtype='float32')\n    zero = tf.constant([0],dtype='float32')\n    rotation_matrix = tf.reshape(tf.concat([c1,s1,zero, -s1,c1,zero, zero,zero,one], axis=0), [3,3])\n\n    transformed = affine_transform(image, rotation_matrix)\n    return transformed\n\ndef translate_x(image, level):\n    lvl = int_parameter(sample_level(level), IMAGE_SIZE[0] / 3)\n    rand_var = tf.random.uniform(shape=[], dtype=tf.float32)\n    lvl = tf.cond(rand_var > 0.5, lambda: lvl, lambda: -lvl)\n\n    one = tf.constant([1], dtype='float32')\n    zero = tf.constant([0], dtype='float32')\n    lvl = tf.cast(lvl, tf.float32)\n    translate_x_matrix = tf.reshape(tf.concat([one,zero,zero, zero,one,lvl, zero,zero,one], axis=0), [3,3])\n\n    transformed = affine_transform(image, translate_x_matrix)\n    return transformed\n\ndef translate_y(image, level):\n    lvl = int_parameter(sample_level(level), IMAGE_SIZE[0] / 3)\n    rand_var = tf.random.uniform(shape=[], dtype=tf.float32)\n    lvl = tf.cond(rand_var > 0.5, lambda: lvl, lambda: -lvl)\n\n    one = tf.constant([1], dtype='float32')\n    zero = tf.constant([0], dtype='float32')\n    lvl = tf.cast(lvl, tf.float32)\n    translate_y_matrix = tf.reshape(tf.concat([one,zero,lvl, zero,one,zero, zero,zero,one], axis=0), [3,3])\n\n    transformed = affine_transform(image, translate_y_matrix)\n    return transformed\n\ndef shear_x(image, level):\n    lvl = float_parameter(sample_level(level), 0.3)\n    rand_var = tf.random.uniform(shape=[], dtype=tf.float32)\n    lvl = tf.cond(rand_var > 0.5, lambda: lvl, lambda: -lvl)\n\n    one = tf.constant([1], dtype='float32')\n    zero = tf.constant([0], dtype='float32')\n    s2 = tf.math.sin(lvl)\n    shear_x_matrix = tf.reshape(tf.concat([one,s2,zero, zero,one,zero, zero,zero,one],axis=0), [3,3])   \n\n    transformed = affine_transform(image, shear_x_matrix)\n    return transformed\n\ndef shear_y(image, level):\n    lvl = float_parameter(sample_level(level), 0.3)\n    rand_var = tf.random.uniform(shape=[], dtype=tf.float32)\n    lvl = tf.cond(rand_var > 0.5, lambda: lvl, lambda: -lvl)\n\n    one = tf.constant([1], dtype='float32')\n    zero = tf.constant([0], dtype='float32')\n    c2 = tf.math.cos(lvl)\n    shear_y_matrix = tf.reshape(tf.concat([one,zero,zero, zero,c2,zero, zero,zero,one],axis=0), [3,3])   \n    \n    transformed = affine_transform(image, shear_y_matrix)\n    return transformed\n\ndef solarize(image, level):\n    # For each pixel in the image, select the pixel\n    # if the value is less than the threshold.\n    # Otherwise, subtract 255 from the pixel.\n    threshold = float_parameter(sample_level(level), 1)\n    return tf.where(image < threshold, image, 1 - image)\n\ndef solarize_add(image, level):\n    # For each pixel in the image less than threshold\n    # we add 'addition' amount to it and then clip the\n    # pixel value to be between 0 and 255. The value\n    # of 'addition' is between -128 and 128.\n    threshold = float_parameter(sample_level(level), 1)\n    addition = float_parameter(sample_level(level), 0.5)\n    rand_var = tf.random.uniform(shape=[], dtype=tf.float32)\n    addition = tf.cond(rand_var > 0.5, lambda: addition, lambda: -addition)\n\n    added_image = tf.cast(image, tf.float32) + addition\n    added_image = tf.cast(tf.clip_by_value(added_image, 0, 1), tf.float32)\n    return tf.where(image < threshold, added_image, image)\n\ndef posterize(image, level):\n    lvl = int_parameter(sample_level(level), 8)\n    shift = 8 - lvl\n    shift = tf.cast(shift, tf.uint8)\n    image = tf.cast(tf.math.scalar_mul(255, image), tf.uint8)\n    image = tf.bitwise.left_shift(tf.bitwise.right_shift(image, shift), shift)\n    return tf.cast(tf.clip_by_value(tf.math.divide(image, 255), 0, 1), tf.float32)\n\ndef autocontrast(image, _):\n    image = tf.cast(tf.math.scalar_mul(255, image), tf.uint8)\n\n    def scale_channel(image):\n        # A possibly cheaper version can be done using cumsum/unique_with_counts\n        # over the histogram values, rather than iterating over the entire image.\n        # to compute mins and maxes.\n        lo = tf.cast(tf.reduce_min(image), tf.float32)\n        hi = tf.cast(tf.reduce_max(image), tf.float32)\n\n        # Scale the image, making the lowest value 0 and the highest value 255.\n        def scale_values(im):\n            scale = 255.0 / (hi - lo)\n            offset = -lo * scale\n            im = tf.cast(im, tf.float32) * scale + offset\n            im = tf.clip_by_value(im, 0.0, 255.0)\n            return tf.cast(im, tf.uint8)\n\n        result = tf.cond(hi > lo, lambda: scale_values(image), lambda: image)\n        return result\n\n    # Assumes RGB for now.  Scales each channel independently\n    # and then stacks the result.\n    s1 = scale_channel(image[:, :, 0])\n    s2 = scale_channel(image[:, :, 1])\n    s3 = scale_channel(image[:, :, 2])\n    image = tf.stack([s1, s2, s3], 2)\n    return tf.cast(tf.clip_by_value(tf.math.divide(image, 255), 0, 1), tf.float32)\n\ndef equalize(image, _):\n    image = tf.cast(tf.math.scalar_mul(255, image), tf.uint8)\n\n    def scale_channel(im, c):\n        im = tf.cast(im[:, :, c], tf.int32)\n        # Compute the histogram of the image channel.\n        histo = tf.histogram_fixed_width(im, [0, 255], nbins=256)\n        # For the purposes of computing the step, filter out the nonzeros.\n        nonzero = tf.where(tf.not_equal(histo, 0))\n        nonzero_histo = tf.reshape(tf.gather(histo, nonzero), [-1])\n        step = (tf.reduce_sum(nonzero_histo) - nonzero_histo[-1]) // 255\n\n        def build_lut(histo, step):\n            # Compute the cumulative sum, shifting by step // 2\n            # and then normalization by step.\n            lut = (tf.cumsum(histo) + (step // 2)) // step\n            # Shift lut, prepending with 0.\n            lut = tf.concat([[0], lut[:-1]], 0)\n            # Clip the counts to be in range.  This is done\n            # in the C code for image.point.\n            return tf.clip_by_value(lut, 0, 255)\n\n        # If step is zero, return the original image.  Otherwise, build\n        # lut from the full histogram and step and then index from it.\n        result = tf.cond(tf.equal(step, 0),\n                        lambda: im,\n                        lambda: tf.gather(build_lut(histo, step), im))\n\n        return tf.cast(result, tf.uint8)\n\n    # Assumes RGB for now.  Scales each channel independently\n    # and then stacks the result.\n    s1 = scale_channel(image, 0)\n    s2 = scale_channel(image, 1)\n    s3 = scale_channel(image, 2)\n    image = tf.stack([s1, s2, s3], 2)\n\n    return tf.cast(tf.clip_by_value(tf.math.divide(image, 255), 0, 1), tf.float32)\n\ndef color(image, level):\n    factor = float_parameter(sample_level(level), 1.8) + 0.1\n    image = tf.cast(tf.math.scalar_mul(255, image), tf.uint8)\n    degenerate = tf.image.grayscale_to_rgb(tf.image.rgb_to_grayscale(image))\n    blended = blend(degenerate, image, factor)\n    return tf.cast(tf.clip_by_value(tf.math.divide(blended, 255), 0, 1), tf.float32)\n\ndef brightness(image, level):\n    delta = float_parameter(sample_level(level), 0.5) + 0.1\n    rand_var = tf.random.uniform(shape=[], dtype=tf.float32)\n    delta = tf.cond(rand_var > 0.5, lambda: delta, lambda: -delta) \n    return tf.image.adjust_brightness(image, delta=delta)\n\ndef contrast(image, level):\n    factor = float_parameter(sample_level(level), 1.8) + 0.1\n    factor = tf.reshape(factor, [])\n    rand_var = tf.random.uniform(shape=[], dtype=tf.float32)\n    factor = tf.cond(rand_var > 0.5, lambda: factor, lambda: 1.9 - factor  )\n\n    return tf.image.adjust_contrast(image, factor)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## AugMix"},{"metadata":{"trusted":true},"cell_type":"code","source":"means = {'R': 0.44892993872313053, 'G': 0.4148519066242368, 'B': 0.301880284715257}\nstds = {'R': 0.24393544875614917, 'G': 0.2108791383467354, 'B': 0.220427056859487}\n\ndef substract_means(image):\n    image = image - np.array([means['R'], means['G'], means['B']])\n    return image\n\ndef normalize(image):\n    image = substract_means(image)\n    image = image / np.array([stds['R'], stds['G'], stds['B']])\n    return tf.clip_by_value(image, 0, 1)\n\ndef apply_op(image, level, which):\n    # is there any better way than manually typing all of these conditions? \n    # I tried to randomly select transformation from array of functions, but tensorflow didn't let me to\n    augmented = image\n    augmented = tf.cond(which == tf.constant([0], dtype=tf.int32), lambda: rotate(image, level), lambda: augmented)\n    augmented = tf.cond(which == tf.constant([1], dtype=tf.int32), lambda: translate_x(image, level), lambda: augmented)\n    augmented = tf.cond(which == tf.constant([2], dtype=tf.int32), lambda: translate_y(image, level), lambda: augmented)\n    augmented = tf.cond(which == tf.constant([3], dtype=tf.int32), lambda: shear_x(image, level), lambda: augmented)\n    augmented = tf.cond(which == tf.constant([4], dtype=tf.int32), lambda: shear_y(image, level), lambda: augmented)\n    augmented = tf.cond(which == tf.constant([5], dtype=tf.int32), lambda: solarize_add(image, level), lambda: augmented)\n    augmented = tf.cond(which == tf.constant([6], dtype=tf.int32), lambda: solarize(image, level), lambda: augmented)\n    augmented = tf.cond(which == tf.constant([7], dtype=tf.int32), lambda: posterize(image, level), lambda: augmented)\n    augmented = tf.cond(which == tf.constant([8], dtype=tf.int32), lambda: autocontrast(image, level), lambda: augmented)\n    augmented = tf.cond(which == tf.constant([9], dtype=tf.int32), lambda: equalize(image, level), lambda: augmented)\n    augmented = tf.cond(which == tf.constant([10], dtype=tf.int32), lambda: color(image, level), lambda: augmented)\n    augmented = tf.cond(which == tf.constant([11], dtype=tf.int32), lambda: contrast(image, level), lambda: augmented)\n    augmented = tf.cond(which == tf.constant([12], dtype=tf.int32), lambda: brightness(image, level), lambda: augmented)\n    return augmented\n\ndef augmix(image):\n    # you can play with these parameters\n    severity = 3 # level of transformations as described above in transformations (integer from 1 to 10)\n    width = 3 # number of different chains of transformations to be mixed\n    depth = -1 # number of transformations in one chain, -1 means random from 1 to 3\n    \n    alpha = 1.\n    dir_dist = tfp.distributions.Dirichlet([alpha]*width)\n    ws = tf.cast(dir_dist.sample(), tf.float32)\n    beta_dist = tfp.distributions.Beta(alpha, alpha)\n    m = tf.cast(beta_dist.sample(), tf.float32)\n\n    mix = tf.zeros_like(image, dtype='float32')\n\n    def outer_loop_cond(i, depth, mix):\n        return tf.less(i, width)\n\n    def outer_loop_body(i, depth, mix):\n        image_aug = tf.identity(image)\n        depth = tf.cond(tf.greater(depth, 0), lambda: depth, lambda: tf.random.uniform(shape=[], minval=1, maxval=3, dtype=tf.int32))\n\n        def inner_loop_cond(j, image_aug):\n            return tf.less(j, depth)\n\n        def inner_loop_body(j, image_aug):\n            which = tf.random.uniform(shape=[], minval=0, maxval=3, dtype=tf.int32)\n            image_aug = apply_op(image_aug, severity, which)\n            j = tf.add(j, 1)\n            return j, image_aug\n        \n        j = tf.constant([0], dtype=tf.int32)\n        j, image_aug = tf.while_loop(inner_loop_cond, inner_loop_body, [j, image_aug])\n\n        wsi = tf.gather(ws, i)\n        mix = tf.add(mix, wsi*normalize(image_aug))\n        i = tf.add(i, 1)\n        return i, depth, mix\n\n    i = tf.constant([0], dtype=tf.int32)\n    i, depth, mix = tf.while_loop(outer_loop_cond, outer_loop_body, [i, depth, mix])\n    \n    mixed = tf.math.scalar_mul((1 - m), normalize(image)) + tf.math.scalar_mul(m, mix)\n    return tf.clip_by_value(mixed, 0, 1)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Visualization"},{"metadata":{},"cell_type":"markdown","source":"## Transformations"},{"metadata":{"trusted":true},"cell_type":"code","source":"# get sample\ndataset = load_dataset([TRAINING_FILENAMES[0]])\ndataset = dataset.batch(8)\nbatches = iter(dataset)\nfirst_batch = next(batches)[0]\n\nlevel = 10 # level of strength of transformations, from 1 to 10","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Original"},{"metadata":{"trusted":true},"cell_type":"code","source":"display_batch_of_images(first_batch)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Translate"},{"metadata":{"trusted":true},"cell_type":"code","source":"translated = tf.map_fn(lambda img: translate_x(img, level) if np.random.rand() < 0.5 else translate_y(img, level), first_batch)\ndisplay_batch_of_images(translated)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Shear"},{"metadata":{"trusted":true},"cell_type":"code","source":"sheared = tf.map_fn(lambda img: shear_x(img, level) if np.random.rand() < 0.5 else shear_y(img, level), first_batch)\ndisplay_batch_of_images(sheared)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Rotate"},{"metadata":{"trusted":true},"cell_type":"code","source":"rotated = tf.map_fn(lambda img: rotate(img, level), first_batch)\ndisplay_batch_of_images(rotated)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Solarize"},{"metadata":{"trusted":true},"cell_type":"code","source":"solarized = tf.map_fn(lambda img: solarize(img, level), first_batch)\ndisplay_batch_of_images(solarized)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Solarize Add"},{"metadata":{"trusted":true},"cell_type":"code","source":"solarized = tf.map_fn(lambda img: solarize_add(img, level), first_batch)\ndisplay_batch_of_images(solarized)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Posterize"},{"metadata":{"trusted":true},"cell_type":"code","source":"posterized = tf.map_fn(lambda img: posterize(img, level), first_batch)\ndisplay_batch_of_images(posterized)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Autocontrast"},{"metadata":{"trusted":true},"cell_type":"code","source":"autocontrasted = tf.map_fn(lambda img: autocontrast(img, level), first_batch)\ndisplay_batch_of_images(autocontrasted)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Contrast"},{"metadata":{"trusted":true},"cell_type":"code","source":"contrasted = tf.map_fn(lambda img: contrast(img, level), first_batch)\ndisplay_batch_of_images(contrasted)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Equalize"},{"metadata":{"trusted":true},"cell_type":"code","source":"equalized = tf.map_fn(lambda img: equalize(img, level), first_batch)\ndisplay_batch_of_images(equalized)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Brightness"},{"metadata":{"trusted":true},"cell_type":"code","source":"bright = tf.map_fn(lambda img: brightness(img, level), first_batch)\ndisplay_batch_of_images(bright)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Color"},{"metadata":{"trusted":true},"cell_type":"code","source":"colored = tf.map_fn(lambda img: color(img, level), first_batch)\ndisplay_batch_of_images(colored)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## AugMix"},{"metadata":{"trusted":true},"cell_type":"code","source":"augmented = tf.map_fn(lambda img: augmix(img), first_batch)\ndisplay_batch_of_images(augmented)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### More AugMix"},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset = load_dataset(VALIDATION_FILENAMES)\ndataset = dataset.map(lambda img, label: (augmix(img), label), num_parallel_calls=AUTO)\ndataset = dataset.batch(20)\nbatches = iter(dataset)\ndisplay_batch_of_images(next(batches))","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":4}