{"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"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":22990,"databundleVersionId":2048213,"sourceType":"competition"},{"sourceId":56448703,"sourceType":"kernelVersion"},{"sourceId":62323512,"sourceType":"kernelVersion"}],"dockerImageVersionId":30034,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Keras UFPN+EfficientNet training with TFRecords input\nIn this notebook we will train a U-Net architecture implemented in Keras/TensorFlow. We already created TFRecords of the dataset in this notebook:\n  * [HuBMAP TIF 2 TFRecords](https://www.kaggle.com/mistag/hubmap-image-2-tfrecords-256-512-1024)  \n  \n","metadata":{}},{"cell_type":"code","source":"!pip install segmentation_models -q","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nfrom pylab import *\nimport os\nos.environ['SM_FRAMEWORK'] = 'tf.keras'\nimport sys\nimport segmentation_models as sm\nimport tensorflow as tf\nfrom tensorflow.keras import backend as K\nfrom tensorflow.keras.utils import get_custom_objects\nfrom tensorflow.keras.models import model_from_json\nfrom tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint, ReduceLROnPlateau, CSVLogger, TensorBoard\nimport tensorflow_addons as tfa\nfrom functools import partial\nimport json\nfrom sklearn.model_selection import KFold\nimport matplotlib.pyplot as plt\n%matplotlib inline  ","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DPATH = '../input/fork-of-data-hubmap-image-2-tfrecords-256-512-10'\n# get parameters from the data creation notebook\n#with open(DPATH+'/dparams.json') as json_file:\n#    dparams = json.load(json_file)\n#dparams","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Dataset creation\nWe will use TFRecords with a resolution of 256x256 (downscaled from 1024x1024).","metadata":{}},{"cell_type":"code","source":"FILENAMES = tf.io.gfile.glob(DPATH+\"/*-256.tfrecord\")\nK_SPLITS = 5 # number of folds","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# augmentation\ndef data_augment(image, mask):\n    p_rotation = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    p_rotate = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    p_cutout = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    p_shear = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    p_crop = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    \n    # grayscale\n    if tf.random.uniform(()) > 0.9: \n        image = tf.image.rgb_to_grayscale(image)\n        image = tf.image.grayscale_to_rgb(image)\n\n    # x*90deg rotation\n    if p_rotate > .5:\n        image, mask = data_augment_rotate(image, mask)\n    \n    # flip\n    image, mask = data_augment_spatial(image, mask)\n    \n    # brightness/contrast\n    if tf.random.uniform(()) > 0.5:\n        if tf.random.uniform(()) > 0.5:\n            image = tf.image.random_contrast(image, 0.8, 1.2)\n        else:\n            image = tf.image.random_brightness(image, 0.1)\n    \n    # hue/saturation\n    if tf.random.uniform(()) > 0.5:\n        if tf.random.uniform(()) > 0.5:    \n            image = tf.image.random_saturation(image, 0.7, 1.3)\n        else:      \n            a = tf.random.uniform((), -0.1, 0.1)\n            image = tf.image.adjust_hue(image, a)\n            \n    # noise \n    if tf.random.uniform(()) > 0.5: \n        if tf.random.uniform(()) > 0.5:\n            gnoise = tf.random.normal(shape=tf.shape(image), mean=0.0, stddev=0.1, dtype=tf.float32)\n            image = tf.add(image, gnoise)\n        else:\n            image = tf.image.random_jpeg_quality(image, 40, 80, seed=None)\n    \n    # black pixels\n    if tf.random.uniform(()) > 0.75: \n        r = tf.random.uniform((IMG_SIZE, IMG_SIZE,3), minval=0, maxval=1, dtype=tf.dtypes.float32)      \n        r = (r > tf.random.uniform([], 0., .2, dtype=tf.float32))\n        image = tf.math.multiply(image, tf.cast(r, dtype=tf.float32))\n        \n    # cutout    \n    if tf.random.uniform(()) > 0.75:\n        image = data_augment_cutout(image, 10, 10, 48)\n            \n    image = tf.clip_by_value(image, 0.0, 1.0)\n\n    \n    return image, mask\n\ndef one_cut(image, min_size, max_size):\n    image = tf.squeeze(tfa.image.random_cutout(tf.raw_ops.Pack(values=[image]),\n                                               (tf.random.uniform((), minval=min_size, maxval=max_size, dtype=tf.dtypes.int32),\n                                                tf.random.uniform((), minval=min_size, maxval=max_size, dtype=tf.dtypes.int32)),\n                                               #tf.random.uniform((), minval=min_size, maxval=max_size, dtype=tf.dtypes.int32),\n                                               constant_values=tf.random.uniform(())))\n    return image\n\ndef data_augment_cutout(image, max_cuts, min_size, max_size):\n    cuts = tf.random.uniform((), minval=1, maxval=max_cuts, dtype=tf.dtypes.int32)\n    image = one_cut(image, min_size, max_size)\n    if cuts > 1:\n        image = one_cut(image, min_size, max_size)\n    if cuts > 2:\n        image = one_cut(image, min_size, max_size)\n    if cuts > 3:\n        image = one_cut(image, min_size, max_size)\n    if cuts > 4:\n        image = one_cut(image, min_size, max_size)\n    if cuts > 5:\n        image = one_cut(image, min_size, max_size)\n    if cuts > 6:\n        image = one_cut(image, min_size, max_size)\n    if cuts > 7:\n        image = one_cut(image, min_size, max_size)\n    if cuts > 8:\n        image = one_cut(image, min_size, max_size)\n    if cuts > 9:\n        image = one_cut(image, min_size, max_size)\n    return image\n\ndef data_augment_spatial(image, mask):\n    if tf.random.uniform(()) > 0.5:\n        image = tf.image.flip_left_right(image)\n        mask = tf.image.flip_left_right(mask)\n    if tf.random.uniform(()) > 0.5:\n        image = tf.image.flip_up_down(image)\n        mask = tf.image.flip_up_down(mask)\n\n    return image, mask\n\ndef data_augment_rotate(image, mask):\n    p_rotate = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    \n    if p_rotate > .66:\n        image = tf.image.rot90(image, k=3) # rotate 270º\n        mask = tf.image.rot90(mask, k=3) # rotate 270º\n    elif p_rotate > .33:\n        image = tf.image.rot90(image, k=2) # rotate 180º\n        mask = tf.image.rot90(mask, k=2) # rotate 180º\n    else:\n        image = tf.image.rot90(image, k=1) # rotate 90º\n        mask = tf.image.rot90(mask, k=1) # rotate 90º\n\n    return image, mask\n\ndef data_augment_crop(image, mask):\n    p_crop = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    crop_size = tf.random.uniform([], int(IMG_SIZE*.8), IMG_SIZE, dtype=tf.int32)\n    \n    if p_crop > .5:\n        ox = tf.random.uniform([], 0, IMG_SIZE-crop_size, dtype=tf.int32)\n        oy = tf.random.uniform([], 0, IMG_SIZE-crop_size, dtype=tf.int32)\n        image = tf.image.crop_to_bounding_box(image, oy, ox, crop_size, crop_size)\n        mask = tf.image.crop_to_bounding_box(mask, oy, ox, crop_size, crop_size)\n    else:\n        if p_crop > .4:\n            image = tf.image.central_crop(image, central_fraction=.7)\n            mask = tf.image.central_crop(mask, central_fraction=.7)\n        elif p_crop > .2:\n            image = tf.image.central_crop(image, central_fraction=.8)\n            mask = tf.image.central_crop(mask, central_fraction=.8)\n        else:\n            image = tf.image.central_crop(image, central_fraction=.9)\n            mask = tf.image.central_crop(mask, central_fraction=.9)\n    \n    image = tf.image.resize(image, size=[IMG_SIZE, IMG_SIZE])\n    mask = tf.image.resize(mask, size=[IMG_SIZE, IMG_SIZE], method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)\n\n    return image, mask","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"IMG_SIZE = 256\nIMAGE_SIZE = [IMG_SIZE, IMG_SIZE]\nAUTOTUNE = tf.data.experimental.AUTOTUNE\nBATCH_SIZE = 16\n\n# hyperparameters saved for later use during inference\nhparams = {\n    \"IMG_SIZE\": IMG_SIZE,\n    \"SCALE_FACTOR\": 1024//IMG_SIZE,  # TFRecords dataset uses tiles of 1024x1024\n    \"BATCH_SIZE\": BATCH_SIZE,\n    \"K_SPLITS\": K_SPLITS}\nwith open(\"hparams.json\", \"w\") as json_file:\n    json_file.write(json.dumps(hparams, indent = 4))\n\n# decode image or mask\ndef decode_image(image, isjpeg=True):\n    if isjpeg:\n        ch = 3\n        image = tf.image.decode_jpeg(image, channels=ch)\n    else:\n        ch = 1\n        image = tf.image.decode_png(image, channels=ch)\n        image = tf.expand_dims(image, -1)\n    image = tf.cast(image, tf.float32)\n    image = image /255.\n    image = tf.reshape(image, [*IMAGE_SIZE, ch])\n    return image\n\n# read a single record \ndef read_tfrecord(example):\n    tfrecord_format = ( # only extract features we are interested in\n        {\n            \"image/encoded\": tf.io.FixedLenFeature([], tf.string),\n            \"mask/encoded\": tf.io.FixedLenFeature([], tf.string),\n        }\n    )\n    example = tf.io.parse_single_example(example, tfrecord_format)\n    image = decode_image(example[\"image/encoded\"], True) # jpeg format\n    mask = decode_image(example[\"mask/encoded\"], False) # png format\n    return image, mask\n\ndef load_dataset(filenames, IsTrain=True):\n    ignore_order = tf.data.Options()\n    ignore_order.experimental_deterministic = False  # disable order, increase speed\n    # automatically interleaves reads from multiple files\n    dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads = AUTOTUNE)\n    # uses data as soon as it streams in, rather than in its original order\n    dataset = dataset.with_options(ignore_order)  \n    dataset = dataset.map(partial(read_tfrecord), num_parallel_calls=AUTOTUNE)\n    if IsTrain: # augmentation        \n        dataset = dataset.map(partial(data_augment), num_parallel_calls=AUTOTUNE)\n    # returns a dataset of (image, mask) pairs \n    return dataset\n\ndef get_dataset(filenames, IsTrain=True):\n    dataset = load_dataset(filenames, IsTrain)\n    dataset = dataset.shuffle(BATCH_SIZE*256)\n    dataset = dataset.prefetch(buffer_size=AUTOTUNE)\n    dataset = dataset.batch(BATCH_SIZE)\n    dataset = dataset.repeat()\n    return dataset","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Plot a few images from the dataset to check that everything is OK (including augmentation):","metadata":{}},{"cell_type":"code","source":"train_dataset = get_dataset(FILENAMES, True)\n\nimage_batch, mask_batch = next(iter(train_dataset))\n\ndef show_batch(image_batch, mask_batch):\n    plt.figure(figsize=(16, 16))\n    for n in range(min(BATCH_SIZE,16)):\n        ax = plt.subplot(4, 4, n + 1)\n        plt.imshow(image_batch[n])\n        plt.imshow(np.squeeze(mask_batch[n]), alpha=0.25)\n        plt.axis(\"off\")\n\nshow_batch(image_batch.numpy(), mask_batch.numpy())","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"There is a Pandas pickle file accompanying the TFRecords files containing the number of images per TFRecord.","metadata":{}},{"cell_type":"code","source":"# calculate number of images in train/val sets for each fold\nkf = KFold(n_splits=K_SPLITS)\ndf = pd.read_pickle(DPATH+'/record_stats.pkl')\ntcnt, vcnt = np.zeros(K_SPLITS, dtype=int), np.zeros(K_SPLITS, dtype=int)\nidx = 0\nfor train_index, test_index in kf.split(FILENAMES):\n    for i in train_index:\n        fname = FILENAMES[i].split('/')[-1].split('\\\\')[-1]\n        tcnt[idx] += df[df.File == fname].ImgCount.iloc[0]\n    for i in test_index:\n        fname = FILENAMES[i].split('/')[-1].split('\\\\')[-1]\n        vcnt[idx] += df[df.File == fname].ImgCount.iloc[0]\n    idx += 1\n\nprint('Train images: {}, Validation images: {}'.format(tcnt, vcnt))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Let's have a look at some statistics from the [quick reference](https://www.kaggle.com/mistag/hubmap-quick-reference):","metadata":{}},{"cell_type":"code","source":"dff = pd.read_pickle('../input/hubmap-quick-reference/image_stats.pkl')\ntot = dff[dff.dataset == 'train'].glomeruli.sum()\ndff.head(15)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Loss Functions\nThere are several loss functions to choose from, a few ones are defined below. The Focal Tversky loss is known to perform well on many segmentation tasks. We could also use the Dice coefficient loss or the Tversky loss (experiment to find the best one). ","metadata":{}},{"cell_type":"code","source":"# metrics and loss functions\nsmooth = 1.\n\ndef dice_coef(y_true, y_pred):\n    y_true_f = K.flatten(y_true)\n    y_pred_f = K.flatten(y_pred)\n    intersection = K.sum(y_true_f * y_pred_f)\n    return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)\n\n\ndef dice_coef_loss(y_true, y_pred):\n    return 1.-dice_coef(y_true, y_pred)\n\n\ndef iou(y_true, y_pred):\n    def f(y_true, y_pred):\n        intersection = (y_true * y_pred).sum()\n        union = y_true.sum() + y_pred.sum() - intersection\n        x = (intersection + 1e-15) / (union + 1e-15)\n        x = x.astype(np.float32)\n        return x\n    return tf.numpy_function(f, [y_true, y_pred], tf.float32)\n\ndef tversky(y_true, y_pred, smooth=1, alpha=0.7):\n    y_true_pos = K.flatten(y_true)\n    y_pred_pos = K.flatten(y_pred)\n    true_pos = K.sum(y_true_pos * y_pred_pos)\n    false_neg = K.sum(y_true_pos * (1 - y_pred_pos))\n    false_pos = K.sum((1 - y_true_pos) * y_pred_pos)\n    return (true_pos + smooth) / (true_pos + alpha * false_neg + (1 - alpha) * false_pos + smooth)\n\ndef tversky_loss(y_true, y_pred):\n    return 1 - tversky(y_true, y_pred)\n\ndef focal_tversky_loss(y_true, y_pred, gamma=0.75):\n    tv = tversky(y_true, y_pred)\n    return K.pow((1 - tv), gamma)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Compile & train model\nNote that we save the built models as a .json file for use during inference, one for each fold.","metadata":{}},{"cell_type":"code","source":"MNAME = 'FPN-model43e'\n\ndef get_callbacks(idx):\n    mc = ModelCheckpoint(MNAME+\"-{}.h5\".format(idx), save_best_only=True)\n    rp = ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=5, min_lr=0.00001)\n    #cl = CSVLogger(\"train_log-{}.csv\".format(idx))\n    #tb = TensorBoard()\n    es = EarlyStopping(monitor='val_loss', patience=6, restore_best_weights=False)\n    return [mc, rp, es]","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#%%capture\nhistory = []\nloss, vloss, dice, vdice = np.zeros(K_SPLITS),np.zeros(K_SPLITS),np.zeros(K_SPLITS),np.zeros(K_SPLITS)\nidx = 0\nlr = 5e-4\nselected_folds = [4] # select a few folds only to reduce execution time\nfor train_index, test_index in kf.split(FILENAMES):\n    if idx in selected_folds: \n        train_dataset = get_dataset([FILENAMES[i] for i in train_index], True)\n        valid_dataset = get_dataset([FILENAMES[i] for i in test_index], False)\n        model = sm.FPN('efficientnetb4', classes=1, encoder_weights='imagenet', activation = 'sigmoid')\n        with open(MNAME+\"-{}.json\".format(idx), \"w\") as json_file:\n            json_file.write(model.to_json())    \n        opt = tf.keras.optimizers.Adam(lr)\n        metrics = [\"acc\", iou, dice_coef, tversky]\n        model.compile(loss=dice_coef_loss, optimizer=opt, metrics=metrics)\n        callbacks = get_callbacks(idx)\n        hist = model.fit(train_dataset, \n                            validation_data=valid_dataset,\n                            epochs=25,\n                            steps_per_epoch=1+tcnt[idx]//BATCH_SIZE,\n                            validation_steps=1+vcnt[idx]//BATCH_SIZE,\n                            callbacks=callbacks)\n        loss[idx] = hist.history['loss'][-1]\n        vloss[idx] = min(hist.history['val_loss'])\n        dice[idx] = hist.history['dice_coef'][-1]\n        vdice[idx] = max(hist.history['val_dice_coef'])\n        history.append(hist)\n    idx += 1","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pd.DataFrame({'loss': loss, 'min val.loss': vloss, 'dice coef.': dice, 'max val. dice coef.': vdice})","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Inspect learning curves\nBelow we plot the loss for both training and validation along with the Dice coefficients for the last fold. It is important to keep an eye on these curves to verify that our model and training are setup correctly.","metadata":{}},{"cell_type":"code","source":"colors = ['tab:blue', 'tab:orange', 'tab:green', 'tab:red', 'tab:purple']\nplt.figure(figsize=(16,10))\nfor i in range(len(selected_folds)):\n    plt.plot(history[i].history['loss'], linestyle='-', color=colors[i], label='Train loss fold #{}'.format(selected_folds[i]))\nfor i in range(len(selected_folds)):\n    plt.plot(history[i].history['val_loss'], linestyle='--', color=colors[i], label='Validation loss fold #{}'.format(selected_folds[i]))\nplt.title('Model Loss')\nplt.ylabel('Loss')\nplt.xlabel('Epoch')\nplt.legend()\nplt.show();","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"colors = ['tab:blue', 'tab:orange', 'tab:green', 'tab:red', 'tab:purple']\nplt.figure(figsize=(16,10))\nfor i in range(len(selected_folds)):\n    plt.plot(history[i].history['dice_coef'], linestyle='-', color=colors[i], label='Train dice coef. fold #{}'.format(selected_folds[i]))\nfor i in range(len(selected_folds)):\n    plt.plot(history[i].history['val_dice_coef'], linestyle='--', color=colors[i], label='Validation dice coef. fold #{}'.format(selected_folds[i]))\nplt.title('Model Dice Coef.')\nplt.ylabel('Dice coef.')\nplt.xlabel('Epoch')\nplt.legend()\nplt.show();","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Looks pretty good! We are now ready for the final step - [making predictions with the saved model](https://www.kaggle.com/mistag/inference-hubmap-fpn-single-model-ii). ","metadata":{}},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}