U_Net_Image_Segmentation.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "U_Net_Image_Segmentation.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python378jvsc74a57bd0e7a6a2a0faef9b3dca8eee8d2f0177f3278be1cb967e7fc435687a2394b4cb7a",
      "display_name": "Python 3.7.8 64-bit"
    },
    "metadata": {
      "interpreter": {
        "hash": "e7a6a2a0faef9b3dca8eee8d2f0177f3278be1cb967e7fc435687a2394b4cb7a"
      }
    }
  },
  "cells": [
    {
      "source": [
        "This U-Net Implementation Borrowed heavily from Keras Docs Tutorial https://keras.io/examples/vision/oxford_pets_image_segmentation/"
      ],
      "cell_type": "markdown",
      "metadata": {}
    },
    {
      "source": [
        "load datasets"
      ],
      "cell_type": "markdown",
      "metadata": {}
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import os\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "ROOT_DIR = os.getcwd()\n",
        "BUOY_DIR = os.path.join(ROOT_DIR, \"buoy_data\")\n",
        "input_dir = os.path.join(BUOY_DIR, \"Train/img\")\n",
        "target_dir = os.path.join(BUOY_DIR, \"Train/masks_machine\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0rVj7Hb1lnzx"
      },
      "source": [
        "**Prepare paths of input images and target segmentation masks**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wIx33Dx9A15e",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 262
        },
        "outputId": "73e8be31-4402-4a2a-bdbd-15360181475b"
      },
      "source": [
        "img_size = (128, 128)\n",
        "\n",
        "num_classes = 1\n",
        "batch_size = 32\n",
        "\n",
        "input_img_paths = sorted(\n",
        "    [\n",
        "        os.path.join(input_dir, fname)\n",
        "        for fname in os.listdir(input_dir)\n",
        "        if fname.endswith(\".jpg\")\n",
        "    ]\n",
        ")\n",
        "target_img_paths = sorted(\n",
        "    [\n",
        "        os.path.join(target_dir, fname)\n",
        "        for fname in os.listdir(target_dir)\n",
        "        if fname.endswith(\".png\") and not fname.startswith(\".\")\n",
        "    ]\n",
        ")\n",
        "\n",
        "print(\"Number of samples:\", len(input_img_paths))\n",
        "\n",
        "for input_path, target_path in zip(input_img_paths[:10], target_img_paths[:10]):\n",
        "    print(input_path, \"|\", target_path)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "osSb9PfxltE0"
      },
      "source": [
        "**Display an input image and a segmentation mask**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dR7hymxMBIjW",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 229
        },
        "outputId": "266d4f23-738f-458a-e4f1-4d5bfe107aba"
      },
      "source": [
        "from IPython.display import Image, display\n",
        "from tensorflow.keras.preprocessing.image import load_img, img_to_array\n",
        "import PIL\n",
        "from PIL import ImageOps\n",
        "import numpy as np\n",
        "\n",
        "\n",
        "# Display input image #7\n",
        "display(Image(filename=input_img_paths[9]))\n",
        "\n",
        "# Display auto-contrast version of corresponding target (per-pixel categories)\n",
        "img = PIL.ImageOps.autocontrast(load_img(target_img_paths[9]))\n",
        "display(img)\n",
        "print(np.unique(load_img(target_img_paths[9])))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VWDXozytl9jk"
      },
      "source": [
        "**Prepare Sequence class to load & vectorize batches of data**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WKz0t9_7BJbE"
      },
      "source": [
        "from tensorflow import keras\n",
        "import numpy as np\n",
        "from tensorflow.keras.preprocessing.image import load_img\n",
        "\n",
        "\n",
        "class Buoys(keras.utils.Sequence):\n",
        "    \"\"\"Helper to iterate over the data (as Numpy arrays).\"\"\"\n",
        "\n",
        "    def __init__(self, batch_size, img_size, input_img_paths, target_img_paths):\n",
        "        self.batch_size = batch_size\n",
        "        self.img_size = img_size\n",
        "        self.input_img_paths = input_img_paths\n",
        "        self.target_img_paths = target_img_paths\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.target_img_paths) // self.batch_size\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        \"\"\"Returns tuple (input, target) correspond to batch #idx.\"\"\"\n",
        "        i = idx * self.batch_size\n",
        "        batch_input_img_paths = self.input_img_paths[i : i + self.batch_size]\n",
        "        batch_target_img_paths = self.target_img_paths[i : i + self.batch_size]\n",
        "        x = np.zeros((batch_size,) + self.img_size + (3,), dtype=\"float32\")\n",
        "        for j, path in enumerate(batch_input_img_paths):\n",
        "            img = load_img(path, target_size=self.img_size)\n",
        "            x[j] = img\n",
        "        y = np.zeros((batch_size,) + self.img_size + (1,), dtype=\"uint8\")\n",
        "        for j, path in enumerate(batch_target_img_paths):\n",
        "            img = load_img(path, target_size=self.img_size, color_mode=\"grayscale\")\n",
        "            y[j] = np.expand_dims(img, 2)\n",
        "        return x, y"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aGOa1pRamEdM"
      },
      "source": [
        "**U-Net blocks**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0DVFHu-KBNyc"
      },
      "source": [
        "def down_block(x, filters, kernel_size=(3, 3), padding=\"same\", strides=1):\n",
        "    c = keras.layers.Conv2D(filters, kernel_size, padding=padding, strides=strides, activation=\"relu\")(x)\n",
        "    c = keras.layers.Conv2D(filters, kernel_size, padding=padding, strides=strides, activation=\"relu\")(c)\n",
        "    p = keras.layers.MaxPool2D((2, 2), (2, 2))(c)\n",
        "    return c, p\n",
        "\n",
        "def up_block(x, skip, filters, kernel_size=(3, 3), padding=\"same\", strides=1):\n",
        "    us = keras.layers.UpSampling2D((2, 2))(x)\n",
        "    concat = keras.layers.Concatenate()([us, skip])\n",
        "    c = keras.layers.Conv2D(filters, kernel_size, padding=padding, strides=strides, activation=\"relu\")(concat)\n",
        "    c = keras.layers.Conv2D(filters, kernel_size, padding=padding, strides=strides, activation=\"relu\")(c)\n",
        "    return c\n",
        "\n",
        "def bottleneck(x, filters, kernel_size=(3, 3), padding=\"same\", strides=1):\n",
        "    c = keras.layers.Conv2D(filters, kernel_size, padding=padding, strides=strides, activation=\"relu\")(x)\n",
        "    c = keras.layers.Conv2D(filters, kernel_size, padding=padding, strides=strides, activation=\"relu\")(c)\n",
        "    return c"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "26yPBe9CmIqp"
      },
      "source": [
        "**U-Net architecture**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CCp2qUDzBYon"
      },
      "source": [
        "def UNet():\n",
        "    f = [16, 32, 64, 128, 256]\n",
        "    inputs = keras.layers.Input((img_size[0], img_size[1], 3))\n",
        "    \n",
        "    p0 = inputs\n",
        "    c1, p1 = down_block(p0, f[0]) #128 -> 64\n",
        "    c2, p2 = down_block(p1, f[1]) #64 -> 32\n",
        "    c3, p3 = down_block(p2, f[2]) #32 -> 16\n",
        "    c4, p4 = down_block(p3, f[3]) #16->8\n",
        "    \n",
        "    bn = bottleneck(p4, f[4])\n",
        "    \n",
        "    u1 = up_block(bn, c4, f[3]) #8 -> 16\n",
        "    u2 = up_block(u1, c3, f[2]) #16 -> 32\n",
        "    u3 = up_block(u2, c2, f[1]) #32 -> 64\n",
        "    u4 = up_block(u3, c1, f[0]) #64 -> 128\n",
        "    \n",
        "    outputs = keras.layers.Conv2D(3, (1, 1), padding=\"same\", activation=\"softmax\")(u4)\n",
        "    model = keras.models.Model(inputs, outputs)\n",
        "    return model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SLa9lMdGBc5M",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "e2e5fc97-9f6a-46d6-f662-6109a1d2c15f"
      },
      "source": [
        "model = UNet()\n",
        "model.summary()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WhAvtDtGmZYb"
      },
      "source": [
        "**Set aside a validation split**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SAL16bPrC5oO"
      },
      "source": [
        "import random\n",
        "# Split our img paths into a training and a validation set\n",
        "val_samples = 1000\n",
        "random.Random(1337).shuffle(input_img_paths)\n",
        "random.Random(1337).shuffle(target_img_paths)\n",
        "train_input_img_paths = input_img_paths[:-val_samples]\n",
        "train_target_img_paths = target_img_paths[:-val_samples]\n",
        "val_input_img_paths = input_img_paths[-val_samples:]\n",
        "val_target_img_paths = target_img_paths[-val_samples:]\n",
        "\n",
        "# Instantiate data Sequences for each split\n",
        "train_gen = Buoys(\n",
        "    batch_size, img_size, train_input_img_paths, train_target_img_paths\n",
        ")\n",
        "val_gen = Buoys(batch_size, img_size, val_input_img_paths, val_target_img_paths)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ef2teaCemf0D"
      },
      "source": [
        "**Train the model**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2xaL0LlwC869",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "1e8f86a0-9b64-4894-86f3-9f3938ab341f"
      },
      "source": [
        "# Configure the model for training.\n",
        "# We use the \"sparse\" version of categorical_crossentropy\n",
        "# because our target data is integers.\n",
        "model.compile(optimizer=\"rmsprop\", loss=\"sparse_categorical_crossentropy\")\n",
        "\n",
        "callbacks = [\n",
        "    keras.callbacks.ModelCheckpoint(\"buoy_detection.h5\", save_best_only=True)\n",
        "]\n",
        "\n",
        "# Train the model, doing validation at the end of each epoch.\n",
        "epochs = 10\n",
        "history = model.fit(train_gen, epochs=epochs, validation_data=val_gen, callbacks=callbacks)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "source": [
        "Optionally Load the Model if model is already trained, don't want to waste time/resources constantly retraining"
      ],
      "cell_type": "markdown",
      "metadata": {}
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "model = keras.models.load_model(\"buoy_detection.h5\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qjUE-kRommLC"
      },
      "source": [
        "**Visualize predictions**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "L7Ew0UV5DA5F",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 697
        },
        "outputId": "05779f5b-7a14-4152-a886-4d98a9d45df2"
      },
      "source": [
        "# Generate predictions for all images in the validation set\n",
        "\n",
        "val_gen = Buoys(batch_size, img_size, val_input_img_paths, val_target_img_paths)\n",
        "val_preds = model.predict(val_gen)\n",
        "\n",
        "def display_mask(i):\n",
        "    \"\"\"Quick utility to display a model's prediction.\"\"\"\n",
        "    mask = np.argmax(val_preds[i], axis=-1)\n",
        "    mask = np.expand_dims(mask, axis=-1)\n",
        "    img = PIL.ImageOps.autocontrast(keras.preprocessing.image.array_to_img(mask))\n",
        "    display(img)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "for i in range(10):\n",
        "    display_mask(i)"
      ]
    },
    {
      "source": [
        "plot loss and accuracy"
      ],
      "cell_type": "markdown",
      "metadata": {}
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "plt.plot(history.history[\"loss\"])\n",
        "plt.plot(history.history[\"val_loss\"])\n",
        "plt.title(\"UNet Model Loss\")\n",
        "plt.legend(['train','val'], loc='upper left')\n",
        "plt.show()"
      ]
    },
    {
      "source": [
        "Compute Dice"
      ],
      "cell_type": "markdown",
      "metadata": {}
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from keras import backend as K\n",
        "from PIL import Image, ImageOps\n",
        "import tensorflow as tf\n",
        "\n",
        "def dice_coef(y_true, y_pred):\n",
        "    y_true_f = K.flatten(y_true)\n",
        "    y_pred_f = K.flatten(y_pred)\n",
        "    intersection = K.sum(y_true_f * y_pred_f)\n",
        "    return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)\n",
        "d = 0\n",
        "c= 0\n",
        "for i in range(len(val_preds)):\n",
        "    truth = Image.open(target_img_paths[i])\n",
        "    truth = np.asarray(truth, dtype=np.uint8)\n",
        "    truth = tf.convert_to_tensor(truth, dtype=np.uint8)\n",
        "    # truth = img_to_array(load_img(target_img_paths[i]))\n",
        "    pred = val_preds[i]\n",
        "    # truth = truth.resize((128,128))\n",
        "    print(pred.shape)\n",
        "    print(truth.shape)\n",
        "    d += dice_coef(truth, pred)\n",
        "    c +=1\n",
        "print(1-(d/c))\n"
      ]
    },
    {
      "source": [
        "Compute meanIoU"
      ],
      "cell_type": "markdown",
      "metadata": {}
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "from keras import backend as K\n",
        "m = tf.keras.metrics.MeanIoU(num_classes=1)\n",
        "\n",
        "for i in range(4):\n",
        "    truth = Image.open(target_img_paths[i])\n",
        "    truth = truth.resize((128,128))\n",
        "    truth = np.asarray(truth, dtype=np.uint8)\n",
        "    # truth = tf.convert_to_tensor(truth, dtype=np.uint8)\n",
        "    # truth = img_to_array(load_img(target_img_paths[i]))\n",
        "    pred = val_preds[i]\n",
        "    y_true_f = K.flatten(truth)\n",
        "    y_pred_f = K.flatten(pred)\n",
        "    m.update_state(truth,pred)\n",
        "\n"
      ]
    },
    {
      "source": [
        "Show Masks"
      ],
      "cell_type": "markdown",
      "metadata": {}
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def display_masks(im, mask):\n",
        "    plt.figure()\n",
        "    plt.imshow(im, 'gray', interpolation='none')\n",
        "\n",
        "    plt.show()\n",
        "display_masks(load_img(input_img_paths[0]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": []
    }
  ]
}