{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyOA6q4xuqXb0BaHFu0tHC2h"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "#Pet Image Segementation Using Modified U-Nets built on the Oxford IIIT Pets Dataset"
      ],
      "metadata": {
        "id": "JiICHTic5xBd"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Image Segementation is the process on taking an image as an input and indivually labling if each pixel is part of the object, bording the object, or is not part of the object. The Oxford IIIT Pets Dataset is perfect for this. It is a 37 category database with around 200 images in each category. We can further increase this with data augmentation. All of the images have a corresponding mask which has all the pixels divided into 3 classes: on the pet, bordering the pet, or outside the pet. Using this, we can train a modified U-Net to predict these masks when faced with new images. This model has acheived a 92% accuracy on the validation data, which is very high considering that we can further improve this with more epoches and model tuning."
      ],
      "metadata": {
        "id": "kokOKokS7GY8"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "We must first begin by importing the nessesary libaries into our program. We will be using TensorFlow and Keras to build and train the model, MatPlotLib to show our images and masks, and TensorFlow Datasets to access our dataset."
      ],
      "metadata": {
        "id": "AmBgUS1G96ve"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "##Preparing Our Data To Be Processed"
      ],
      "metadata": {
        "id": "CZdG-a5yBBlR"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install git+https://github.com/tensorflow/examples.git\n",
        "\n",
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "from tensorflow_examples.models.pix2pix import pix2pix\n",
        "import matplotlib.pyplot as plt\n"
      ],
      "metadata": {
        "id": "jERRrqR--f1G"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now that we have imported all of our libraries, we can start to load the dataset and get ready for the data to be processed."
      ],
      "metadata": {
        "id": "D6BjANjW_aFS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "dataset, info = tfds.load('oxford_iiit_pet:3.*.*', with_info=True)\n",
        "\n",
        "TRAIN_LENGTH = info.splits['train'].num_examples\n",
        "BATCH_SIZE = 64\n",
        "BUFFER_SIZE = 1000\n",
        "STEPS_PER_EPOCH = TRAIN_LENGTH // BATCH_SIZE"
      ],
      "metadata": {
        "id": "dkemyqs7_q-q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Next, we will make a load_image function that will resize the image and the mask to (128, 128). We will also make a function to normizlize the image, which will reduce the value range of the image from (0, 255) to (0, 1)."
      ],
      "metadata": {
        "id": "jeN_rXKz_5Qr"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def load_image(datapoint):\n",
        "  input_image = tf.image.resize(datapoint['image'], (128, 128))\n",
        "  input_mask = tf.image.resize(datapoint['segmentation_mask'], (128, 128), method = tf.image.ResizeMethod.NEAREST_NEIGHBOR)\n",
        "  input_image, input_mask = normalize_image(input_image, input_mask)\n",
        "  return input_image, input_mask\n",
        "\n",
        "def normalize_image(input_image, input_mask):\n",
        "  input_image = tf.cast(input_image, tf.float32) / 255.0\n",
        "  input_mask -= 1\n",
        "  return input_image, input_mask"
      ],
      "metadata": {
        "id": "l5NFCx12AqDx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We will also now need to add our data augmentation. This will be impletemented in the form of a class and will randomly flip our images and masks using a set seed."
      ],
      "metadata": {
        "id": "K2en1ZvABmMP"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class Augment(tf.keras.layers.Layer):\n",
        "  def __init__(self, seed=42):\n",
        "    super().__init__()\n",
        "    self.augment_inputs = tf.keras.layers.RandomFlip(mode=\"horizontal\", seed=seed)\n",
        "    self.augment_labels = tf.keras.layers.RandomFlip(mode=\"horizontal\", seed=seed)\n",
        "\n",
        "  def call(self, inputs, labels):\n",
        "    inputs = self.augment_inputs(inputs)\n",
        "    labels = self.augment_labels(labels)\n",
        "    return inputs, labels"
      ],
      "metadata": {
        "id": "XBVo-uRzB8Vc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now we can group our images into batches to be processed."
      ],
      "metadata": {
        "id": "7YBY0UvfA5El"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "train_images = dataset['train'].map(load_image, num_parallel_calls=tf.data.AUTOTUNE)\n",
        "test_images = dataset['test'].map(load_image, num_parallel_calls=tf.data.AUTOTUNE)\n",
        "\n",
        "train_batches = (train_images.cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE).repeat().map(Augment()).prefetch(buffer_size=tf.data.AUTOTUNE))\n",
        "test_batches = test_images.batch(BATCH_SIZE)"
      ],
      "metadata": {
        "id": "fAvQHyxABNH3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "##Creating and Training the Model"
      ],
      "metadata": {
        "id": "GIcff6fzCSKR"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "First we will begin by creating our base model. It will be based on the pretrained MobileNetV2 model. We will also create the second model, which will be based on pix2pix.\n"
      ],
      "metadata": {
        "id": "GyY1EnvXCgKY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "base_model = tf.keras.applications.MobileNetV2(input_shape=[128, 128, 3], include_top=False)\n",
        "\n",
        "layer_names = [\n",
        "    'block_1_expand_relu',\n",
        "    'block_3_expand_relu',\n",
        "    'block_6_expand_relu',\n",
        "    'block_13_expand_relu',\n",
        "    'block_16_project',\n",
        "]\n",
        "base_model_outputs = [base_model.get_layer(name).output for name in layer_names]\n",
        "\n",
        "down_stack = tf.keras.Model(inputs=base_model.input, outputs=base_model_outputs)\n",
        "\n",
        "down_stack.trainable = False\n",
        "\n",
        "up_stack = [pix2pix.upsample(512, 3), pix2pix.upsample(256, 3), pix2pix.upsample(128, 3), pix2pix.upsample(64, 3),]\n"
      ],
      "metadata": {
        "id": "KdEkW6IRC0tf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now we can define the function that will make our U-Net Model."
      ],
      "metadata": {
        "id": "udgKbAzTDclx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def U_net_model(output_channels:int, down_stack, up_stack):\n",
        "  inputs = tf.keras.layers.Input(shape=[128, 128, 3])\n",
        "  skips = down_stack(inputs)\n",
        "  outputs = skips[-1]\n",
        "  skips = reversed(skips[:-1])\n",
        "\n",
        "  for up, skip in zip(up_stack, skips):\n",
        "    outputs = up(outputs)\n",
        "    concatenate = tf.keras.layers.Concatenate()\n",
        "    outputs = concatenate([outputs, skip])\n",
        "\n",
        "  last = tf.keras.layers.Conv2DTranspose(filters=output_channels, kernel_size=3, strides=2, padding='same')\n",
        "  outputs = last(outputs)\n",
        "  return tf.keras.Model(inputs=inputs, outputs=outputs)"
      ],
      "metadata": {
        "id": "bfkx81pCDjCe"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now we can finally create and compile the model."
      ],
      "metadata": {
        "id": "0oSXkajXDsQk"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "OUTPUT_CLASSES = 3\n",
        "\n",
        "model = U_net_model(OUTPUT_CLASSES, down_stack, up_stack)\n",
        "model.compile(optimizer='adam', loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), metrics=['accuracy'])"
      ],
      "metadata": {
        "id": "922GC6FhD6cw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "It is now time to train the model we have just built using the train and test batches."
      ],
      "metadata": {
        "id": "0VJaMpmvD9k7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "EPOCHS = 20\n",
        "VAL_SUBSPLITS = 5\n",
        "VALIDATION_STEPS = info.splits['test'].num_examples//BATCH_SIZE//VAL_SUBSPLITS\n",
        "\n",
        "model.fit(train_batches, epochs=EPOCHS, steps_per_epoch=STEPS_PER_EPOCH, validation_steps=VALIDATION_STEPS, validation_data=test_batches)"
      ],
      "metadata": {
        "id": "He629X5qE4j5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now we can save the model in three different formats for later use. We will be saving it in the legacy .h5 format, the new .keras format, and the TensorFlow SavedModel."
      ],
      "metadata": {
        "id": "5jV9jQwyFKeb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model.save(\"pets.h5\")\n",
        "model.save(\"pets.keras\")\n",
        "model.save(\"model/dogs\")"
      ],
      "metadata": {
        "id": "4PCorPKVG-4_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "##Predict Using the Model"
      ],
      "metadata": {
        "id": "3YSb402_NTwm"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now we will use our model to predict a mask against new data. But first, we must load the model off the disk."
      ],
      "metadata": {
        "id": "xmb_EUTQNjEh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def attempt_load(i):\n",
        "  try:\n",
        "    model = tf.keras.models.load_model('pets'+names[i])\n",
        "    return model\n",
        "  except:\n",
        "    attempt_load(i+1)\n",
        "\n",
        "names = ['.keras', '', '.h5']\n",
        "\n",
        "model = attempt_load(0)"
      ],
      "metadata": {
        "id": "jXiX1ghnN2py"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now we need to create our predicted mask."
      ],
      "metadata": {
        "id": "bdb37QuKRyfG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def create_mask(pred_mask):\n",
        "  pred_mask = tf.math.argmax(pred_mask, axis=-1)\n",
        "  pred_mask = pred_mask[..., tf.newaxis]\n",
        "  return pred_mask[0]"
      ],
      "metadata": {
        "id": "7QUSGR3YSPna"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We also need a function that can display the image and the predicted mask."
      ],
      "metadata": {
        "id": "ogVrMFQFSS6t"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def display(display_list):\n",
        "  plt.figure(figsize=(15, 15))\n",
        "  titles = ['Input Image', 'Predicted Mask']\n",
        "  for i in range(len(display_list)):\n",
        "    plt.subplot(1, len(display_list), i+1)\n",
        "    plt.title(titles[i])\n",
        "    plt.imshow(tf.keras.utils.array_to_img(display_list[i]))\n",
        "    plt.axis('off')\n",
        "  plt.show()"
      ],
      "metadata": {
        "id": "fyAdwN8KStsb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Finally, we need a way to get the users' picture and use our model to predict the mask for that picture."
      ],
      "metadata": {
        "id": "ZSoCGr2hSvJN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def show_predictions(image_url, model):\n",
        "  image = tf.keras.utils.get_file(origin=image_url)\n",
        "  image = tf.keras.utils.load_img(image)\n",
        "  image = tf.keras.utils.img_to_array(image)\n",
        "  image = tf.image.resize(image, (128,128))\n",
        "  image = tf.cast(image, tf.float32) / 255.0\n",
        "  image = tf.expand_dims(image, axis=0)\n",
        "  pred_mask = model.predict(image)\n",
        "  display([image[0], create_mask(pred_mask)])"
      ],
      "metadata": {
        "id": "QrApNNhCeRYK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "And now lets wrap it all together!"
      ],
      "metadata": {
        "id": "eyqUFmFfeUT_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "while True:\n",
        "  url = input(\"Please enter an image url:\")\n",
        "  try:\n",
        "    image = tf.keras.utils.get_file(origin=url)\n",
        "    image = tf.keras.utils.load_img(image)\n",
        "    break\n",
        "  except:\n",
        "    print(\"That is not a valid link\")\n",
        "\n",
        "show_predictions(url, model)"
      ],
      "metadata": {
        "id": "is94T8gBeYle"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}