prisma_segmentation.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "prismaseg.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 Prisma-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, SeparableConv2D\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": "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 from directory.\n",
        "\n",
        "Ensure the images are in **RGB** format and masks (**ALPHA**) have pixel values **0 or 255**."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pd0l80v49c32",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "IMGDS=\"/content/portrait256/images\";\n",
        "MSKDS=\"/content/portrait256/masks\";\n",
        "\n",
        "# Total number of images\n",
        "num_images=len(os.listdir(IMGDS+\"/img\"))"
      ],
      "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",
        "CHECKPOINT=\"/content/drive/My Drive/portrait256/prisma-net-{epoch:02d}-{val_loss:.2f}.hdf5\"\n",
        "LOGS='./logs'\n",
        "BATCH_SIZE=64"
      ],
      "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:** The keras 'flow_from_directory' expects a specific directory structure for loading datasets. Your parent data-set directory should contain two sub-directories 'images' and 'masks'. Now, each of these directories should have a sub-directory(say 'img' and 'msk) for storing images or masks."
      ]
    },
    {
      "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.2,\n",
        "                     height_shift_range=0.2,\n",
        "                     zoom_range=0.2,\n",
        "                     horizontal_flip=True,\n",
        "                     validation_split=0.2\n",
        "                    )\n",
        "\n",
        "image_datagen = ImageDataGenerator(**data_gen_args)\n",
        "mask_datagen = ImageDataGenerator(**data_gen_args)\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_from_directory(\n",
        "    IMGDS,\n",
        "    batch_size=batch_sz,\n",
        "    shuffle=True,\n",
        "    subset='training',\n",
        "    color_mode=\"rgb\",\n",
        "    class_mode=None,\n",
        "    seed=seed)\n",
        "\n",
        "train_mask_generator = mask_datagen.flow_from_directory(\n",
        "    MSKDS,\n",
        "    batch_size=batch_sz,\n",
        "    shuffle=True,\n",
        "    subset='training',\n",
        "    color_mode=\"grayscale\",\n",
        "    class_mode=None,\n",
        "    seed=seed)\n",
        "\n",
        "\n",
        "val_image_generator = image_datagen.flow_from_directory(\n",
        "    IMGDS, \n",
        "batch_size = batch_sz,\n",
        "shuffle=True,\n",
        "subset='validation',\n",
        "color_mode=\"rgb\",\n",
        "class_mode=None,\n",
        "seed=seed)\n",
        "\n",
        "val_mask_generator = mask_datagen.flow_from_directory(\n",
        "     MSKDS,\n",
        "batch_size = batch_sz,\n",
        "shuffle=True,\n",
        "subset='validation',\n",
        "color_mode=\"grayscale\",\n",
        "class_mode=None,\n",
        "seed=seed)\n",
        "                   \n",
        "# combine generators into one which yields image and masks\n",
        "train_generator = zip(train_image_generator, train_mask_generator)\n",
        "val_generator = zip(val_image_generator, val_mask_generator)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3P0sZZUq_FQS",
        "colab_type": "text"
      },
      "source": [
        "**Model Architecture**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4iC43M9CE0hi",
        "colab_type": "text"
      },
      "source": [
        "The prisma-net basically uses a **U-Net** encoder-decoder structure. However, the architecture incorporates a few significant **changes**. Firstly, we replace the concatenation of features after upsampling with **element-wise addition**. Further, instead of normal Conv+ReLu block, we use a **residual block with depth-wise separable convolutions**. Finally, to improve the accuracy the **decoder** part contains **more blocks** than encoder."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UrBstYmvzrns",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def residual_block(x, nfilters):\n",
        "\n",
        "  y = Activation(\"relu\")(x)\n",
        "  y= SeparableConv2D(filters=nfilters, kernel_size=3, padding=\"same\")(y)\n",
        "  y = Activation(\"relu\")(y)\n",
        "  y= SeparableConv2D(filters=nfilters, kernel_size=3, padding=\"same\")(y)\n",
        "\n",
        "  z = Add()([x, y])\n",
        "\n",
        "  return z"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7lXscVnbkaZW",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def prisma_unet(finetuene=False, alpha=1):\n",
        "\n",
        "    input = Input(shape=(256,256,3))\n",
        "    \n",
        "    # Encoder part\n",
        "    x = Conv2D(filters=8, kernel_size=3,padding = 'same' )(input)\n",
        "\n",
        "    res1= residual_block(x, nfilters=8)\n",
        "    x = Conv2D(filters=32, kernel_size=3, strides=2, padding = 'same' )(res1)\n",
        "\n",
        "    res2= residual_block(x, nfilters=32)\n",
        "    x = Conv2D(filters=64, kernel_size=3, strides=2, padding = 'same' )(res2)\n",
        "\n",
        "    res3= residual_block(x, nfilters=64)\n",
        "    x = Conv2D(filters=128, kernel_size=3, strides=2, padding = 'same' )(res3)\n",
        "\n",
        "\n",
        "    x= residual_block(x, nfilters=128)\n",
        "    res4= residual_block(x, nfilters=128)\n",
        "    x = Conv2D(filters=128, kernel_size=3, strides=2, padding = 'same' )(res4)\n",
        "\n",
        "    x= residual_block(x, nfilters=128)\n",
        "    x= residual_block(x, nfilters=128)\n",
        "    x= residual_block(x, nfilters=128)\n",
        "    x= residual_block(x, nfilters=128)\n",
        "    x= residual_block(x, nfilters=128)\n",
        "    x= residual_block(x, nfilters=128)\n",
        "   \n",
        "\n",
        "    # Decoder part\n",
        "    x=Conv2DTranspose(filters=128, kernel_size=3, strides=2, padding = \"same\")(x)\n",
        "    x = Add()([x, res4 ])\n",
        "\n",
        "    x= residual_block(x, nfilters=128)\n",
        "    x= residual_block(x, nfilters=128)\n",
        "    x= residual_block(x, nfilters=128)\n",
        "\n",
        "    x = Conv2DTranspose(filters=64, kernel_size=3, strides=2, padding = 'same' )(x)\n",
        "    x = Add()([x, res3 ])\n",
        "   \n",
        "    x= residual_block(x, nfilters=64)\n",
        "    x= residual_block(x, nfilters=64)\n",
        "    x= residual_block(x, nfilters=64)\n",
        "\n",
        "    x = Conv2DTranspose(filters=32, kernel_size=3, strides=2, padding = 'same' )(x)\n",
        "    x = Add()([x, res2 ])\n",
        "   \n",
        "    x= residual_block(x, nfilters=32)\n",
        "    x= residual_block(x, nfilters=32)\n",
        "    x= residual_block(x, nfilters=32)\n",
        "\n",
        "    x = Conv2DTranspose(filters=8, kernel_size=3, strides=2, padding = 'same' )(x)\n",
        "    x = Add()([x, res1 ])\n",
        "  \n",
        "    x= residual_block(x, nfilters=8)\n",
        "    x= residual_block(x, nfilters=8)\n",
        "    x= residual_block(x, nfilters=8)\n",
        "  \n",
        "    x = Conv2DTranspose(1, (1,1), padding='same')(x)\n",
        "    x = Activation('sigmoid', name=\"op\")(x) \n",
        "\n",
        "    model = Model(inputs=input, outputs=x)\n",
        "    model.compile(loss='binary_crossentropy', optimizer=Adam(lr=1e-3),metrics=['accuracy'])\n",
        "\n",
        "    return model"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-_2WhLoR0wZj",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Get prisma network\n",
        "model = prisma_unet()\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='prisma-net.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=3, 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]"
      ],
      "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": [
        "Initially train the model for **300 epochs** using **supervisely person** dataset. Finally, train the model on **portrait datasets**, using the result of the previous step as initial values for **weights**.\n",
        "\n",
        "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/portrait256/prisma-net-07-0.09.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=4,\n",
        "    callbacks=callbacks_list)"
      ],
      "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/baby.jpg')\n",
        "\n",
        "# Load the model\n",
        "model=load_model('/content/drive/My Drive/portrait256/prisma-net-15-0.08.hdf5')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "joNyBzWJNZKI",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Inference\n",
        "im=im.resize((256,256),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.predict(img.reshape(1,256,256,3))\n",
        "out=np.float32((out>0.5))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5TyPqFEdBNyQ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Output mask\n",
        "plt.imshow(np.squeeze(out.reshape((256,256))))"
      ],
      "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.output\n",
        "newout=Reshape((65536,))(output)\n",
        "new_model=Model(model.input,newout)\n",
        "\n",
        "new_model.save('prisma-net.h5')\n",
        "\n",
        "# For Float32 Model\n",
        "converter = tf.lite.TFLiteConverter.from_keras_model_file('/content/prisma-net.h5')\n",
        "tflite_model = converter.convert()\n",
        "open(\"prisma-net.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/prisma-net.h5')\n",
        "converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_SIZE]\n",
        "tflite_model = converter.convert()\n",
        "open(\"prisma-net_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/prisma-net.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(\"prisma-net_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/prisma-net.h5',compile=False)\n",
        "test_imgs=np.load('/content/kids.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,256,256,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": []
    }
  ]
}