{"cells":[{"metadata":{},"cell_type":"markdown","source":"This notebook shows an implementation of\nthe [Fmix](https://arxiv.org/pdf/2002.12047.pdf) data augmentation by Tensorflow.\nFmix is a kind of Mixed Sample Data Augmentation.\nIt uses binary masks obtained by applying a threshold to low frequency images sampled from Fourier space.\n\nI referred the followings:\n* [Understanding and Enhancing Mixed Sample Data Augmentation](https://arxiv.org/pdf/2002.12047.pdf) -- Original paper.\n* https://github.com/ecs-vlc/FMix -- Original code.\n* [Getting started with 100+ flowers on TPU](https://www.kaggle.com/mgornergoogle/getting-started-with-100-flowers-on-tpu) -- based on this notebook.\n* https://github.com/numpy/numpy/blob/master/numpy/fft/helper.py -- fftfreq() implementation."},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import math, re, os\nimport tensorflow as tf\nimport numpy as np\nfrom matplotlib import pyplot as plt\nfrom PIL import Image\nfrom kaggle_datasets import KaggleDatasets\nprint(\"Tensorflow version \" + tf.__version__)\nAUTO = tf.data.experimental.AUTOTUNE","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# TPU or GPU detection"},{"metadata":{"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":{},"cell_type":"markdown","source":"# Competition data access"},{"metadata":{"trusted":true},"cell_type":"code","source":"GCS_DS_PATH = KaggleDatasets().get_gcs_path() # 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] # At this size, a GPU will run out of memory. Use the TPU.\n                        # For GPU training, please select 224 x 224 px image size.\nBATCH_SIZE = 16 * strategy.num_replicas_in_sync\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')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Datasets"},{"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    \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    # disregarding 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) # uses 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)\n    # returns a dataset of (image, label) pairs if labeled=True or (image, id) pairs if labeled=False\n    return dataset","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Fmix"},{"metadata":{},"cell_type":"markdown","source":"The steps of Fmix are as follows:\n1. We  first  sample  a  random  complex  tensor  for which both the real and imaginary part are independent and Gaussian.\n2. We then scale each component according to its frequency via the parameter δ such that higher values of δ correspond to increased decay of high frequency information.\n3. Next, we perform an inverse Fourier transform on the complex tensor and take the real part to obtain a grey-scale image.\n4. Finally, we set the top proportion of the image to have value ‘1’ and the rest to have value ‘0’ to obtain our binary mask, then mix two images by using this mask."},{"metadata":{"trusted":true},"cell_type":"code","source":"# I could not find an equivalent in Tensorflow, so I made this.\n# Original: https://github.com/numpy/numpy/blob/master/numpy/fft/helper.py\n\ndef fftfreq(n, d=1.0):\n    val = 1.0 / (n * d)\n    N = (n - 1) // 2 + 1\n    p1 = tf.range(0, N, dtype=tf.float32)\n    p2 = tf.range(-(n // 2), 0, dtype=tf.float32)\n    results = tf.concat([p1, p2], 0)\n    return results * val","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# https://github.com/ecs-vlc/FMix/blob/master/fmix.py\n\n# Parameters are 'h' and 'w' for simplicity.\ndef fftfreqnd(h, w):\n    \"\"\" Get bin values for discrete fourier transform of size (h, w)\n    \"\"\"\n    # In the original implementation, '[: w // 2 + 2]' or '[: w // 2 + 1]' is\n    # applied to fx.  However, I don't do this here, because tf.signal.ifft2d\n    # returns the same shape as the input.  tf.signal.ifft2d does not accept\n    # shape, as in np.fft.irfftn used in the original code.  I think that\n    # a tensor of width by height is necessary here for tf.signal.ifft2d.\n    fx = fftfreq(w)  # [: w // 2 + 2] or [ : w // 2 + 1]\n    fy = fftfreq(h)\n    \n    fx_square = fx * fx\n    fy_square = fy * fy\n    return tf.math.sqrt(fx_square[ tf.newaxis, : ] + fy_square[ : , tf.newaxis ])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Here is the result of 'fftfreqnd(512, 512)'.\nFour corners correspond to low frequency,\nand center is high frequency.\nNow, the values at corners are small\nand those at center are big.\nThe result here will be used as a divisor,\nso finally values for low frequency become big and\nthose for high frequency are small."},{"metadata":{"trusted":true},"cell_type":"code","source":"ffreqnd = fftfreqnd(512, 512)\n\nplt.pcolor(ffreqnd)\nplt.colorbar()\nplt.title('fftfreqnd(512, 512) result')\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"One difficult point in Tensorflow program is\nto find the shape of a tensor.\nTo clarify this point, I use suffix.\nThe characters used in suffix are:\n* b -- batch\n* h, w, c -- image height, width, and color.\n* p -- pixels in an image (height x width)\n* t -- total pixels in a batch (batch x pixels in a image)\n* 1, 2 -- constant\n\nFor example:\n* ...._bhwc means rank 4 tensor and each dimensions are for batch, height, width, and color.\n* tf.expand_dims(..._bhw, -1) expands one dimention at last, so the result is ..._bhw1."},{"metadata":{"trusted":true},"cell_type":"code","source":"IMAGE_MAX_W_H = max(IMAGE_SIZE[0], IMAGE_SIZE[1])\n\ndef get_spectrum(data_count, freqs, decay_power):\n    # Make a tensor to scale frequencies, low frequencies are bigger\n    # and high frequencies are smaller.\n    # Make freqs greater than 0 to avoid division by 0.\n    lowest_freq = tf.constant(1. / IMAGE_MAX_W_H)\n    freqs_gt_zero = tf.math.maximum(freqs, lowest_freq)\n    scale_hw = 1.0 / tf.math.pow(freqs_gt_zero, decay_power)\n\n    # Generate random Gaussian distribution numbers of data_count x height x width x 2.\n    # 2 in the last dimension is for real and imaginary part of a complex number.\n    # In the original program, the first dimention is used for channels.\n    # In this program, it is used for data in a batch.\n    param_size = [data_count] + list(freqs.shape) + [2]\n    param_bhw2 = tf.random.normal(param_size)\n\n    # Make a spectrum by multiplying scale and param.  For scale,\n    # expand first and last dimension for batch and real/imaginary part.\n    scale_1hw1 = tf.expand_dims(scale_hw, -1)[ tf.newaxis, : ]\n    spectrum_bhw2 = scale_1hw1 * param_bhw2\n    return spectrum_bhw2","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def make_low_freq_images(data_count, decay):\n    # Make a mask image by inverse Fourier transform of a spectrum,\n    # which is generated by get_spectrum().\n    freqs = fftfreqnd(*IMAGE_SIZE)\n    spectrum_bhw2 = get_spectrum(data_count, freqs, decay)\n    spectrum_re_bhw = spectrum_bhw2[:, :, :, 0]\n    spectrum_im_bhw = spectrum_bhw2[:, :, :, 1]\n    spectrum_comp_bhw = tf.complex(spectrum_re_bhw, spectrum_im_bhw)\n    mask_bhw = tf.math.real(tf.signal.ifft2d(spectrum_comp_bhw))\n\n    # Scale the mask values from 0 to 1.\n    mask_min_b = tf.reduce_min(mask_bhw, axis=(-2, -1))\n    mask_min_b11 = mask_min_b[ :, tf.newaxis, tf.newaxis]\n    mask_shift_to_0_bhw = mask_bhw - mask_min_b11\n    mask_max_b = tf.reduce_max(mask_shift_to_0_bhw, axis=(-2, -1))\n    mask_max_b11 = mask_max_b[ :, tf.newaxis, tf.newaxis]\n    mask_scaled_bhw = mask_shift_to_0_bhw / mask_max_b11\n    return mask_scaled_bhw","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Helper function to draw a tensor as images.\n# Expected tensor shape is [batch, height, width, and optinal color]\n# Expected values are from 0 to 1.\ndef plot_as_images(tf_0_1, title, figsize=(12, 2), rows=1):\n    plt.figure(figsize=figsize)\n    # If the shape is 'bhw1', then make it 'bhw'.\n    if tf_0_1.ndim == 4 and tf_0_1.shape[-1] == 1:\n        tf_0_1 = tf.reshape(tf_0_1, tf_0_1.shape[:3])\n    np_0_255 = (tf_0_1.numpy() * 255).astype(np.uint8)\n    image_count = np_0_255.shape[0]\n    cols = (image_count + rows - 1) // rows\n    for i in range(image_count):\n        plt.subplot(rows, cols, i + 1)\n        im_pil = Image.fromarray(np_0_255[i])\n        plt.imshow(im_pil, cmap='gray')\n        plt.title(title + str(i + 1))\n        plt.axis(\"off\")\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Here are some samples of low frequency images."},{"metadata":{"trusted":true},"cell_type":"code","source":"lfimages = make_low_freq_images(5, 3.0)\nplot_as_images(lfimages, \"Low Freq Image \")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"IMAGE_PIXEL_COUNT = IMAGE_SIZE[0] * IMAGE_SIZE[1]\n\ndef make_binary_masks(data_count, low_freq_images_bhw, mix_ratios_b):\n    # The goal is \"top proportion of the image to have value ‘1’ and the rest to have value ‘0’\".\n    # To make this I use tf.scatter_nd().  For tf.scatter_nd(), indices and values\n    # are necessary.\n    \n    # For each image, get indices of an image whose order is sorted from top to bottom.\n    # These are used for row indices.  To combine with column indices, they are reshaped to 1D.\n    low_freq_images_bp = tf.reshape(low_freq_images_bhw, [data_count, -1])\n    row_indices_bp = tf.argsort(low_freq_images_bp, axis=-1, direction='DESCENDING', stable=True)\n    row_indices_t = tf.reshape(row_indices_bp, [-1])\n    \n    # Make column indices, col_indices_t looks like\n    # '[ 0 ... 0 1 ... 1 ..... data_count-1 ... data_count-1]'\n    col_indices_b = tf.range(data_count, dtype=tf.int32)\n    col_indices_t = tf.repeat(col_indices_b, IMAGE_PIXEL_COUNT, axis=-1)\n    \n    # Combine column and row indices for scatter_nd.\n    scatter_indices_2t = tf.stack([col_indices_t, row_indices_t])\n    scatter_indices_t2 = tf.transpose(scatter_indices_2t)\n\n    # Make a tensor which looks like:\n    # [[ 0.0 ... 1.0 ]   \\  <-- tf.linspace(0.0, 1.0, IMAGE_PIXEL_COUNT)\n    #   ...               | data_count\n    #  [ 0.0 ... 1.0 ]]  /\n    linspace_0_1_p = tf.linspace(0.0, 1.0, IMAGE_PIXEL_COUNT)\n    linspace_0_1_1p = linspace_0_1_p[ tf.newaxis, : ]\n    linspace_0_1_bp = tf.repeat(linspace_0_1_1p, data_count, axis=0)\n    \n    # Make mix_ratio of the top elements in each data '1' and the rest '0'\n    # This looks like:\n    # [[ 1.0 1.0 ... 0.0 ]   \\    <-- top mix_ratios_b[0] elements are 1.0\n    #   ...                   | data_count\n    #  [ 1.0 1.0 ... 0.0 ]]  /    <-- top mix_ratios_b[data_count - 1] elements are 1.0\n    mix_ratios_b1 = mix_ratios_b[ :, tf.newaxis]\n    scatter_updates_bp = tf.where(linspace_0_1_bp <= mix_ratios_b1, 1.0, 0.0)\n    scatter_updates_t = tf.reshape(scatter_updates_bp, [-1])\n    \n    # Make binary masks by using tf.scatter_nd(), then reshape.\n    bin_masks_bp = tf.scatter_nd(scatter_indices_t2, scatter_updates_t, [data_count, IMAGE_PIXEL_COUNT])\n    bin_masks_bhw1 = tf.reshape(bin_masks_bp, [data_count, *IMAGE_SIZE, 1])\n    return bin_masks_bhw1","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Here are some binary mask samples."},{"metadata":{"trusted":true},"cell_type":"code","source":"lfimages = make_low_freq_images(5, 3.0)\nbin_masks = make_binary_masks(5, lfimages, tf.constant([ 0.1, 0.3, 0.5, 0.7, 0.9 ]))\nplot_as_images(bin_masks, \"Binary mask \")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def do_mix(mix_ratios, orig_batch, mixing_batch):\n    # Mix original and mixing batch data with the specified ratios.\n    mixed_batch = (mix_ratios * orig_batch) + ((1.0 - mix_ratios) * mixing_batch)\n    return mixed_batch","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The followings are samples for low frequency gray-scale images,\nbinary masks, images to be mixed, and Fmix images."},{"metadata":{"trusted":true},"cell_type":"code","source":"lfimages = make_low_freq_images(5, 3.0)\nbin_masks = make_binary_masks(5, lfimages, tf.constant([ 0.5 ] * 5))\ndataset = iter(load_dataset(TRAINING_FILENAMES, labeled=True).batch(5))\nimages_1, labels_1 = next(dataset)\nimages_2, labels_2 = next(dataset)\nmixed_images = do_mix(bin_masks, images_1, images_2)\n\nplot_as_images(lfimages, \"Gray-scale \")\nplot_as_images(bin_masks, \"Binary mask \")\nplot_as_images(images_1, \"Image 1-\")\nplot_as_images(images_2, \"Image 2-\")\nplot_as_images(mixed_images, \"Fmix \")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import tensorflow_probability as tfp\ntfd = tfp.distributions\n\n# Hyper parameters and a beta distribution to generate mix ratios.\nFMIX_ALPHA = 1.0\nFMIX_DECAY = 3.0\nfmix_beta_dist = tfd.Beta(FMIX_ALPHA, FMIX_ALPHA)\n\ndef fmix(orig_images):\n    # orig_images.shape[0] is None when making computational graph.\n    # So, this should be a tensor.\n    data_count = tf.shape(orig_images)[0]\n    \n    # Randomly select mixing images.\n    data_indices = tf.range(data_count, dtype=tf.int32)\n    mixing_indices = tf.random.shuffle(data_indices)\n    mixing_images = tf.gather(orig_images, mixing_indices)\n\n    # Generate mix ratios by beta distribution.\n    mix_ratios = fmix_beta_dist.sample([data_count])\n\n    # Generate binary masks, then mix images.\n    low_freq_images = make_low_freq_images(data_count, FMIX_DECAY)\n    bin_masks = make_binary_masks(data_count, low_freq_images, mix_ratios)\n    mixed_images = do_mix(bin_masks, orig_images, mixing_images)\n    return mixed_images","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Here are Fmix samples."},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset = iter(load_dataset(TRAINING_FILENAMES, labeled=True).batch(15))\norig_images, orig_labels = next(dataset)\nmixed_images = fmix(orig_images)\n\nplot_as_images(mixed_images, \"Fmix \", figsize=(12, 7), rows=3)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Thanks for reading!\nIf you have any questions or find mistakes, please let me know."},{"metadata":{},"cell_type":"markdown","source":"# History\n* 2020/03/29 -- First version\n* 2020/04/25 -- Fixed an issue for \"data_count\" in \"fmix()\". It should be a tensor. Otherwise it might generate an error if data count was not the predefined constant \"BATCH_SIZE\". "}],"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}