{"cells":[{"metadata":{},"cell_type":"markdown","source":"## Keras U-Net+MobileNetV2 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 JPG+TFRecords](https://www.kaggle.com/mistag/data-hubmap-tif-2-jpg-tfrecords-128-256-512-1024/edit/run/47859494)  \n  \nThen trained a U-Net only in [this notebook](https://www.kaggle.com/mistag/inference-hubmap-u-net-512x512?scriptVersionId=48280335).  \n\nIn this notebook we will use MobileNetV2 as the encoder part of the U-Net model."},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"from tensorflow.keras.models import Model\nfrom tensorflow.keras.layers import *\nfrom tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint, ReduceLROnPlateau, CSVLogger, TensorBoard\nfrom tensorflow.keras import backend as K\nfrom tensorflow.keras.applications import MobileNetV2\nimport tensorflow as tf\nfrom functools import partial\nimport numpy as np\nimport pandas as pd\nimport json\nfrom sklearn.model_selection import KFold\nimport matplotlib.pyplot as plt\n%matplotlib inline  ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# get parameters from the data creation notebook\nwith open('../input/data-hubmap-tif-2-jpg-tfrecords-128-256-512-1024/dparams.json') as json_file:\n    dparams = json.load(json_file)\ndparams","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Dataset creation\nWe will use TFRecords with a resolution of 256x256 (downscaled from 1024x1024)."},{"metadata":{"trusted":true},"cell_type":"code","source":"FILENAMES = tf.io.gfile.glob(\"/kaggle/input/data-hubmap-tif-2-jpg-tfrecords-128-256-512-1024/Tissue-256-*.tfrecord\")\nK_SPLITS = 5 # number of folds","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"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\n# read a single record and do augmentation\ndef read_tfrecord_tr(example):\n    image, mask = read_tfrecord(example)\n    # basic augmentation  (expand as desired)\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    if tf.random.uniform(()) > 0.75:\n        if tf.random.uniform(()) > 0.5:\n            k = 1\n        else:\n            k = 3\n        image = tf.image.rot90(image, k)\n        mask = tf.image.rot90(mask, k)\n    if tf.random.uniform(()) > 0.75: # random contrast/brightness\n        if tf.random.uniform(()) > 0.5:\n            a = tf.random.uniform((), 0.7, 1.3)\n            image = tf.image.adjust_contrast(image, a)\n        else:\n            a = tf.random.uniform((), 0., 0.5)\n            image = tf.image.adjust_brightness(image, a)\n    if tf.random.uniform(()) > 0.7: # add noise\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    if tf.random.uniform(()) > 0.7: # change hue\n        a = tf.random.uniform((), -0.2, 0.2)\n        image = tf.image.adjust_hue(image, a)  \n        \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    dataset = tf.data.TFRecordDataset(\n        filenames\n    )  # automatically interleaves reads from multiple files\n    dataset = dataset.with_options(\n        ignore_order\n    )  # uses data as soon as it streams in, rather than in its original order\n    if IsTrain: # augmentation\n        dataset = dataset.map(\n            partial(read_tfrecord_tr), num_parallel_calls=AUTOTUNE\n        )\n    else: # no augmentation\n        dataset = dataset.map(\n            partial(read_tfrecord), num_parallel_calls=AUTOTUNE\n        )\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.prefetch(buffer_size=AUTOTUNE)\n    dataset = dataset.batch(BATCH_SIZE)\n    dataset = dataset.repeat()\n    return dataset","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Plot a few images from the dataset to check that everything is OK (including augmentation):"},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset = get_dataset(FILENAMES, True)\n\nimage_batch, mask_batch = next(iter(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())","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"There is a Pandas pickle file accompanying the TFRecords files containing the number of images per TFRecord."},{"metadata":{"trusted":true},"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('../input/data-hubmap-tif-2-jpg-tfrecords-128-256-512-1024/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))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Build U-Net model\nThe U-Net model is really simple to build, and easy to modify as well. Code snippet from [UNET Segmentation with Pretrained MobileNetV2 as Encoder](https://idiotdeveloper.com/unet-segmentation-with-pretrained-mobilenetv2-as-encoder/)."},{"metadata":{"trusted":true},"cell_type":"code","source":"def unet_vgg_model():\n    inputs = Input((IMG_SIZE, IMG_SIZE, 3), name=\"input_image\")\n    \n    encoder = MobileNetV2(input_tensor=inputs, weights=\"imagenet\", include_top=False, alpha=0.35)\n    skip_connection_names = [\"input_image\", \"block_1_expand_relu\", \"block_3_expand_relu\", \"block_6_expand_relu\"]\n    encoder_output = encoder.get_layer(\"block_13_expand_relu\").output\n    \n    f = [16, 32, 64, 128]\n    x = encoder_output\n    for i in range(1, len(skip_connection_names)+1, 1):\n        x_skip = encoder.get_layer(skip_connection_names[-i]).output\n        x = UpSampling2D((2, 2))(x)\n        x = Concatenate()([x, x_skip])\n        \n        x = Conv2D(f[-i], (3, 3), padding=\"same\")(x)\n        x = BatchNormalization()(x)\n        x = Activation(\"relu\")(x)\n        \n        x = Conv2D(f[-i], (3, 3), padding=\"same\")(x)\n        x = BatchNormalization()(x)\n        x = Activation(\"relu\")(x)\n        \n    x = Conv2D(1, (1, 1), padding=\"same\")(x)\n    x = Activation(\"sigmoid\")(x)\n    \n    model = Model(inputs, x)\n    return model","execution_count":null,"outputs":[]},{"metadata":{},"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":{"trusted":true},"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)","execution_count":null,"outputs":[]},{"metadata":{},"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":{"trusted":true},"cell_type":"code","source":"def get_callbacks(idx):\n    mc = ModelCheckpoint(\"model{}.h5\".format(idx), save_best_only=True)\n    rp = ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=8, min_lr=0.00001)\n    cl = CSVLogger(\"train_log{}.csv\".format(idx))\n    tb = TensorBoard()\n    #es = EarlyStopping(monitor='val_loss', patience=10, restore_best_weights=False)\n    return [mc, rp, cl, tb]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%capture\nloss, vloss, dice, vdice = np.zeros(K_SPLITS),np.zeros(K_SPLITS),np.zeros(K_SPLITS),np.zeros(K_SPLITS)\nidx = 0\nfor train_index, test_index in kf.split(FILENAMES):\n    lr = 5e-4\n    model = unet_vgg_model()\n    with open(\"model{}.json\".format(idx), \"w\") as json_file:\n        json_file.write(model.to_json())\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    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    history = model.fit(train_dataset, \n                        validation_data=valid_dataset,\n                        epochs=60,\n                        steps_per_epoch=tcnt[idx]//BATCH_SIZE,\n                        validation_steps=vcnt[idx]//BATCH_SIZE,\n                        callbacks=callbacks)\n    model.save('./model{}'.format(idx))\n    loss[idx] = history.history['loss'][-1]\n    vloss[idx] = history.history['val_loss'][-1]\n    dice[idx] = history.history['dice_coef'][-1]\n    vdice[idx] = history.history['val_dice_coef'][-1]   \n    idx += 1","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dfh = pd.DataFrame({'loss': loss, 'validation loss': vloss, 'dice coef.': dice, 'validation dice coef.': vdice})\ndfh","execution_count":null,"outputs":[]},{"metadata":{},"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":{"trusted":true},"cell_type":"code","source":"plt.figure(figsize=(16, 8))\nx = np.arange(1,len(history.history['loss'])+1)\nplt.plot(x, history.history['loss'], label='Train loss')\nplt.plot(x, history.history['val_loss'], label='Validation loss')\nplt.plot(x, history.history['dice_coef'], label='Train Dice coef.')\nplt.plot(x, history.history['val_dice_coef'], label='Validation Dice coef.')\nplt.xlabel('Epoch')\nplt.suptitle('Learning curves')\nplt.legend();","execution_count":null,"outputs":[]},{"metadata":{},"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-u-net-mobilenetv2-256x256). "},{"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}