portrait_segmentation_v3.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Oldseg.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1Swu0Yhy8lbz",
        "colab_type": "text"
      },
      "source": [
        "**Portrait Segmentation Using Mobile-Unet**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yxR8NdhbEcGh",
        "colab_type": "text"
      },
      "source": [
        "Set up the GPU runtime"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ncUboS6EcRJH",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        " # Check GPU\n",
        "!nvidia-smi"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "s9gDKHrJhsRM",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Mount G-drive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YmZnqybY_XI4",
        "colab_type": "text"
      },
      "source": [
        "**Imports**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KXFHcZ7Ub-cr",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Import libraries\n",
        "import os\n",
        "import tensorflow as tf\n",
        "import keras\n",
        "from keras.models import Model\n",
        "from keras.layers import Dense, Input,Flatten, concatenate,Reshape, Conv2D, MaxPooling2D, Lambda,Activation,Conv2DTranspose\n",
        "from keras.layers import UpSampling2D, Conv2DTranspose, BatchNormalization, Dropout, DepthwiseConv2D, Add\n",
        "from keras.callbacks import TensorBoard, ModelCheckpoint, Callback, ReduceLROnPlateau\n",
        "from keras.regularizers import l1\n",
        "from keras.optimizers import SGD, Adam\n",
        "import keras.backend as K\n",
        "from keras.utils import plot_model\n",
        "from keras.callbacks import TensorBoard, ModelCheckpoint, Callback\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "from scipy.ndimage.filters import gaussian_filter\n",
        "from random import randint\n",
        "from keras.models import load_model\n",
        "from keras.preprocessing.image import ImageDataGenerator\n",
        "from PIL import Image\n",
        "import matplotlib.pyplot as plt\n",
        "from random import randint\n",
        "%matplotlib inline"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "435W142g_dFt",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Keras optimization library\n",
        "!pip install kito\n",
        "from kito import reduce_keras_model"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vuT3oK359WhC",
        "colab_type": "text"
      },
      "source": [
        "**Load dataset**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v4rYsYWLOkaH",
        "colab_type": "text"
      },
      "source": [
        "Load the datset for training the model.\n",
        "\n",
        "Ensure the images are in **RGB** format and masks (**ALPHA**) have pixel values **0 or 255**."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HOmrS9429b54",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Load the dataset\n",
        "x_train=np.load(\"/content/drive/My Drive/finsegds/img_uint8.npy\")\n",
        "y_train=np.load(\"/content/drive/My Drive/finsegds/msk_uint8.npy\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pd0l80v49c32",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Verify the mask shape and values\n",
        "print(np.unique(y_train))\n",
        "print(y_train.shape)\n",
        "\n",
        "# Total number of images\n",
        "num_images=x_train.shape[0]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3ts7trncShBU",
        "colab_type": "text"
      },
      "source": [
        "Copy pretrained model to local runtime disk. Save the checkpoints to your google drive (safe)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0RSXSLX3RW7j",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Configure save paths and batch size\n",
        "PRETRAINED='/content/pretrained_model.hdf5'\n",
        "CHECKPOINT=\"/content/drive/My Drive/finsegds/munet_mnv3_wm10-{epoch:02d}-{val_loss:.2f}.hdf5\"\n",
        "LOGS='./logs'\n",
        "BATCH_SIZE=64"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ojpg0uXx_NMo",
        "colab_type": "text"
      },
      "source": [
        "**Preprocessing**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mtvn3SMeEpUc",
        "colab_type": "text"
      },
      "source": [
        "Normalize the source images  at runtime; but do not modify the masks"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "c7QQjir_vK6A",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Preprocessing function (runtime)\n",
        "def normalize_batch(imgs):\n",
        "    if imgs.shape[-1] > 1 :\n",
        "      return (imgs -  np.array([0.50693673, 0.47721124, 0.44640532])) /np.array([0.28926975, 0.27801928, 0.28596011])\n",
        "    else:\n",
        "      return imgs.round()\n",
        "def denormalize_batch(imgs,should_clip=True):\n",
        "    imgs= (imgs * np.array([0.28926975, 0.27801928, 0.28596011])) + np.array([0.50693673, 0.47721124, 0.44640532])\n",
        "    \n",
        "    if should_clip:\n",
        "        imgs= np.clip(imgs,0,1)\n",
        "    return imgs"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fuyYcl-k_KQb",
        "colab_type": "text"
      },
      "source": [
        "**Data Generator**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ogHJrBXXE8cX",
        "colab_type": "text"
      },
      "source": [
        "Create a data generator to load images and masks together at runtime. \n",
        "Use same seed for performing run-time augmentation for images and masks. Here we use  80/20 tran-val split.\n",
        "\n",
        "**Note:** Currently the data generator loads the entire augmented dataset(npy) into memory (RAM) , so there is a good chance that the application would crash if there is not enough memory. You may alternatively use a different data generator for loading images from directories, after ensuring a proper directory structure."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Npm8xJCMcdcm",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Data generator for training and validation\n",
        "\n",
        "data_gen_args = dict(rescale=1./255,\n",
        "                     width_shift_range=0.1,\n",
        "                     height_shift_range=0.1,\n",
        "                     zoom_range=0.2,\n",
        "                     horizontal_flip=True,\n",
        "                     validation_split=0.2\n",
        "                    )\n",
        "\n",
        "image_datagen = ImageDataGenerator(**data_gen_args, preprocessing_function=normalize_batch)\n",
        "mask_datagen = ImageDataGenerator(**data_gen_args,  preprocessing_function=normalize_batch)\n",
        "\n",
        "# Provide the same seed and keyword arguments to the fit and flow methods\n",
        "seed = 1\n",
        "batch_sz=BATCH_SIZE\n",
        "\n",
        "# Train-val split (80-20)\n",
        "num_train=int(num_images*0.8)\n",
        "num_val=int(num_images*0.2) \n",
        "\n",
        "\n",
        "train_image_generator = image_datagen.flow(\n",
        "    x_train,\n",
        "    batch_size=batch_sz,\n",
        "    shuffle=True,\n",
        "    subset='training',\n",
        "    seed=seed)\n",
        "\n",
        "train_mask_generator = mask_datagen.flow(\n",
        "    y_train,\n",
        "    batch_size=batch_sz,\n",
        "    shuffle=True,\n",
        "    subset='training',\n",
        "    seed=seed)\n",
        "\n",
        "\n",
        "val_image_generator = image_datagen.flow(\n",
        "    x_train, \n",
        "batch_size = batch_sz,\n",
        "shuffle=True,\n",
        "subset='validation',\n",
        "seed=seed)\n",
        "\n",
        "val_mask_generator = mask_datagen.flow(\n",
        "     y_train,\n",
        "batch_size = batch_sz,\n",
        "shuffle=True,\n",
        "subset='validation',\n",
        "seed=seed)\n",
        "\n",
        "                     \n",
        "# combine generators into one which yields image and masks\n",
        "\n",
        "train_generator = zip(train_image_generator, train_mask_generator)\n",
        "val_generator = zip(val_image_generator, val_mask_generator)\n",
        "\n",
        "# Delete numpy arrays to free memory\n",
        "del(x_train)\n",
        "del(y_train)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3P0sZZUq_FQS",
        "colab_type": "text"
      },
      "source": [
        "**Model Architecture**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y3VOYGIcGb5d",
        "colab_type": "text"
      },
      "source": [
        "Here we use the minimalistic version  of **Mobilent v3** with **width multiplier 1.0 or 0.5** as encoder (feature extractor).  \n",
        "\n",
        "For the **decoder part**, we can use a upsampling block with **Transpose Convolution** of **stride of 2**. Ensure proper **skip connections** between encoder and decoder parts for better results."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ie7ZO7X52oLP",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# MobilenetV3 Blocks\n",
        " \n",
        "def deconv_block(tensor, nfilters, size=3, padding='same', kernel_initializer = 'he_normal'):\n",
        "    \n",
        "    y = Conv2DTranspose(filters=nfilters, kernel_size=size, strides=2, padding = padding, kernel_initializer = kernel_initializer)(tensor)\n",
        "    y = BatchNormalization()(y)\n",
        "    y = Dropout(0.5)(y)\n",
        "    y = Activation(\"relu\")(y)\n",
        "    \n",
        "    return y\n",
        "\n",
        "def conv_block(tensor, filters, strides, size=3):\n",
        "    x = Conv2D(filters=filters, kernel_size=size, strides=strides, padding = 'same' )(tensor)\n",
        "    x = BatchNormalization()(x)\n",
        "    x = Activation(\"relu\")(x)\n",
        "\n",
        "    return x\n",
        "\n",
        "def bottleneck(tensor, nfilters, exp_ch, strides=1, alpha=1,residual=False):\n",
        "\n",
        "  x = conv_block(tensor, filters=exp_ch, size=1, strides=1)\n",
        "\n",
        "  y = DepthwiseConv2D(kernel_size=3, strides=strides, depth_multiplier=1, padding='same')(x)\n",
        "  y = BatchNormalization()(y)\n",
        "  y = Activation(\"relu\")(y)\n",
        "\n",
        "  z = Conv2D(filters=int(nfilters*alpha), kernel_size=1, strides=1, padding = 'same' )(y)\n",
        "  z = BatchNormalization()(z)\n",
        "  \n",
        "  if residual:\n",
        "      z = Add()([z, tensor])\n",
        "\n",
        "  return z\n",
        "\n",
        "\n",
        "# MobilenetV3 Base\n",
        "def get_mobilenetv3(pretrained=False, alpha=1):\n",
        "     \n",
        "     input = Input(shape=(224,224,4))\n",
        "\n",
        "     x = conv_block(input, filters=16, size=3, strides=2)\n",
        "\n",
        "     x = DepthwiseConv2D(kernel_size=3, strides=2, depth_multiplier=1, padding='same')(x)\n",
        "     x = BatchNormalization()(x)\n",
        "     x = Activation(\"relu\")(x)\n",
        "\n",
        "     x = Conv2D(filters=int(16*alpha), kernel_size=1, strides=1, padding = 'same' )(x)\n",
        "     x = BatchNormalization()(x)\n",
        "\n",
        "     x = bottleneck(x, nfilters=24, exp_ch= 72, strides=2, alpha=alpha)\n",
        "     x = bottleneck(x, nfilters=24, exp_ch= 88, strides=1, alpha=alpha, residual=True)\n",
        "     \n",
        "     x = bottleneck(x, nfilters=40,exp_ch= 96, strides=2, alpha=alpha)\n",
        "     x = bottleneck(x, nfilters=40,exp_ch= 240, strides=1, alpha=alpha, residual=True)\n",
        "     x = bottleneck(x, nfilters=40,exp_ch= 240, strides=1, alpha=alpha, residual=True)\n",
        "     x = bottleneck(x, nfilters=48,exp_ch= 120, strides=1, alpha=alpha)\n",
        "     x = bottleneck(x, nfilters=48,exp_ch= 144, strides=1, alpha=alpha, residual=True)\n",
        "     x = bottleneck(x, nfilters=96, exp_ch=288,strides=2, alpha=alpha)\n",
        "     x = bottleneck(x, nfilters=96,exp_ch= 576, strides=1, alpha=alpha, residual=True)\n",
        "     x = bottleneck(x, nfilters=96,exp_ch= 576, strides=1, alpha=alpha, residual=True)\n",
        "\n",
        "     x=conv_block(x, 576, strides=1, size=1)\n",
        "    \n",
        "     model = Model(inputs=input, outputs=x)\n",
        "     return model\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jVD8elwhqBhI",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#  Model architecture: mnv3_unet\n",
        "\n",
        "def mnv3_unet(finetuene=False, pretrained=False, alpha=1):\n",
        "\n",
        "    # Load pretrained model (if any)\n",
        "    if (pretrained):\n",
        "       model=load_model(PRETRAINED)\n",
        "       print(\"Loaded pretrained model ...\\n\")\n",
        "       return model\n",
        "\n",
        "    # Encoder\n",
        "    mnv3 = get_mobilenetv3(pretrained=False, alpha=alpha)\n",
        "\n",
        "    # Decoder\n",
        "    x=mnv3.output\n",
        "\n",
        "    x = deconv_block(x, int(256*alpha))\n",
        "    x = concatenate([x, mnv3.layers[71].output], axis = 3)\n",
        "    \n",
        "    x = deconv_block(x, int(128*alpha))\n",
        "    x = concatenate([x, mnv3.layers[28].output], axis = 3)\n",
        "                \n",
        "    x = deconv_block(x, int(64*alpha))\n",
        "    x = concatenate([x, mnv3.layers[11].output], axis = 3)\n",
        "    \n",
        "    x = deconv_block(x, int(32*alpha))\n",
        "    x = concatenate([x, mnv3.layers[3].output], axis = 3)\n",
        "                \n",
        "\n",
        "    x = Conv2DTranspose(filters=int(16*alpha), kernel_size=3, strides=2, padding='same', kernel_initializer = 'he_normal')(x)\n",
        "    x = BatchNormalization()(x)\n",
        "    x = Activation(\"relu\")(x)\n",
        "    \n",
        "   \n",
        "    x = Conv2DTranspose(1, (1,1), padding='same')(x)\n",
        "    x = Activation('sigmoid', name=\"op\")(x)\n",
        "    \n",
        "    \n",
        "    model = Model(inputs=mnv3.input, outputs=x)\n",
        "    \n",
        "    \n",
        "    model.compile(loss='binary_crossentropy', optimizer=Adam(lr=1e-3),metrics=['accuracy'])\n",
        "    return model\n",
        "\n",
        "model=mnv3_unet(finetuene=False, pretrained=False, alpha=1)\n",
        "\n",
        "# Model summary\n",
        "model.summary()\n",
        "\n",
        "# Layer specifications\n",
        "for i, layer in enumerate(model.layers):\n",
        "    print(i, layer.output.name, layer.output.shape)\n",
        "\n",
        "# Plot model architecture\n",
        "plot_model(model, to_file='portrait_mnv3.png')\n",
        "\n",
        "# Save checkpoints\n",
        "checkpoint = ModelCheckpoint(CHECKPOINT, monitor='val_loss', verbose=1, save_weights_only=False , save_best_only=True, mode='min')\n",
        "\n",
        "# Callbacks \n",
        "reduce_lr = ReduceLROnPlateau(factor=0.5, patience=15, min_lr=0.000001, verbose=1)\n",
        "tensorboard = TensorBoard(log_dir=LOGS, histogram_freq=0,\n",
        "                          write_graph=True, write_images=True)\n",
        "\n",
        "callbacks_list = [checkpoint, tensorboard, reduce_lr]\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sUABPgIS3tpV",
        "colab_type": "text"
      },
      "source": [
        "**Train**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-A5lcnZ0Iohs",
        "colab_type": "text"
      },
      "source": [
        "Train the model for **300 epochs** with our custom data generator. Use keras callbacks for **tensorboard** visulaization and **learning rate decay** as shown below. You can resume your training from a previous session by loading the entire **pretrained model** (weights  & optimzer state) as a hdf5 file."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "t005QITaiorA",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Load pretrained model (if any)\n",
        "model=load_model('/content/drive/My Drive/finsegds/munet_mnv3_wm10-81-0.07.hdf5')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mCb3a6BHj60k",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Train the model\n",
        "model.fit_generator(\n",
        "    train_generator,\n",
        "    epochs=300,\n",
        "    steps_per_epoch=num_train/batch_sz,\n",
        "    validation_data=val_generator, \n",
        "    validation_steps=num_val/batch_sz,\n",
        "    use_multiprocessing=True,\n",
        "    workers=2,\n",
        "    callbacks=callbacks_list)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tj0oSNBJKLe4",
        "colab_type": "text"
      },
      "source": [
        "**Evaluate**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vBbRoF1vKmI5",
        "colab_type": "text"
      },
      "source": [
        "Evalute the performance of the model on a test data-set."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wng8tdoDKPQZ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Load a trained model checkpoint\n",
        "model=load_model('/content/munet_mnv3_wm10-81-0.07.hdf5')\n",
        "\n",
        "# Load a test dataset\n",
        "new_xtest=x_train[14958:,...]\n",
        "new_ytest=y_train[14958:,...]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UGE76DBQKP2b",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Evaluate model \n",
        "score = model.evaluate(np.float32(new_xtest/255.0), np.float32(new_ytest/255.0), verbose=0)\n",
        "# Print loss and accuracy\n",
        "print('Test loss:', score[0])\n",
        "print('Test accuracy:', score[1])\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Eg4KEzgz3Wyz",
        "colab_type": "text"
      },
      "source": [
        "**Optimize**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B_81SoSvKvX4",
        "colab_type": "text"
      },
      "source": [
        "Using the kito library, you can optimize the model by folding the batch norms. This does not change the model behaviour or accuracy; but helps us to reduce the number of layers."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2J5spTGWlB8A",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Optimize model by folding batch-norms\n",
        "model_reduced = reduce_keras_model(model)\n",
        "model_reduced.summary()\n",
        "model_reduced.save('munet_mnv3_wm10_bnoptimized.h5')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KdpmP99C3b5l",
        "colab_type": "text"
      },
      "source": [
        "**Test**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "06PZInENLLti",
        "colab_type": "text"
      },
      "source": [
        "Test the model on a new portrait image and plot the results."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wl3gUNQkNYbb",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Load a test image\n",
        "im=Image.open('/content/sf99.png')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "joNyBzWJNZKI",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Inference\n",
        "im=im.resize((224,224),Image.ANTIALIAS)\n",
        "img=np.float32(np.array(im)/255.0)\n",
        "plt.imshow(img[:,:,0:3])\n",
        "img=img[:,:,0:3]\n",
        "\n",
        "# Reshape input and threshold output\n",
        "out=model_reduced.predict(img.reshape(1,224,224,3))\n",
        "out=np.float32((out>0.5))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DBIcgqNWNgYp",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Output mask\n",
        "plt.imshow(np.squeeze(out.reshape((224,224))))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mo-1rEzx4MH7",
        "colab_type": "text"
      },
      "source": [
        "**Export Model**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gOSneVrHLgwO",
        "colab_type": "text"
      },
      "source": [
        "Export the model to **tflite** format for **real-time** inference on a **smart-phone**."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZlRxm8m5kLxx",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Flatten output and save model\n",
        "output = model_reduced.output\n",
        "newout=Reshape((50176,))(output)\n",
        "new_model=Model(model_reduced.input,newout)\n",
        "\n",
        "new_model.save('munet_mnv3_wm10.h5')\n",
        "\n",
        "# For Float32 Model\n",
        "converter = tf.lite.TFLiteConverter.from_keras_model_file('/content/munet_mnv3_wm10.h5')\n",
        "tflite_model = converter.convert()\n",
        "open(\"munet_mnv3_wm10.tflite\", \"wb\").write(tflite_model)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ouCJJyl4yqhE",
        "colab_type": "text"
      },
      "source": [
        "**Post-training Quantization**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E7Ilquw8L5zM",
        "colab_type": "text"
      },
      "source": [
        "We can **reduce the model size and latency** by performing post training quantization. Fixed precison conversion (**UINT8**) allows us to reduce the model size significantly by quantizing the model weights.We can run this model on the mobile **CPU**. The **FP16** (experimental) conversion allows us to reduce the model size by half and the corresponding model can be run directly on mobile **GPU**."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TMvRzTYVbnzZ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# For UINT8 Quantization\n",
        "\n",
        "converter = tf.lite.TFLiteConverter.from_keras_model_file('/content/munet_mnv3_wm10.h5')\n",
        "converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_SIZE]\n",
        "tflite_model = converter.convert()\n",
        "open(\"munet_mnv3_wm10_uint8.tflite\", \"wb\").write(tflite_model)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "x5BMOoZBk1Sq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# For Float16 Quantization (Experimental)\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        "converter = tf.lite.TFLiteConverter.from_keras_model_file('/content/munet_mnv3_wm10.h5')\n",
        "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "converter.target_spec.supported_types = [tf.lite.constants.FLOAT16]\n",
        "tflite_model = converter.convert()\n",
        "open(\"munet_mnv3_wm10_fp16.tflite\", \"wb\").write(tflite_model)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7b8gobWwY8z8",
        "colab_type": "text"
      },
      "source": [
        "**Plot sample output**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jx9mu6umZVjn",
        "colab_type": "text"
      },
      "source": [
        "Load the test data as a batch using a numpy array. \n",
        "\n",
        "Crop the image using the output mask and plot the result."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eHT4q8teyFmX",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Load test images and model\n",
        "model=load_model('/content/munet_mnv3_wm10.h5',compile=False)\n",
        "test_imgs=np.load('/content/timg_uint8.npy')\n",
        "test_imgs= np.float32(np.array(test_imgs)/255.0)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IYRxEoKvMFXT",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Perform batch prediction\n",
        "out=model.predict(test_imgs)\n",
        "out=np.float32((out>0.5))\n",
        "out=out.reshape((4,224,224,1))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bBSM2BGlMUNr",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Plot the output using matplotlib\n",
        "fig=plt.figure(figsize=(16, 16))\n",
        "columns = 4\n",
        "rows = 2\n",
        "\n",
        "for i in range(1, columns+1):\n",
        "    img = test_imgs[i-1].squeeze()\n",
        "    fig.add_subplot(rows, columns, i)\n",
        "    plt.imshow(img)\n",
        "plt.show()\n",
        "\n",
        "fig=plt.figure(figsize=(16, 16))\n",
        "columns = 4\n",
        "rows = 2\n",
        "\n",
        "for i in range(1, columns+1):\n",
        "    img = out[i-1].squeeze()/255.0\n",
        "    fig.add_subplot(rows, columns, 4+i)\n",
        "    plt.imshow(out[i-1]*test_imgs[i-1])\n",
        "\n",
        "plt.show()"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}