{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Introduction"},{"metadata":{},"cell_type":"markdown","source":"In this notebook, I demonstrate and provide the necessary functions required to integrate Albumentations augmentations into Tensorflow Datasets (tf.data.Dataset) pipelines that originated from .tfrec shards.\n\nIn comparison to tf.image's preprocessing module and keras's ImageDataGenerator class, Albumentations' augment library is equipped with a wider variety of augmentation functions, catering to more complex augmentations.\n\nOn top of being GPU-capable, this method also enables on-the-fly batch augmentation which makes computations exponentially faster and less memory-intensive.\n\nTo make the script future-proof and customizable, I have included locations where additional augmentation or preprocessing steps can be inserted.\n\nFor educational purposes various snippets have been included throughout the script elaborating on the steps performed. \n\nLastly, to understand the data flow, it is advised to read in a reversed manner; i.e from get_train_dataset to decode_func."},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# Import dependencies\nimport tensorflow as tf\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom kaggle_datasets import KaggleDatasets\nimport albumentations as A\nimport os\n\nfrom functools import partial\nfrom collections import Counter\n\n# Selecting 1st GPU\nos.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\"\n\n# Paths/Data\nPATH = \"../input/hpa-single-cell-image-classification\"\n\n# For GPU use this\nFILENAMES = tf.io.gfile.glob(PATH + \"/train_tfrecords/train*.tfrec\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Initialize variables\nAUTOTUNE = tf.data.experimental.AUTOTUNE\nBATCH_SIZE = 32\nIMAGE_SIZE = [256, 256]\nSEED = 42\n\nprint('Number of shards:', len(FILENAMES))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Displaying .tfrec feature dictionary"},{"metadata":{"trusted":true},"cell_type":"code","source":"def list_record_features(tfrecords_path):\n    # Dict of extracted feature information\n    features = {}\n    # Iterate records\n    for rec in tf.data.TFRecordDataset([str(tfrecords_path)]):\n        # Get record bytes\n        example_bytes = rec.numpy()\n        # Parse example protobuf message\n        example = tf.train.Example()\n        example.ParseFromString(example_bytes)\n        # Iterate example features\n        for key, value in example.features.feature.items():\n            # Kind of data in the feature\n            kind = value.WhichOneof('kind')\n            # Size of data in the feature\n            size = len(getattr(value, kind).value)\n            # Check if feature was seen before\n            if key in features:\n                # Check if values match, use None otherwise\n                kind2, size2 = features[key]\n                if kind != kind2:\n                    kind = None\n                if size != size2:\n                    size = None\n            # Save feature data\n            features[key] = (kind, size)\n    return features\n\n# Print extracted feature information\nfeatures = list_record_features(FILENAMES[0])\nprint(*features.items(), sep='\\n')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def decode_func(image):\n    # Preprocessing functions  \n    image = tf.image.decode_jpeg(image, channels=3)\n    image = tf.cast(image, tf.float32) / 255.0  \n    image = tf.image.resize(image, size=[*IMAGE_SIZE], \n            method=tf.image.ResizeMethod.NEAREST_NEIGHBOR) \n    image = tf.reshape(image, [*IMAGE_SIZE, 3])\n    return image","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def transform(image, image_size):\n    \"\"\"Customizable augmentation using Albumentations library.\"\"\"\n    \"\"\"Set always_apply to instantiate transform as a\n    preprocess transformation.\"\"\"\n    transforms = A.Compose([\n            A.Rotate(limit=40),\n            A.RandomBrightness(),\n            A.HueSaturationValue(hue_shift_limit=100,\n                                 val_shift_limit=0),\n            A.Cutout(),\n            A.RandomContrast(),\n            A.RandomCrop(height=220, width=220),\n            A.HorizontalFlip(always_apply=True)])\n    \n    # Converting image to tf.float32 type\n    \"\"\"Albumentations augmentation functions only supports\n    .uint8 and .float32 data types therefore conversion is\n    required.\"\"\"\n    image = tf.cast(x=image, dtype=tf.float32).numpy()\n\n    # Apply augmentation on each image instance\n    \"\"\"transforms function is an Albumentations Compose()\n    function that outputs a dictionary in a form of \n    {'image': image}, and image is in uint8 form.\"\"\"\n    image = transforms(image=image)['image']\n\n    return image","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def augment_func(image, label):\n    \"\"\"Function is formulated to convert our image\n    from a numpy array to a tf numpy array.\"\"\"\n    image = tf.numpy_function(func=transform,\n                              inp=[image, IMAGE_SIZE],\n                              Tout=tf.float32)\n    return image, label","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def read_tfrecord(dataset, labelled):\n    \"\"\"Function is fed with a tf.data.Dataset data type and returns\n    a singular image that has been augmented.\"\"\"\n    # Feature mapping for either labelled or unlabelled datasets\n    labelled_map = {\"image\": tf.io.FixedLenFeature([], tf.string),\n                    \"image_name\": tf.io.FixedLenFeature([], tf.string),\n                    \"target\": tf.io.FixedLenFeature([], tf.string)}\n    unlabelled_map = {\"image\": tf.io.FixedLenFeature([], tf.string),\n                     \"image_name\": tf.io.FixedLenFeature([], tf.string)}\n    tfrecord_format = (labelled_map if labelled else unlabelled_map)\n    \n    # Read singular images in a form of dictionary\n    \"\"\"parse_single_example function inputs a serialized dataset\n    and outputs singular examples as a form of a dict \n    {'image': image, 'target': target}\"\"\"\n    example = tf.io.parse_single_example(serialized=dataset,\n                                  features=tfrecord_format)\n\n    # Extracting the image from example dict\n    \"\"\"image is in a form of a Tensor.\"\"\"\n    image = example['image']\n\n    # Preprocess images\n    image = decode_func(image=image)\n\n    # Converting labels into int32 data type\n    \"\"\"If labelled is True, function returns label together\n    with the image.\"\"\"\n    if labelled:\n        label = example['target']\n        label = tf.cast(x=label, dtype=tf.string)\n        return image, label      \n        \n    return image","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def load_dataset(filenames, labelled=True, ordered=False):\n    \"\"\"Function is fed with filenames that are in a form of\n     lists of (.tfrec) shards (eg. PATH\\\\train00-1234.tfrec)\n     and outputs as a tf.data.Dataset data type.\"\"\"\n    # Initialize ignore_order variable\n    ignore_order = tf.data.Options()\n\n    if ordered==True:\n        # Reinstate order if ordered is set to True when calling dataset\n        ignore_order.experimental_deterministic = False\n\n    # Automatically interleaves reads from multiple files\n    dataset = tf.data.TFRecordDataset(filenames=filenames, \n                                    num_parallel_reads=AUTOTUNE)\n\n    # if ordered, ignore_order will be True\n    dataset = dataset.with_options(options=ignore_order)\n\n    # Apply read_tfrecord function to each element of the dataset\n    dataset = dataset.map(partial(read_tfrecord, labelled=labelled), \n                          num_parallel_calls=AUTOTUNE)\n\n    return dataset","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_train_dataset(filenames, labelled=True, ordered=False, \n                      augment=False, label_smoothing=True):\n    \"\"\"Function takes in the list of filenames that is\n    input into the load_dataset custom function, followed by\n    preprocessing the data at dataset level and outputs\n    as a tf.data.Dataset data type.\"\"\"\n    # Loads dataset from the training set (.tfrec) shard list\n    dataset = load_dataset(filenames=filenames, \n                           labelled=labelled, \n                           ordered=ordered)\n    \n    # Perform one hot encode for label smoothing\n    \"\"\"Set label_smoothing to True only when using \n    Categorical Crossentropy loss to simulate Sparse\n    Categorical Crossentropy since there is no label\n    smoothing parameter available.\"\"\"\n    if label_smoothing==True:\n        \"\"\"To obtain number of classes within the same function,\n        we take 50 batches for security to obtain all unique\n        labels\"\"\"\n        # Obtaining a subset of the dataset\n        subset = dataset.take(50).as_numpy_iterator()\n        # Extracting the labels\n        label_list = [label for (image, label) in subset]\n        # Obtaining the number of unique labels\n        classes = len(Counter(label_list).keys())\n        # Applying One Hot Encoding to the dataset labels\n        \"\"\"tf.data.Dataset comes in a form of paired list;\n        (image, label)\"\"\"\n        def one_hot(image, label):\n            return image, tf.one_hot(indices=label, \n                                     depth=classes,\n                                     dtype=tf.float32)    \n        dataset = dataset.map(map_func=one_hot,\n                              num_parallel_calls=AUTOTUNE)\n    \n    # Performing image preprocessing using tf.image functions\n    if augment==True:\n        dataset = dataset.map(map_func=augment_func, \n                              num_parallel_calls=AUTOTUNE)\n        \n      ##################################################\n      ## Insert other forms of advanced augmentations ##\n      ## eg. Attentive CutMix, CutMix, Cutout, MixUp, ##\n      ##          ShakeDrop, DropBlock, etc           ##\n      ##################################################\n#         dataset = dataset.map(map_func=cut_mix,\n#                               num_parallel_calls=AUTOTUNE)\n\n    # Repeats the dataset indefinitely during model training\n    dataset = dataset.repeat(count=None)\n\n    # Shuffles the dataset sampled by buffer_size\n    dataset = dataset.shuffle(buffer_size=2048, seed=SEED)\n\n    # Batching of the dataset as per BATCH_SIZE specified\n    dataset = dataset.batch(batch_size=BATCH_SIZE)\n\n    # Prefetch dataset in batches to reduce buffering time\n    dataset = dataset.prefetch(buffer_size=AUTOTUNE)\n\n    return dataset","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_valid_dataset(filenames, labelled=True, \n                      ordered=False, label_smoothing=True):\n    \"\"\"Function takes in the list of filenames that is\n    input into the load_dataset custom function, followed by\n    preprocessing the data at dataset level and outputs\n    as a tf.data.Dataset data type.\"\"\"\n    # Loads dataset from the validation set (.tfrec) shard list\n    dataset = load_dataset(filenames=filenames, \n                           labelled=labelled,\n                           ordered=ordered)\n\n    # Perform one hot encode for label smoothing\n    \"\"\"Set label_smoothing to True only when using \n    Categorical Crossentropy loss to simulate Sparse\n    Categorical Crossentropy since there is no label\n    smoothing parameter available.\"\"\"\n    if label_smoothing==True:\n        \"\"\"To obtain number of classes within the same function,\n        we take 50 batches for security to obtain all unique\n        labels\"\"\"\n        # Obtaining a subset of the dataset\n        subset = dataset.take(50).as_numpy_iterator()\n        # Extracting the labels\n        label_list = [label for (image, label) in subset]\n        # Obtaining the number of unique labels\n        classes = len(Counter(label_list).keys())\n        # Applying One Hot Encoding to the dataset labels\n        \"\"\"tf.data.Dataset comes in a form of paired list;\n        (image, label)\"\"\"\n        def one_hot(image, label):\n            return image, tf.one_hot(indices=label, \n                                     depth=classes,\n                                     dtype=tf.float32)          \n        dataset = dataset.map(map_func=one_hot,\n                              num_parallel_calls=AUTOTUNE)\n\n    # Batching of the dataset as per BATCH_SIZE specified\n    dataset = dataset.batch(batch_size=BATCH_SIZE)\n\n    # Saves subsequent preloaded batches into memory\n    dataset = dataset.cache()\n\n    # Prefetch dataset in batches to reduce buffering time\n    dataset = dataset.prefetch(buffer_size=AUTOTUNE)\n\n    return dataset","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_test_dataset(filenames, labelled=False, ordered=True):\n    \"\"\"Function takes in the list of filenames that is\n    input into the load_dataset custom function, followed by\n    preprocessing the data at dataset level and outputs\n    as a tf.data.Dataset data type.\"\"\"\n    # Loads dataset from the test set (.tfrec) shard list\n    dataset = load_dataset(filenames=filenames, \n                           labelled=labelled,\n                           ordered=ordered)\n\n    # Batching of the dataset as per BATCH_SIZE specified\n    dataset = dataset.batch(batch_size=BATCH_SIZE)\n\n    # Prefetch dataset in batches to reduce buffering time\n    dataset = dataset.prefetch(buffer_size=AUTOTUNE)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def count_data_items(filenames):\n    \"\"\"Obtaining total number of images in dataset from\n    the tfrecord shards.\"\"\"\n    n = [int(re.compile(r\"-([0-9]*)\\.\").search(filename).group(1))\n         for filename in filenames]\n    return np.sum(n)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Before Augmentation"},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n\nrow = 6; col = 5\n\nall_elements = get_train_dataset(FILENAMES,\n                                 augment=False,\n                                 label_smoothing=False)\n\nfor (image, label) in all_elements:\n    plt.figure(figsize=(15, int(15*row/col)))\n    for j in range(row*col):\n        plt.subplot(row, col, j+1)\n        plt.axis('off')\n        plt.imshow(image[j,])\n\n    plt.tight_layout()\n    plt.show()\n    break","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## After Augmentation\n\n*Note that the sample augmentations were exaggerated to accentuate the changes done to the images.*"},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n\nrow = 6; col = 5\n\nall_elements = get_train_dataset(FILENAMES,\n                                 augment=True,\n                                 label_smoothing=False)\n\nfor (image, label) in all_elements:\n    plt.figure(figsize=(15, int(15*row/col)))\n    for j in range(row*col):\n        plt.subplot(row, col, j+1)\n        plt.axis('off')\n        plt.imshow(image[j,])\n\n    plt.tight_layout()\n    plt.show()\n    break","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Do give me an upvote if this notebook helped. Thanks!!!"},{"metadata":{},"cell_type":"markdown","source":"### References\n\nHow to train a Keras model on TFRecord files: \nhttps://keras.io/examples/keras_recipes/tfrecord/\n\nUsing Albumentations with Tensorflow:\nhttps://albumentations.ai/docs/examples/tensorflow-example/"}],"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}