slim512.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "slim512.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LE3iLGtVrGAA"
      },
      "source": [
        "**SlimNet: Real-time Portrait Segmentation on High Resolution Images**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UoQ0Qv8FrtvT"
      },
      "source": [
        "Slim-net is a light weight CNN for performing **real-time portrait segmentation** on mobile devices, using high resolution images. We were able to achieve **99% training accuracy** on the aisegment portrait dataset and run the model(**1.5MB**) on a mid-range android smartphone at **20 fps** on deployment. Using the high resolution input images, we were able to preserve **fine details** and **avoid sharp edges** on segmentation masks, during inference .The architecture is heavily inspired from the mediapipe **hair-segmentation** model for android and the tflite model runs on any **android** device without additional API's."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gVew8k6Ubygn"
      },
      "source": [
        "**Environment and Datset**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mIEKFDHPZ2BO"
      },
      "source": [
        "Choose a **GPU runtime** in colab, for training the network."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "j-n95-nM0ZsJ"
      },
      "source": [
        "!nvidia-smi"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gdNzqjXjaYOu"
      },
      "source": [
        "Install **tensorflow 1.15** version, using pip command."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xndLfFgL8KZw"
      },
      "source": [
        "# Install TensorFlow 2.0 (GPU)\n",
        "!pip install tensorflow-gpu==1.15"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D5VnwwTHa27O"
      },
      "source": [
        "Extract the **dataset** for training, from google-drive."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8thgmfRjnmBO"
      },
      "source": [
        "!unzip /content/drive/My\\ Drive/slim512/aiseg_small.zip"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LP6eNNQ0b5KC"
      },
      "source": [
        "**Packages and Libraries**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1uc_i_Nt8VEZ"
      },
      "source": [
        "# Import the packages\n",
        "from tensorflow.keras.preprocessing.image import ImageDataGenerator\n",
        "from tensorflow.keras.callbacks import TensorBoard, ModelCheckpoint, Callback, ReduceLROnPlateau\n",
        "from tensorflow.keras.layers import Input, Conv2D, Add, Multiply, Reshape , MaxPool2D, Conv2DTranspose, PReLU, concatenate, Lambda\n",
        "from tensorflow.keras.layers import Dropout, BatchNormalization, concatenate, Activation, AveragePooling2D, UpSampling2D\n",
        "from tensorflow.keras.layers import Flatten, Dense, GlobalAveragePooling2D, BatchNormalization, DepthwiseConv2D, SeparableConv2D\n",
        "from tensorflow.keras.callbacks import TensorBoard, ModelCheckpoint, Callback, ReduceLROnPlateau\n",
        "from tensorflow.keras.models import Model\n",
        "from tensorflow.keras.utils import plot_model\n",
        "from tensorflow.keras.optimizers import SGD, Adam\n",
        "from tensorflow.keras.regularizers import l2\n",
        "from tensorflow.keras.models import load_model\n",
        "from tensorflow.keras.callbacks import TensorBoard, ModelCheckpoint, Callback\n",
        "import matplotlib.pyplot as plt\n",
        "import tensorflow as tf\n",
        "import os, cv2, imageio\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "from PIL import Image, ImageEnhance\n",
        "from IPython.display import clear_output\n",
        "from sklearn.utils import shuffle\n",
        "from sklearn.model_selection import train_test_split"
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "svO8Q_dmi7lz"
      },
      "source": [
        "**Data-loader and Augmentations**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_79iT--J3eDn"
      },
      "source": [
        "Use the **custom data-loader** to perform  **augmentaion** on-the-fly and feed the network with **batch of images**. Here we use the augmentaion like **brightness, saturation, contrast, cropping, flipping** etc for training the model. If your masks are not in **raw format**, then you need to convert them into sparse labels(color indexed) for training with **SparseCategoricalCrossentropy** loss (i.e 0 for bg and 1 for fg). If you want to provide labels using **one-hot** representation, please use **CategoricalCrossentropy** loss."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YpnYEtUeXkkf"
      },
      "source": [
        "import tensorflow as tf\n",
        "import random\n",
        "\n",
        "class DataLoader(object):\n",
        "    \"\"\"A TensorFlow Dataset API based loader for semantic segmentation problems.\"\"\"\n",
        "\n",
        "    def __init__(self, image_paths, mask_paths, image_size, channels=[3, 3], crop_percent=None, palette=None, seed=None):\n",
        "        \"\"\"\n",
        "        Initializes the data loader object\n",
        "        Args:\n",
        "            image_paths: List of paths of train images.\n",
        "            mask_paths: List of paths of train masks (segmentation masks)\n",
        "            image_size: Tuple of (Height, Width), the final height \n",
        "                        of the loaded images.\n",
        "            channels: List of ints, first element is number of channels in images,\n",
        "                      second is the number of channels in the mask image (needed to\n",
        "                      correctly read the images into tensorflow.)\n",
        "            crop_percent: Float in the range 0-1, defining percentage of image \n",
        "                          to randomly crop.\n",
        "            palette: A list of RGB pixel values in the mask. If specified, the mask\n",
        "                     will be one hot encoded along the channel dimension.\n",
        "            seed: An int, if not specified, chosen randomly. Used as the seed for \n",
        "                  the RNG in the data pipeline.\n",
        "        \"\"\"\n",
        "        self.image_paths = image_paths\n",
        "        self.mask_paths = mask_paths\n",
        "        self.palette = palette\n",
        "        self.image_size = image_size\n",
        "        if crop_percent is not None:\n",
        "            if 0.0 < crop_percent <= 1.0:\n",
        "                self.crop_percent = tf.constant(crop_percent, tf.float32)\n",
        "            elif 0 < crop_percent <= 100:\n",
        "                self.crop_percent = tf.constant(crop_percent / 100., tf.float32)\n",
        "            else:\n",
        "                raise ValueError(\"Invalid value entered for crop size. Please use an \\\n",
        "                                  integer between 0 and 100, or a float between 0 and 1.0\")\n",
        "        else:\n",
        "            self.crop_percent = None\n",
        "        self.channels = channels\n",
        "        if seed is None:\n",
        "            self.seed = random.randint(0, 1000)\n",
        "        else:\n",
        "            self.seed = seed\n",
        "\n",
        "    def _corrupt_brightness(self, image, mask):\n",
        "        \"\"\"\n",
        "        Radnomly applies a random brightness change.\n",
        "        \"\"\"\n",
        "        cond_brightness = tf.cast(tf.random.uniform(\n",
        "            [], maxval=2, dtype=tf.int32), tf.bool)\n",
        "        image = tf.cond(cond_brightness, lambda: tf.image.random_brightness(\n",
        "            image, 0.1), lambda: tf.identity(image))\n",
        "        return image, mask\n",
        "\n",
        "\n",
        "    def _corrupt_contrast(self, image, mask):\n",
        "        \"\"\"\n",
        "        Randomly applies a random contrast change.\n",
        "        \"\"\"\n",
        "        cond_contrast = tf.cast(tf.random.uniform(\n",
        "            [], maxval=2, dtype=tf.int32), tf.bool)\n",
        "        image = tf.cond(cond_contrast, lambda: tf.image.random_contrast(\n",
        "            image, 0.1, 0.8), lambda: tf.identity(image))\n",
        "        return image, mask\n",
        "\n",
        "\n",
        "    def _corrupt_saturation(self, image, mask):\n",
        "        \"\"\"\n",
        "        Randomly applies a random saturation change.\n",
        "        \"\"\"\n",
        "        cond_saturation = tf.cast(tf.random.uniform(\n",
        "            [], maxval=2, dtype=tf.int32), tf.bool)\n",
        "        image = tf.cond(cond_saturation, lambda: tf.image.random_saturation(\n",
        "            image, 0.1, 0.8), lambda: tf.identity(image))\n",
        "        return image, mask\n",
        "\n",
        "\n",
        "    def _crop_random(self, image, mask):\n",
        "        \"\"\"\n",
        "        Randomly crops image and mask in accord.\n",
        "        \"\"\"\n",
        "        cond_crop_image = tf.cast(tf.random.uniform(\n",
        "            [], maxval=2, dtype=tf.int32, seed=self.seed), tf.bool)\n",
        "        cond_crop_mask = tf.cast(tf.random.uniform(\n",
        "            [], maxval=2, dtype=tf.int32, seed=self.seed), tf.bool)\n",
        "\n",
        "        shape = tf.cast(tf.shape(image), tf.float32)\n",
        "        h = tf.cast(shape[0] * self.crop_percent, tf.int32)\n",
        "        w = tf.cast(shape[1] * self.crop_percent, tf.int32)\n",
        "\n",
        "        image = tf.cond(cond_crop_image, lambda: tf.image.random_crop(\n",
        "            image, [h, w, self.channels[0]], seed=self.seed), lambda: tf.identity(image))\n",
        "        mask = tf.cond(cond_crop_mask, lambda: tf.image.random_crop(\n",
        "            mask, [h, w, self.channels[1]], seed=self.seed), lambda: tf.identity(mask))\n",
        "\n",
        "        return image, mask\n",
        "\n",
        "\n",
        "    def _flip_left_right(self, image, mask):\n",
        "        \"\"\"\n",
        "        Randomly flips image and mask left or right in accord.\n",
        "        \"\"\"\n",
        "        image = tf.image.random_flip_left_right(image, seed=self.seed)\n",
        "        mask = tf.image.random_flip_left_right(mask, seed=self.seed)\n",
        "\n",
        "        return image, mask\n",
        "\n",
        "\n",
        "    def _resize_data(self, image, mask):\n",
        "        \"\"\"\n",
        "        Resizes images to specified size and normalizes the image: [0...1]\n",
        "        \"\"\"\n",
        "        image = tf.image.resize(image, [self.image_size, self.image_size]) /255.0\n",
        "        mask = tf.image.resize(mask, [self.image_size, self.image_size], method='nearest')//tf.reduce_max(mask) # masks should be binary with 0 representing background\n",
        "        \n",
        "        return image, mask\n",
        "\n",
        "\n",
        "    def _parse_data(self, image_paths, mask_paths):\n",
        "        \"\"\"\n",
        "        Reads image and mask files depending on\n",
        "        specified extension.\n",
        "        \"\"\"\n",
        "        image_content = tf.io.read_file(image_paths)\n",
        "        mask_content = tf.io.read_file(mask_paths)\n",
        "\n",
        "        images = tf.image.decode_jpeg(image_content, channels=self.channels[0])\n",
        "        masks = tf.image.decode_jpeg(mask_content, channels=self.channels[1])\n",
        "\n",
        "        return images, masks\n",
        "\n",
        "\n",
        "    def _one_hot_encode(self, image, mask):\n",
        "        \"\"\"\n",
        "        Converts mask to a one-hot encoding specified by the semantic map.\n",
        "        \"\"\"\n",
        "        one_hot_map = []\n",
        "        for colour in self.palette:\n",
        "            class_map = tf.reduce_all(tf.equal(mask, colour), axis=-1)\n",
        "            one_hot_map.append(class_map)\n",
        "        one_hot_map = tf.stack(one_hot_map, axis=-1)\n",
        "        one_hot_map = tf.cast(one_hot_map, tf.float32)\n",
        "        \n",
        "        return image, one_hot_map\n",
        "\n",
        "    def data_batch(self, batch_size, augment, shuffle=False, one_hot_encode=False):\n",
        "        \"\"\"\n",
        "        Reads data, normalizes it, shuffles it, then batches it, returns a\n",
        "        the next element in dataset op and the dataset initializer op.\n",
        "        Inputs:\n",
        "            batch_size: Number of images/masks in each batch returned.\n",
        "            augment: Boolean, whether to augment data or not.\n",
        "            shuffle: Boolean, whether to shuffle data in buffer or not.\n",
        "            one_hot_encode: Boolean, whether to one hot encode the mask image or not.\n",
        "                            Encoding will done according to the palette specified when\n",
        "                            initializing the object.\n",
        "        Returns:\n",
        "            data: A tf dataset object.\n",
        "        \"\"\"\n",
        "\n",
        "        # Create dataset out of the 2 files:\n",
        "        data = tf.data.Dataset.from_tensor_slices((self.image_paths, self.mask_paths))\n",
        "\n",
        "        # Parse images and labels\n",
        "        data = data.map(self._parse_data, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "\n",
        "        # If augmentation is to be applied\n",
        "        if augment:\n",
        "            data = data.map(self._corrupt_brightness,\n",
        "                            num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "\n",
        "            data = data.map(self._corrupt_contrast,\n",
        "                            num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "\n",
        "            data = data.map(self._corrupt_saturation,\n",
        "                            num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "\n",
        "            if self.crop_percent is not None:\n",
        "                data = data.map(self._crop_random, \n",
        "                                num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "\n",
        "            data = data.map(self._flip_left_right,\n",
        "                            num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "\n",
        "        # Resize to smaller dims for speed\n",
        "        data = data.map(self._resize_data, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "\n",
        "        # One hot encode the mask\n",
        "        if one_hot_encode:\n",
        "            if self.palette is None:\n",
        "                raise ValueError('No Palette for one-hot encoding specified in the data loader! \\\n",
        "                                  please specify one when initializing the loader.')\n",
        "            data = data.map(self._one_hot_encode, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "\n",
        "        if shuffle:\n",
        "            # Shuffle, repeat, batch and prefetch\n",
        "            data = data.shuffle(1000).repeat().batch(batch_size).prefetch(tf.data.experimental.AUTOTUNE)\n",
        "        else:\n",
        "            # Batch and prefetch\n",
        "            data = data.repeat().batch(batch_size).prefetch(tf.data.experimental.AUTOTUNE)\n",
        "\n",
        "        return data"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0rTkh4KE6cFy"
      },
      "source": [
        "Configure the **data loader, image paths and logging** options."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RG43zijLnd80"
      },
      "source": [
        "import tensorflow as tf\n",
        "import os\n",
        "\n",
        "IMAGE_DIR_PATH = '/content/imgs'\n",
        "MASK_DIR_PATH = '/content/msks'\n",
        "\n",
        "image_paths = [os.path.join(IMAGE_DIR_PATH, x) for x in sorted(os.listdir(IMAGE_DIR_PATH)) if x.endswith('.jpg')]\n",
        "mask_paths = [os.path.join(MASK_DIR_PATH, x) for x in sorted(os.listdir(MASK_DIR_PATH)) if x.endswith('.png')]\n",
        "\n",
        "train_image_paths, val_image_paths, train_mask_paths, val_mask_paths = train_test_split(image_paths, mask_paths, test_size = 0.2, random_state = 0)\n",
        "\n",
        "CHECKPOINT=\"/content/drive/My Drive/slim512/ckpt/slim-net-{epoch:02d}-{val_loss:.2f}.hdf5\"\n",
        "LOGS='./logs'\n",
        "\n",
        "num_train=len(train_image_paths)\n",
        "num_val=len(val_image_paths)\n",
        "batch_sz=64\n",
        "epochs=100\n",
        "\n",
        "\n",
        "# Initialize the dataloader object\n",
        "train_dataset = DataLoader(image_paths=train_image_paths,\n",
        "                     mask_paths=train_mask_paths,\n",
        "                     image_size=512,\n",
        "                     crop_percent=0.8,\n",
        "                     channels=[3, 1],\n",
        "                     seed=47)\n",
        "val_dataset = DataLoader(image_paths=val_image_paths,\n",
        "                     mask_paths=val_mask_paths,\n",
        "                     image_size=512,\n",
        "                     crop_percent=0.8,\n",
        "                     channels=[3, 1],\n",
        "                     seed=47)\n",
        "\n",
        "# Parse the images and masks, and return the data in batches, augmented optionally.\n",
        "train_dataset = train_dataset.data_batch(batch_size=batch_sz,\n",
        "                             augment=True, \n",
        "                             shuffle=True)\n",
        "val_dataset = val_dataset.data_batch(batch_size=batch_sz,\n",
        "                             augment=False, \n",
        "                             shuffle=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ziL9lq2bjHcj"
      },
      "source": [
        "**Model Architecture**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P3UxQhmF8p8L"
      },
      "source": [
        "The following is a brief  **summary** of the **architectural features** of the model:-\n",
        "\n",
        "1. The model is based on **encoder-decoder** architecture and uses **PReLU** activation throught the network. It hepls us to achieve **faster convergence** and **improved accuracy**.\n",
        "\n",
        "2. The inputs are initially **downsampled** from a size of 512 to 128 (i,e 1/4'th). This helps us to **reduce** the overall **computation** costs; while preseving the details.\n",
        "\n",
        "3. It uses **skip connections** between the encoder and decoder blocks (like unet) and helps us to extract **fine details** and improves **gradient flow** across layers. \n",
        "\n",
        "4. Further, it uses **bottleneck** layers (like resnet) with **depthwise** convolutions for **faster inference**.\n",
        "\n",
        "5. Also, it uses **dilated** convolution(like deeplab) and helps us to maintain **larger receptive field** with **same computation and memory costs**, while also **preserving resolution**.\n",
        "\n",
        "6. Finally, the features are **upsampled** to full resolution(512) with the help of **transposed convolutions**.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iQYfH9TYXkQY"
      },
      "source": [
        "\n",
        "def bilinear_resize(x, rsize):\n",
        "  return tf.image.resize_bilinear(x, [rsize,rsize], align_corners=True)\n",
        "\n",
        "def encode_bottleneck(x, proj_ch, out_ch, strides=1, dilation=1,separable=True, depthwise=True, preluop=False, pool=False):\n",
        "\n",
        "  x = PReLU(shared_axes=[1, 2])(x)\n",
        "  y = Conv2D(filters=proj_ch, kernel_size=strides, strides=strides, padding='same')(x)\n",
        "  y = PReLU(shared_axes=[1, 2])(y)\n",
        "  \n",
        "  if separable==True:\n",
        "          \n",
        "      if depthwise==True:\n",
        "          y = SeparableConv2D(filters=proj_ch, kernel_size=3, strides=1, padding='same')(y)\n",
        "          y = PReLU(shared_axes=[1, 2])(y)\n",
        "          y = DepthwiseConv2D(kernel_size=3, padding='same')(y)\n",
        "      else:\n",
        "        y= SeparableConv2D(filters=proj_ch, kernel_size=5, strides=1, padding='same')(y)\n",
        "  else:  \n",
        "      y = Conv2D(filters=out_ch, kernel_size=3, dilation_rate= dilation ,strides=1, padding='same')(y)\n",
        "\n",
        "  y = PReLU(shared_axes=[1, 2])(y)\n",
        "  y = Conv2D(filters=out_ch, kernel_size=1, strides=1, padding='same')(y)\n",
        "\n",
        "  if pool == True:\n",
        "      m = MaxPool2D((2, 2), padding='same')(x)\n",
        "      if m.shape[-1] != 128:\n",
        "          x = Conv2D(filters=out_ch, kernel_size=1, strides=1, padding='same')(m)\n",
        "      else:\n",
        "          x = m\n",
        "      z = Add()([x, y])\n",
        "      return z, m \n",
        "\n",
        "  z = Add()([x, y])\n",
        "  \n",
        "  if preluop==True:\n",
        "    return z, x\n",
        "\n",
        "  return z\n",
        "\n",
        "\n",
        "def decode_bottleneck(x,res1, res2 ,proj_ch1, out_ch, proj_ch2, strides=1, rsize=32,pconv=True):\n",
        "\n",
        "  x = PReLU(shared_axes=[1, 2])(x)\n",
        "  y = Conv2D(filters=proj_ch1, kernel_size=strides, strides=strides, padding='same')(x)\n",
        "  y = PReLU(shared_axes=[1, 2])(y)\n",
        "  \n",
        "  y = Conv2DTranspose(filters=8, kernel_size=3, strides=2, padding = 'same' )(y)\n",
        "\n",
        "  y = PReLU(shared_axes=[1, 2])(y)\n",
        "  y = Conv2D(filters=out_ch, kernel_size=1, strides=1, padding='same')(y)\n",
        "\n",
        "  \n",
        "  x = Conv2D(filters=out_ch, kernel_size=1, strides=1, padding='same')(x)\n",
        "  r = Add()([x,res1])\n",
        "  x = Lambda(lambda r: bilinear_resize(r, rsize))(r)\n",
        "\n",
        "  z = Add()([x, y])\n",
        "  z = PReLU(shared_axes=[1, 2])(z)\n",
        "  if pconv == True:\n",
        "    z = concatenate([z,res2])\n",
        "  b = Conv2D(filters=proj_ch2, kernel_size=strides, strides=strides, padding='same')(z)\n",
        "  b = PReLU(shared_axes=[1, 2])(b)\n",
        "  b = Conv2D(filters=proj_ch2, kernel_size=3, strides=1, padding='same')(b)\n",
        "  b = PReLU(shared_axes=[1, 2])(b)\n",
        "  b = Conv2D(filters=out_ch, kernel_size=1, strides=1, padding='same')(b)\n",
        "  \n",
        "  if pconv == True:\n",
        "    z = Conv2D(filters=out_ch, kernel_size=1, strides=1, padding='same')(z)\n",
        "\n",
        "  c = Add()([z, b])\n",
        "\n",
        "  return c\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S-ptHgWiCc4X"
      },
      "source": [
        "**Note:** Here, we need to ensure that the **bilinear resize** has the option **align_corners=True** for proper upsampling of image and for avoiding the shifting problem (TF 1.15 Inference). Also, for **reducing** the number of **parameters**, set the **PReLU** option **shared_axes=[1,2]** to  share the parameters along axes 1 and 2."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eBKtiiwU2yiq"
      },
      "source": [
        "# Define the network using the basic bottleneck layers\n",
        "def slim_net():\n",
        "\n",
        "  # Initial spatial phase [Reduces input by a factor 1/4]\n",
        "  input = Input(shape=(512,512,3), name='ip')\n",
        "  x = Conv2D(filters=8, kernel_size=2, strides=2, padding='valid')(input)\n",
        "  x = PReLU(shared_axes=[1, 2])(x)\n",
        "  x = Conv2D(filters=32, kernel_size=2, strides=2, padding='valid')(x)\n",
        "\n",
        "  b1, r1 = encode_bottleneck(x, proj_ch=16, out_ch=64, strides=2, separable=True, depthwise=True, pool=True)\n",
        "  b2, p1 = encode_bottleneck(b1, proj_ch=16, out_ch=64, strides=1, separable=True, depthwise=True, preluop=True, pool=False)\n",
        "  b3 = encode_bottleneck(b2, proj_ch=16, out_ch=64, strides=1, separable=True, depthwise=True, pool=False)\n",
        "  b4, r2 = encode_bottleneck(b3, proj_ch=32, out_ch=128, strides=2, separable=True, depthwise=True, pool=True)\n",
        "\n",
        "  b5, p2 = encode_bottleneck(b4, proj_ch=16, out_ch=128, strides=1, separable=True, depthwise=True, preluop=True,pool=False)\n",
        "  b6 = encode_bottleneck(b5, proj_ch=16, out_ch=128, strides=1, separable=True, depthwise=True, pool=False)\n",
        "  b7 = encode_bottleneck(b6, proj_ch=16, out_ch=128, strides=1, separable=True, depthwise=True, pool=False)\n",
        "  b8 = encode_bottleneck(b7, proj_ch=16, out_ch=128, strides=1, separable=True, depthwise=True, pool=False)\n",
        "\n",
        "  b9, r3 = encode_bottleneck(b8, proj_ch=16, out_ch=128, strides=2, separable=True, depthwise=True, pool=True)\n",
        "  b10 = encode_bottleneck(b9, proj_ch=8, out_ch=128, strides=1, separable=True, depthwise=True, pool=False)\n",
        "  b11 = encode_bottleneck(b10, proj_ch=8, out_ch=128, strides=1, dilation=2, separable=False, depthwise=False, pool=False) # dil -2\n",
        "  b12 = encode_bottleneck(b11, proj_ch=8, out_ch=128, strides=1, separable=True, depthwise=False, pool=False)\n",
        "  b13 = encode_bottleneck(b12, proj_ch=8, out_ch=128, strides=1, dilation=4, separable=False, depthwise=False, pool=False) # dil -4\n",
        "  b14 = encode_bottleneck(b13, proj_ch=8, out_ch=128, strides=1, separable=True, depthwise=True, pool=False)\n",
        "  b15 = encode_bottleneck(b14, proj_ch=8, out_ch=128, strides=1, dilation=8, separable=False, depthwise=False, pool=False) # dil -8 \n",
        "  b16 = encode_bottleneck(b15, proj_ch=8, out_ch=128, strides=1, separable=True, depthwise=True, pool=False)\n",
        "  b17 = encode_bottleneck(b16, proj_ch=8, out_ch=128, strides=1, dilation=2, separable=False, depthwise=False, pool=False) # dil -2\n",
        "  b18 = encode_bottleneck(b17, proj_ch=8, out_ch=128, strides=1, separable=True, depthwise=False, pool=False)\n",
        "  b19 = encode_bottleneck(b18, proj_ch=8, out_ch=128, strides=1, dilation=4, separable=False, depthwise=False, pool=False) # dil -4\n",
        "  b20 = encode_bottleneck(b19, proj_ch=8, out_ch=128, strides=1, separable=True, depthwise=True, pool=False)\n",
        "  b21 = encode_bottleneck(b20, proj_ch=8, out_ch=128, strides=1, dilation=8, separable=False, depthwise=False, pool=False) # dil -8 \n",
        "\n",
        "  b22 = encode_bottleneck(b21, proj_ch=4, out_ch=128, strides=1, separable=False, depthwise=False, pool=False) # dil -1\n",
        "\n",
        " \n",
        "  d1 = decode_bottleneck(b22,res1=r3, res2=p2 ,proj_ch1=8, proj_ch2=8, out_ch=128, strides=1, rsize=32, pconv=True)\n",
        "  d2 = decode_bottleneck(d1,res1=r2, res2=p1 ,proj_ch1=8, proj_ch2=4 , out_ch=64, strides=1, rsize=64,pconv=True)\n",
        "  d3 = decode_bottleneck(d2,res1=r1, res2=None ,proj_ch1=4, proj_ch2=4, out_ch=32, strides=1, rsize=128,pconv=False)\n",
        "\n",
        "  pout1 = PReLU(shared_axes=[1, 2])(d3)\n",
        "  cout1 = Conv2DTranspose(filters=8, kernel_size=2, strides=2, padding = 'same' )(pout1) # output size: 256\n",
        "  pout2 = PReLU(shared_axes=[1, 2])(cout1)\n",
        "  cout2 = Conv2DTranspose(filters=2, kernel_size=2, strides=2, padding = 'same' )(pout2) # output size: 512\n",
        "\n",
        "  model = Model(inputs=input, outputs=cout2)\n",
        "  model.compile(optimizer='adam',\n",
        "              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "              metrics=['accuracy']) # Ensure you have sparse labels\n",
        "\n",
        "  return model\n",
        "\n",
        "# Initialize the model and plot summary\n",
        "model = slim_net()\n",
        "model.summary()\n",
        "\n",
        "# Plot model architecture\n",
        "plot_model(model, to_file='slim-net.png')\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MUHzczMljQos"
      },
      "source": [
        "**Training and Callbacks**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "05xAC5oD7iYl"
      },
      "source": [
        "Configure the **callbacks** for learning rate decay, logging and checkpoint."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3bklc47ykSKE"
      },
      "source": [
        "# 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=10, 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": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hWnhUK1J742x"
      },
      "source": [
        "Perform **training and validation** on the model, using the dataset and save the results."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Kb4kmSi_n7Fi"
      },
      "source": [
        "# Train the model\n",
        "model_history = model.fit(train_dataset, epochs=epochs,\n",
        "                          steps_per_epoch=num_train//batch_sz,\n",
        "                          validation_steps=num_val//batch_sz,\n",
        "                          validation_data=val_dataset,\n",
        "                          callbacks=callbacks_list)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0rtWTAOlLh_9"
      },
      "source": [
        "**Testing**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WtV3-3Tnx_PG"
      },
      "source": [
        "Load the model and **test images**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A_36W2Dh7oiE"
      },
      "source": [
        "# Load the final model for inference\n",
        "model=load_model('/content/drive/My Drive/slim512/ckpt/slim-net-157-0.02.hdf5',compile=False)\n",
        "\n",
        "# Prepare the batch of images as lists\n",
        "TEST_DIR='/content/test/'\n",
        "images=os.listdir(TEST_DIR)\n",
        "inputs=[]\n",
        "for img in images:\n",
        "    im=Image.open(TEST_DIR+img)\n",
        "    im=im.resize((512,512), Image.ANTIALIAS)\n",
        "    inputs.append(np.array(im))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k_S9vN0CyNAR"
      },
      "source": [
        "Perform **inferene** on input image batch"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pxk9V_jJALZe"
      },
      "source": [
        "# Perform batch prediction and obtain masks\n",
        "batch=np.float32(inputs)/255.0\n",
        "result=model.predict(batch)\n",
        "argmax=np.argmax(result, axis=3)\n",
        "outputs=list(argmax[...,np.newaxis]*batch)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Sp91M2ZK0BEL"
      },
      "source": [
        "Plot the **input and output** images using matplotlib"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FbUY4w7oBSr3"
      },
      "source": [
        "# Combine the input & output list of images\n",
        "plot_lists=inputs+outputs\n",
        "\n",
        "# Plot all the images using matplotlib \n",
        "fig=plt.figure(figsize=(16, 8))\n",
        "columns = 4\n",
        "rows = 2\n",
        "\n",
        "# Show all four inputs and corresponding outputs\n",
        "for i in range(1, columns*rows+1):\n",
        "    img = plot_lists[i-1].squeeze()\n",
        "    fig.add_subplot(rows, columns, i)\n",
        "    plt.imshow(img)\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qQygALBgr7Re"
      },
      "source": [
        "**Exporting**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xv4g0IjVXfL0"
      },
      "source": [
        "Load the trained **model chekpoint** for export."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "n0AkImsR_Gh-"
      },
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "def bilinear_resize(x, rsize):\n",
        "  return tf.image.resize_bilinear(x, [rsize,rsize], align_corners=True)\n",
        "\n",
        "def slice_foreground(x):\n",
        "    return tf.strided_slice(x, [0,0, 0, 1], [1,512, 512, 2], [1, 1, 1, 1])\n",
        "\n",
        "model=load_model('/content/drive/My Drive/slim512/ckpt/slim-net-157-0.02.hdf5',compile=False)\n",
        "model.summary()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ax7F6qc4XiuD"
      },
      "source": [
        "Add new layers and **modify** the network."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gXFz5k4-RNMe"
      },
      "source": [
        "sm=tf.keras.layers.Softmax()(model.output) # softmax\n",
        "str_slice=Lambda(slice_foreground, name=\"strided_slice\")(sm) # strided slice\n",
        "newout=Reshape((262144,))(str_slice) # reshape\n",
        "reshape_model=Model(model.input,newout)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1EvdWZM5YANF"
      },
      "source": [
        "Save the final **keras model** for deployment."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8fEY6L82XY7l"
      },
      "source": [
        "reshape_model.summary()\n",
        "reshape_model.save('/content/slim_reshape_v2.h5')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_QGmEURwYck7"
      },
      "source": [
        "Convert the final keras model to **tflite** format."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "g6qvHgN_EkeR"
      },
      "source": [
        "converter = tf.lite.TFLiteConverter.from_keras_model_file('/content/slim_reshape_v2.h5')\n",
        "tflite_model = converter.convert()\n",
        "open(\"slim_reshape_v2.tflite\", \"wb\").write(tflite_model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HDPfMq7teOBE"
      },
      "source": [
        "**Verification**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rzZ15TIkeqN3"
      },
      "source": [
        "Load the exported slim-net **tflite model** for portait segmentation."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "T6GoRU7_epWD"
      },
      "source": [
        "!wget -O slim_reshape_v2.tflite https://github.com/anilsathyan7/Portrait-Segmentation/blob/master/models/slim_seg_512/slim_reshape%20v2.tflite?raw=true"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GfqsWeMne-Za"
      },
      "source": [
        "Download a sample **portrait image** for testing the model."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MiJx2LEKd-d-"
      },
      "source": [
        "import cv2\n",
        "import numpy as np\n",
        "from skimage import io\n",
        "import tensorflow as tf\n",
        "from matplotlib import pyplot as plt\n",
        "\n",
        "image=io.imread('https://images.newindianexpress.com/uploads/user/imagelibrary/2020/6/9/w600X390/Javed_Akhtar_PTI.jpg')\n",
        "image=cv2.resize(image,(512,512))\n",
        "plt.imshow(image)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xSvlXB1Jfcq7"
      },
      "source": [
        "Perform inference on test images, using **tflite interpreter**."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "byw23hggd1n3"
      },
      "source": [
        "def run_tflite_model(tflite_file, test_image):\n",
        "\n",
        "  # Initialize the interpreter\n",
        "  interpreter = tf.lite.Interpreter(model_path=str(tflite_file))\n",
        "  interpreter.allocate_tensors()\n",
        "\n",
        "  # Get input and output details\n",
        "  input_details = interpreter.get_input_details()[0]\n",
        "  output_details = interpreter.get_output_details()[0]\n",
        "\n",
        "  # Preprocess the input image\n",
        "  test_image = test_image/255.0\n",
        "  test_image = np.expand_dims(test_image, axis=0).astype(input_details[\"dtype\"])\n",
        "\n",
        "  # Run the interpreter and get the output\n",
        "  interpreter.set_tensor(input_details[\"index\"], test_image)\n",
        "  interpreter.invoke()\n",
        "  output = interpreter.get_tensor(output_details[\"index\"])[0]\n",
        "\n",
        "  # Compute mask from segmentaion output\n",
        "  mask = np.reshape(output, (512,512))>0.5\n",
        "\n",
        "  return mask"
      ],
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9_yUqTV_frie"
      },
      "source": [
        "**Crop** the input using output mask and plot the results."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zIj_WenLfaE3"
      },
      "source": [
        "import cv2\n",
        "import numpy as np\n",
        "from skimage import io\n",
        "import tensorflow as tf\n",
        "from matplotlib import pyplot as plt\n",
        "\n",
        "mask=run_tflite_model(tflite_file='/content/slim_reshape_v2.tflite',test_image=image)\n",
        "crop_float=image*mask[...,np.newaxis]\n",
        "plt.imshow(crop_float/255.0)\n",
        "plt.title('Float Model Output')\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}