{ "cells": [ { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "ULKtInm7dAMK" }, "source": [ "# Image Segmentation with `tf.keras`" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "7Plun_k1dAML" }, "source": [ "
\n", "\n", " Run in Google Colab \n", "\n", "View source on GitHub
" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "cl79rk4KKol8" }, "source": [ "In this tutorial we will learn how to segment images. **Segmentation** is the process of generating pixel-wise segmentations giving the class of the object visible at each pixel. For example, we could be identifying the location and boundaries of people within an image or identifying cell nuclei from an image. Formally, image segmentation refers to the process of partitioning an image into a set of pixels that we desire to identify (our target) and the background. \n", "\n", "Specifically, in this tutorial we will be using the [Kaggle Carvana Image Masking Challenge Dataset](https://www.kaggle.com/c/carvana-image-masking-challenge). \n", "\n", "This dataset contains a large number of car images, with each car taken from different angles. In addition, for each car image, we have an associated manually cutout mask; our task will be to automatically create these cutout masks for unseen data. \n", "\n", "## Specific concepts that will be covered:\n", "In the process, we will build practical experience and develop intuition around the following concepts:\n", "* **[Functional API](https://keras.io/getting-started/functional-api-guide/)** - we will be implementing UNet, a convolutional network model classically used for biomedical image segmentation with the Functional API. \n", " * This model has layers that require multiple input/outputs. This requires the use of the functional API\n", " * Check out the original [paper](https://arxiv.org/abs/1505.04597), \n", "U-Net: Convolutional Networks for Biomedical Image Segmentation by Olaf Ronneberger!\n", "* **Custom Loss Functions and Metrics** - We'll implement a custom loss function using binary [**cross entropy**](https://developers.google.com/machine-learning/glossary/#cross-entropy) and **dice loss**. We'll also implement **dice coefficient** (which is used for our loss) and **mean intersection over union**, that will help us monitor our training process and judge how well we are performing. \n", "* **Saving and loading keras models** - We'll save our best model to disk. When we want to perform inference/evaluate our model, we'll load in the model from disk. \n", "\n", "### We will follow the general workflow:\n", "1. Visualize data/perform some exploratory data analysis\n", "2. Set up data pipeline and preprocessing\n", "3. Build model\n", "4. Train model\n", "5. Evaluate model\n", "6. Repeat\n", "\n", "**Audience:** This post is geared towards intermediate users who are comfortable with basic machine learning concepts.\n", "Note that if you wish to run this notebook, it is highly recommended that you do so with a GPU. \n", "\n", "**Time Estimated**: 60 min\n", "\n", "By: Raymond Yuan, Software Engineering Intern" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "bJcQiA3OdCY6" }, "outputs": [], "source": [ "!pip install kaggle" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "ODNLPGHKKgr-" }, "outputs": [], "source": [ "import os\n", "import glob\n", "import zipfile\n", "import functools\n", "\n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", "import matplotlib as mpl\n", "mpl.rcParams['axes.grid'] = False\n", "mpl.rcParams['figure.figsize'] = (12,12)\n", "\n", "from sklearn.model_selection import train_test_split\n", "import matplotlib.image as mpimg\n", "import pandas as pd\n", "from PIL import Image\n" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "YQ9VRReUQxXi" }, "outputs": [], "source": [ "import tensorflow as tf\n", "import tensorflow.contrib as tfcontrib\n", "from tensorflow.python.keras import layers\n", "from tensorflow.python.keras import losses\n", "from tensorflow.python.keras import models\n", "from tensorflow.python.keras import backend as K " ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "RW9gk331S0KA" }, "source": [ "# Get all the files \n", "Since this tutorial will be using a dataset from Kaggle, it requires [creating an API Token](https://github.com/Kaggle/kaggle-api#api-credentials) for your Kaggle account, and uploading it. " ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "sAVM1ZTmdAMR" }, "outputs": [], "source": [ "import os\n", "\n", "# Upload the API token.\n", "def get_kaggle_credentials():\n", " token_dir = os.path.join(os.path.expanduser(\"~\"),\".kaggle\")\n", " token_file = os.path.join(token_dir, \"kaggle.json\")\n", " if not os.path.isdir(token_dir):\n", " os.mkdir(token_dir)\n", " try:\n", " with open(token_file,'r') as f:\n", " pass\n", " except IOError as no_file:\n", " try:\n", " from google.colab import files\n", " except ImportError:\n", " raise no_file\n", " \n", " uploaded = files.upload()\n", " \n", " if \"kaggle.json\" not in uploaded:\n", " raise ValueError(\"You need an API key! see: \"\n", " \"https://github.com/Kaggle/kaggle-api#api-credentials\")\n", " with open(token_file, \"wb\") as f:\n", " f.write(uploaded[\"kaggle.json\"])\n", " os.chmod(token_file, 600)\n", "\n", "get_kaggle_credentials()\n" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "gh6jkMp8dN5B" }, "source": [ "Only import kaggle after adding the credentials." ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "EoWJ1hb9dOV_" }, "outputs": [], "source": [ "import kaggle" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "wC-byMdadAMT" }, "source": [ "### We'll download the data from Kaggle\n", "Caution, large download ahead - downloading all files will require 14GB of diskspace. " ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "6MOTOyU3dAMU" }, "outputs": [], "source": [ "competition_name = 'carvana-image-masking-challenge'" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "3gJSCmWjdAMW" }, "outputs": [], "source": [ "# Download data from Kaggle and unzip the files of interest. \n", "def load_data_from_zip(competition, file):\n", " with zipfile.ZipFile(os.path.join(competition, file), \"r\") as zip_ref:\n", " unzipped_file = zip_ref.namelist()[0]\n", " zip_ref.extractall(competition)\n", "\n", "def get_data(competition):\n", " kaggle.api.competition_download_files(competition, competition)\n", " load_data_from_zip(competition, 'train.zip')\n", " load_data_from_zip(competition, 'train_masks.zip')\n", " load_data_from_zip(competition, 'train_masks.csv.zip')\n", " \n" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "l5SZJKPRdXNX" }, "source": [ "You must [accept the competition rules](https://www.kaggle.com/c/carvana-image-masking-challenge/rules) before downloading the data." ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "_SsQjuN2dWmU" }, "outputs": [], "source": [ "get_data(competition_name)" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "wT1kb3q0ghhi" }, "outputs": [], "source": [ "img_dir = os.path.join(competition_name, \"train\")\n", "label_dir = os.path.join(competition_name, \"train_masks\")" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "9ej-e6cqmRgd" }, "outputs": [], "source": [ "df_train = pd.read_csv(os.path.join(competition_name, 'train_masks.csv'))\n", "ids_train = df_train['img'].map(lambda s: s.split('.')[0])" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "33i4xFXweztH" }, "outputs": [], "source": [ "x_train_filenames = []\n", "y_train_filenames = []\n", "for img_id in ids_train:\n", " x_train_filenames.append(os.path.join(img_dir, \"{}.jpg\".format(img_id)))\n", " y_train_filenames.append(os.path.join(label_dir, \"{}_mask.gif\".format(img_id)))" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "DtutNudKbf70" }, "outputs": [], "source": [ "x_train_filenames, x_val_filenames, y_train_filenames, y_val_filenames = \\\n", " train_test_split(x_train_filenames, y_train_filenames, test_size=0.2, random_state=42)" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "zDycQekHaMqq" }, "outputs": [], "source": [ "num_train_examples = len(x_train_filenames)\n", "num_val_examples = len(x_val_filenames)\n", "\n", "print(\"Number of training examples: {}\".format(num_train_examples))\n", "print(\"Number of validation examples: {}\".format(num_val_examples))" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "Nhda5fkPS3JD" }, "source": [ "### Here's what the paths look like " ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "Di1N83ArilzR" }, "outputs": [], "source": [ "x_train_filenames[:10]" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "Gc-BDv1Zio1z" }, "outputs": [], "source": [ "y_train_filenames[:10]" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "mhvDoZkbcUa1" }, "source": [ "# Visualize\n", "Let's take a look at some of the examples of different images in our dataset. " ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "qUA6SDLhozjj" }, "outputs": [], "source": [ "display_num = 5\n", "\n", "r_choices = np.random.choice(num_train_examples, display_num)\n", "\n", "plt.figure(figsize=(10, 15))\n", "for i in range(0, display_num * 2, 2):\n", " img_num = r_choices[i // 2]\n", " x_pathname = x_train_filenames[img_num]\n", " y_pathname = y_train_filenames[img_num]\n", " \n", " plt.subplot(display_num, 2, i + 1)\n", " plt.imshow(mpimg.imread(x_pathname))\n", " plt.title(\"Original Image\")\n", " \n", " example_labels = Image.open(y_pathname)\n", " label_vals = np.unique(example_labels)\n", " \n", " plt.subplot(display_num, 2, i + 2)\n", " plt.imshow(example_labels)\n", " plt.title(\"Masked Image\") \n", " \n", "plt.suptitle(\"Examples of Images and their Masks\")\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "d4CPgvPiToB_" }, "source": [ "# Set up " ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "HfeMRgyoa2n6" }, "source": [ "Let’s begin by setting up some parameters. We’ll standardize and resize all the shapes of the images. We’ll also set up some training parameters: " ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "oeDoiSFlothe" }, "outputs": [], "source": [ "img_shape = (256, 256, 3)\n", "batch_size = 3\n", "epochs = 5" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "8_d5ATP21npW" }, "source": [ "Using these exact same parameters may be too computationally intensive for your hardware, so tweak the parameters accordingly. Also, it is important to note that due to the architecture of our UNet version, the size of the image must be evenly divisible by a factor of 32, as we down sample the spatial resolution by a factor of 2 with each `MaxPooling2Dlayer`.\n", "\n", "\n", "If your machine can support it, you will achieve better performance using a higher resolution input image (e.g. 512 by 512) as this will allow more precise localization and less loss of information during encoding. In addition, you can also make the model deeper.\n", "\n", "\n", "Alternatively, if your machine cannot support it, lower the image resolution and/or batch size. Note that lowering the image resolution will decrease performance and lowering batch size will increase training time.\n" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "_HONB9JbXxDM" }, "source": [ "# Build our input pipeline with `tf.data`\n", "Since we begin with filenames, we will need to build a robust and scalable data pipeline that will play nicely with our model. If you are unfamiliar with **tf.data** you should check out my other tutorial introducing the concept! \n", "\n", "### Our input pipeline will consist of the following steps:\n", "1. Read the bytes of the file in from the filename - for both the image and the label. Recall that our labels are actually images with each pixel annotated as car or background (1, 0). \n", "2. Decode the bytes into an image format\n", "3. Apply image transformations: (optional, according to input parameters)\n", " * `resize` - Resize our images to a standard size (as determined by eda or computation/memory restrictions)\n", " * The reason why this is optional is that U-Net is a fully convolutional network (e.g. with no fully connected units) and is thus not dependent on the input size. However, if you choose to not resize the images, you must use a batch size of 1, since you cannot batch variable image size together\n", " * Alternatively, you could also bucket your images together and resize them per mini-batch to avoid resizing images as much, as resizing may affect your performance through interpolation, etc.\n", " * `hue_delta` - Adjusts the hue of an RGB image by a random factor. This is only applied to the actual image (not our label image). The `hue_delta` must be in the interval `[0, 0.5]` \n", " * `horizontal_flip` - flip the image horizontally along the central axis with a 0.5 probability. This transformation must be applied to both the label and the actual image. \n", " * `width_shift_range` and `height_shift_range` are ranges (as a fraction of total width or height) within which to randomly translate the image either horizontally or vertically. This transformation must be applied to both the label and the actual image. \n", " * `rescale` - rescale the image by a certain factor, e.g. 1/ 255.\n", "4. Shuffle the data, repeat the data (so we can iterate over it multiple times across epochs), batch the data, then prefetch a batch (for efficiency).\n", "\n", "It is important to note that these transformations that occur in your data pipeline must be symbolic transformations. " ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "EtRA8vILbx2_" }, "source": [ "#### Why do we do these image transformations?\n", "This is known as **data augmentation**. Data augmentation \"increases\" the amount of training data by augmenting them via a number of random transformations. During training time, our model would never see twice the exact same picture. This helps prevent [overfitting](https://developers.google.com/machine-learning/glossary/#overfitting) and helps the model generalize better to unseen data." ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "3aGi28u8Cq9M" }, "source": [ "## Processing each pathname" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "Fb_psznAggwr" }, "outputs": [], "source": [ "def _process_pathnames(fname, label_path):\n", " # We map this function onto each pathname pair \n", " img_str = tf.read_file(fname)\n", " img = tf.image.decode_jpeg(img_str, channels=3)\n", "\n", " label_img_str = tf.read_file(label_path)\n", " # These are gif images so they return as (num_frames, h, w, c)\n", " label_img = tf.image.decode_gif(label_img_str)[0]\n", " # The label image should only have values of 1 or 0, indicating pixel wise\n", " # object (car) or not (background). We take the first channel only. \n", " label_img = label_img[:, :, 0]\n", " label_img = tf.expand_dims(label_img, axis=-1)\n", " return img, label_img" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "Y4UE28JiCuOk" }, "source": [ "## Shifting the image" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "xdY046OqtGVH" }, "outputs": [], "source": [ "def shift_img(output_img, label_img, width_shift_range, height_shift_range):\n", " \"\"\"This fn will perform the horizontal or vertical shift\"\"\"\n", " if width_shift_range or height_shift_range:\n", " if width_shift_range:\n", " width_shift_range = tf.random_uniform([], \n", " -width_shift_range * img_shape[1],\n", " width_shift_range * img_shape[1])\n", " if height_shift_range:\n", " height_shift_range = tf.random_uniform([],\n", " -height_shift_range * img_shape[0],\n", " height_shift_range * img_shape[0])\n", " # Translate both \n", " output_img = tfcontrib.image.translate(output_img,\n", " [width_shift_range, height_shift_range])\n", " label_img = tfcontrib.image.translate(label_img,\n", " [width_shift_range, height_shift_range])\n", " return output_img, label_img" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "qY253aZfCwd2" }, "source": [ "## Flipping the image randomly " ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "OogLSplstur9" }, "outputs": [], "source": [ "def flip_img(horizontal_flip, tr_img, label_img):\n", " if horizontal_flip:\n", " flip_prob = tf.random_uniform([], 0.0, 1.0)\n", " tr_img, label_img = tf.cond(tf.less(flip_prob, 0.5),\n", " lambda: (tf.image.flip_left_right(tr_img), tf.image.flip_left_right(label_img)),\n", " lambda: (tr_img, label_img))\n", " return tr_img, label_img" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "_YIJLIr5Cyyr" }, "source": [ "## Assembling our transformations into our augment function" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "18WA0Sl3olyn" }, "outputs": [], "source": [ "def _augment(img,\n", " label_img,\n", " resize=None, # Resize the image to some size e.g. [256, 256]\n", " scale=1, # Scale image e.g. 1 / 255.\n", " hue_delta=0, # Adjust the hue of an RGB image by random factor\n", " horizontal_flip=False, # Random left right flip,\n", " width_shift_range=0, # Randomly translate the image horizontally\n", " height_shift_range=0): # Randomly translate the image vertically \n", " if resize is not None:\n", " # Resize both images\n", " label_img = tf.image.resize_images(label_img, resize)\n", " img = tf.image.resize_images(img, resize)\n", " \n", " if hue_delta:\n", " img = tf.image.random_hue(img, hue_delta)\n", " \n", " img, label_img = flip_img(horizontal_flip, img, label_img)\n", " img, label_img = shift_img(img, label_img, width_shift_range, height_shift_range)\n", " label_img = tf.to_float(label_img) * scale\n", " img = tf.to_float(img) * scale \n", " return img, label_img" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "tkNqQaR2HQbd" }, "outputs": [], "source": [ "def get_baseline_dataset(filenames, \n", " labels,\n", " preproc_fn=functools.partial(_augment),\n", " threads=5, \n", " batch_size=batch_size,\n", " shuffle=True): \n", " num_x = len(filenames)\n", " # Create a dataset from the filenames and labels\n", " dataset = tf.data.Dataset.from_tensor_slices((filenames, labels))\n", " # Map our preprocessing function to every element in our dataset, taking\n", " # advantage of multithreading\n", " dataset = dataset.map(_process_pathnames, num_parallel_calls=threads)\n", " if preproc_fn.keywords is not None and 'resize' not in preproc_fn.keywords:\n", " assert batch_size == 1, \"Batching images must be of the same size\"\n", "\n", " dataset = dataset.map(preproc_fn, num_parallel_calls=threads)\n", " \n", " if shuffle:\n", " dataset = dataset.shuffle(num_x)\n", " \n", " \n", " # It's necessary to repeat our data for all epochs \n", " dataset = dataset.repeat().batch(batch_size)\n", " return dataset" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "zwtgius5CRKc" }, "source": [ "## Set up train and validation datasets\n", "Note that we apply image augmentation to our training dataset but not our validation dataset. " ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "iu5WmYmOwKrV" }, "outputs": [], "source": [ "tr_cfg = {\n", " 'resize': [img_shape[0], img_shape[1]],\n", " 'scale': 1 / 255.,\n", " 'hue_delta': 0.1,\n", " 'horizontal_flip': True,\n", " 'width_shift_range': 0.1,\n", " 'height_shift_range': 0.1\n", "}\n", "tr_preprocessing_fn = functools.partial(_augment, **tr_cfg)" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "RtzLkDFMpF0T" }, "outputs": [], "source": [ "val_cfg = {\n", " 'resize': [img_shape[0], img_shape[1]],\n", " 'scale': 1 / 255.,\n", "}\n", "val_preprocessing_fn = functools.partial(_augment, **val_cfg)" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "5cNpECdkaafo" }, "outputs": [], "source": [ "train_ds = get_baseline_dataset(x_train_filenames,\n", " y_train_filenames,\n", " preproc_fn=tr_preprocessing_fn,\n", " batch_size=batch_size)\n", "val_ds = get_baseline_dataset(x_val_filenames,\n", " y_val_filenames, \n", " preproc_fn=val_preprocessing_fn,\n", " batch_size=batch_size)" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "Yasuvr5IbFlM" }, "source": [ "## Let's see if our image augmentor data pipeline is producing expected results" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "hjoUqbPdHQej" }, "outputs": [], "source": [ "temp_ds = get_baseline_dataset(x_train_filenames, \n", " y_train_filenames,\n", " preproc_fn=tr_preprocessing_fn,\n", " batch_size=1,\n", " shuffle=False)\n", "# Let's examine some of these augmented images\n", "data_aug_iter = temp_ds.make_one_shot_iterator()\n", "next_element = data_aug_iter.get_next()\n", "with tf.Session() as sess: \n", " batch_of_imgs, label = sess.run(next_element)\n", "\n", " # Running next element in our graph will produce a batch of images\n", " plt.figure(figsize=(10, 10))\n", " img = batch_of_imgs[0]\n", "\n", " plt.subplot(1, 2, 1)\n", " plt.imshow(img)\n", "\n", " plt.subplot(1, 2, 2)\n", " plt.imshow(label[0, :, :, 0])\n", " plt.show()" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "fvtxCncKsoRd" }, "source": [ "# Build the model\n", "We'll build the U-Net model. U-Net is especially good with segmentation tasks because it can localize well to provide high resolution segmentation masks. In addition, it works well with small datasets and is relatively robust against overfitting as the training data is in terms of the number of patches within an image, which is much larger than the number of training images itself. Unlike the original model, we will add batch normalization to each of our blocks. \n", "\n", "The Unet is built with an encoder portion and a decoder portion. The encoder portion is composed of a linear stack of [`Conv`](https://developers.google.com/machine-learning/glossary/#convolution), `BatchNorm`, and [`Relu`](https://developers.google.com/machine-learning/glossary/#ReLU) operations followed by a [`MaxPool`](https://developers.google.com/machine-learning/glossary/#pooling). Each `MaxPool` will reduce the spatial resolution of our feature map by a factor of 2. We keep track of the outputs of each block as we feed these high resolution feature maps with the decoder portion. The Decoder portion is comprised of UpSampling2D, Conv, BatchNorm, and Relus. Note that we concatenate the feature map of the same size on the decoder side. Finally, we add a final Conv operation that performs a convolution along the channels for each individual pixel (kernel size of (1, 1)) that outputs our final segmentation mask in grayscale. \n", "## The Keras Functional API\n", "The Keras functional API is used when you have multi-input/output models, shared layers, etc. It's a powerful API that allows you to manipulate tensors and build complex graphs with intertwined datastreams easily. In addition it makes **layers** and **models** both callable on tensors. \n", " * To see more examples check out the [get started guide](https://keras.io/getting-started/functional-api-guide/). \n", " \n", " \n", " We'll build these helper functions that will allow us to ensemble our model block operations easily and simply. " ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "zfew1i1F6bK-" }, "outputs": [], "source": [ "def conv_block(input_tensor, num_filters):\n", " encoder = layers.Conv2D(num_filters, (3, 3), padding='same')(input_tensor)\n", " encoder = layers.BatchNormalization()(encoder)\n", " encoder = layers.Activation('relu')(encoder)\n", " encoder = layers.Conv2D(num_filters, (3, 3), padding='same')(encoder)\n", " encoder = layers.BatchNormalization()(encoder)\n", " encoder = layers.Activation('relu')(encoder)\n", " return encoder\n", "\n", "def encoder_block(input_tensor, num_filters):\n", " encoder = conv_block(input_tensor, num_filters)\n", " encoder_pool = layers.MaxPooling2D((2, 2), strides=(2, 2))(encoder)\n", " \n", " return encoder_pool, encoder\n", "\n", "def decoder_block(input_tensor, concat_tensor, num_filters):\n", " decoder = layers.Conv2DTranspose(num_filters, (2, 2), strides=(2, 2), padding='same')(input_tensor)\n", " decoder = layers.concatenate([concat_tensor, decoder], axis=-1)\n", " decoder = layers.BatchNormalization()(decoder)\n", " decoder = layers.Activation('relu')(decoder)\n", " decoder = layers.Conv2D(num_filters, (3, 3), padding='same')(decoder)\n", " decoder = layers.BatchNormalization()(decoder)\n", " decoder = layers.Activation('relu')(decoder)\n", " decoder = layers.Conv2D(num_filters, (3, 3), padding='same')(decoder)\n", " decoder = layers.BatchNormalization()(decoder)\n", " decoder = layers.Activation('relu')(decoder)\n", " return decoder" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "xRLp21S_hpTn" }, "outputs": [], "source": [ "inputs = layers.Input(shape=img_shape)\n", "# 256\n", "\n", "encoder0_pool, encoder0 = encoder_block(inputs, 32)\n", "# 128\n", "\n", "encoder1_pool, encoder1 = encoder_block(encoder0_pool, 64)\n", "# 64\n", "\n", "encoder2_pool, encoder2 = encoder_block(encoder1_pool, 128)\n", "# 32\n", "\n", "encoder3_pool, encoder3 = encoder_block(encoder2_pool, 256)\n", "# 16\n", "\n", "encoder4_pool, encoder4 = encoder_block(encoder3_pool, 512)\n", "# 8\n", "\n", "center = conv_block(encoder4_pool, 1024)\n", "# center\n", "\n", "decoder4 = decoder_block(center, encoder4, 512)\n", "# 16\n", "\n", "decoder3 = decoder_block(decoder4, encoder3, 256)\n", "# 32\n", "\n", "decoder2 = decoder_block(decoder3, encoder2, 128)\n", "# 64\n", "\n", "decoder1 = decoder_block(decoder2, encoder1, 64)\n", "# 128\n", "\n", "decoder0 = decoder_block(decoder1, encoder0, 32)\n", "# 256\n", "\n", "outputs = layers.Conv2D(1, (1, 1), activation='sigmoid')(decoder0)" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "luDqDqu8c1AX" }, "source": [ "## Define your model\n", "Using functional API, you must define your model by specifying the inputs and outputs associated with the model. " ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "76QkTzXVczgc" }, "outputs": [], "source": [ "model = models.Model(inputs=[inputs], outputs=[outputs])" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "p0tNnmyOdtyr" }, "source": [ "# Defining custom metrics and loss functions\n", "Defining loss and metric functions are simple with Keras. Simply define a function that takes both the True labels for a given example and the Predicted labels for the same given example. " ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "sfuBVut0fogM" }, "source": [ "Dice loss is a metric that measures overlap. More info on optimizing for Dice coefficient (our dice loss) can be found in the [paper](http://campar.in.tum.de/pub/milletari2016Vnet/milletari2016Vnet.pdf), where it was introduced. \n", "\n", "We use dice loss here because it performs better at class imbalanced problems by design. In addition, maximizing the dice coefficient and IoU metrics are the actual objectives and goals of our segmentation task. Using cross entropy is more of a proxy which is easier to maximize. Instead, we maximize our objective directly. " ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "t_8_hbHECUAW" }, "outputs": [], "source": [ "def dice_coeff(y_true, y_pred):\n", " smooth = 1.\n", " # Flatten\n", " y_true_f = tf.reshape(y_true, [-1])\n", " y_pred_f = tf.reshape(y_pred, [-1])\n", " intersection = tf.reduce_sum(y_true_f * y_pred_f)\n", " score = (2. * intersection + smooth) / (tf.reduce_sum(y_true_f) + tf.reduce_sum(y_pred_f) + smooth)\n", " return score" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "4DgINhlpNaxP" }, "outputs": [], "source": [ "def dice_loss(y_true, y_pred):\n", " loss = 1 - dice_coeff(y_true, y_pred)\n", " return loss" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "qqClGNFJdANU" }, "source": [ "Here, we'll use a specialized loss function that combines binary cross entropy and our dice loss. This is based on [individuals who competed within this competition obtaining better results empirically](https://www.kaggle.com/c/carvana-image-masking-challenge/discussion/40199). Try out your own custom losses to measure performance (e.g. bce + log(dice_loss), only bce, etc.)!" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "udrfi9JGB-bL" }, "outputs": [], "source": [ "def bce_dice_loss(y_true, y_pred):\n", " loss = losses.binary_crossentropy(y_true, y_pred) + dice_loss(y_true, y_pred)\n", " return loss" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "LifmpjXNc9Gz" }, "source": [ "## Compile your model\n", "We use our custom loss function to minimize. In addition, we specify what metrics we want to keep track of as we train. Note that metrics are not actually used during the training process to tune the parameters, but are instead used to measure performance of the training process. " ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "gflcWk2Cc8Bi" }, "outputs": [], "source": [ "model.compile(optimizer='adam', loss=bce_dice_loss, metrics=[dice_loss])\n", "\n", "model.summary()" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "8WG_8iZ_dMbK" }, "source": [ "## Train your model\n", "Training your model with `tf.data` involves simply providing the model's `fit` function with your training/validation dataset, the number of steps, and epochs. \n", "\n", "We also include a Model callback, [`ModelCheckpoint`](https://keras.io/callbacks/#modelcheckpoint) that will save the model to disk after each epoch. We configure it such that it only saves our highest performing model. Note that saving the model capture more than just the weights of the model: by default, it saves the model architecture, weights, as well as information about the training process such as the state of the optimizer, etc." ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "1nHnj6199elZ" }, "outputs": [], "source": [ "save_model_path = '/tmp/weights.hdf5'\n", "cp = tf.keras.callbacks.ModelCheckpoint(filepath=save_model_path, monitor='val_dice_loss', save_best_only=True, verbose=1)" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "vJP_EvuTb4hH" }, "source": [ "Don't forget to specify our model callback in the `fit` function call. " ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "UMZcOrq5aaj1" }, "outputs": [], "source": [ "history = model.fit(train_ds, \n", " steps_per_epoch=int(np.ceil(num_train_examples / float(batch_size))),\n", " epochs=epochs,\n", " validation_data=val_ds,\n", " validation_steps=int(np.ceil(num_val_examples / float(batch_size))),\n", " callbacks=[cp])" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "gCAUsoxfTTrh" }, "source": [ "# Visualize training process" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "AvntxymYn8rM" }, "outputs": [], "source": [ "dice = history.history['dice_loss']\n", "val_dice = history.history['val_dice_loss']\n", "\n", "loss = history.history['loss']\n", "val_loss = history.history['val_loss']\n", "\n", "epochs_range = range(epochs)\n", "\n", "plt.figure(figsize=(16, 8))\n", "plt.subplot(1, 2, 1)\n", "plt.plot(epochs_range, dice, label='Training Dice Loss')\n", "plt.plot(epochs_range, val_dice, label='Validation Dice Loss')\n", "plt.legend(loc='upper right')\n", "plt.title('Training and Validation Dice Loss')\n", "\n", "plt.subplot(1, 2, 2)\n", "plt.plot(epochs_range, loss, label='Training Loss')\n", "plt.plot(epochs_range, val_loss, label='Validation Loss')\n", "plt.legend(loc='upper right')\n", "plt.title('Training and Validation Loss')\n", "\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "dWPhb87GdhkG" }, "source": [ "Even with only 5 epochs, we see strong performance." ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "MGFKf8yCTYbw" }, "source": [ "# Visualize actual performance \n", "We'll visualize our performance on the validation set.\n", "\n", "Note that in an actual setting (competition, deployment, etc.) we'd evaluate on the test set with the full image resolution. " ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "oIddsUcM_KeI" }, "source": [ "To load our model we have two options:\n", "1. Since our model architecture is already in memory, we can simply call `load_weights(save_model_path)`\n", "2. If you wanted to load the model from scratch (in a different setting without already having the model architecture in memory) we simply call \n", "\n", "```model = models.load_model(save_model_path, custom_objects={'bce_dice_loss': bce_dice_loss, 'dice_loss': dice_loss})```, specificing the necessary custom objects, loss and metrics, that we used to train our model. \n", "\n", "If you want to see more examples, check our the [keras guide](https://keras.io/getting-started/faq/#how-can-i-save-a-keras-model)!" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "5Ph7acmrCXm6" }, "outputs": [], "source": [ "# Alternatively, load the weights directly: model.load_weights(save_model_path)\n", "model = models.load_model(save_model_path, custom_objects={'bce_dice_loss': bce_dice_loss,\n", " 'dice_loss': dice_loss})" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab": {}, "colab_type": "code", "id": "0GnwZ7CPaamI" }, "outputs": [], "source": [ "# Let's visualize some of the outputs \n", "data_aug_iter = val_ds.make_one_shot_iterator()\n", "next_element = data_aug_iter.get_next()\n", "\n", "# Running next element in our graph will produce a batch of images\n", "plt.figure(figsize=(10, 20))\n", "for i in range(5):\n", " batch_of_imgs, label = tf.keras.backend.get_session().run(next_element)\n", " img = batch_of_imgs[0]\n", " predicted_label = model.predict(batch_of_imgs)[0]\n", "\n", " plt.subplot(5, 3, 3 * i + 1)\n", " plt.imshow(img)\n", " plt.title(\"Input image\")\n", " \n", " plt.subplot(5, 3, 3 * i + 2)\n", " plt.imshow(label[0, :, :, 0])\n", " plt.title(\"Actual Mask\")\n", " plt.subplot(5, 3, 3 * i + 3)\n", " plt.imshow(predicted_label[:, :, 0])\n", " plt.title(\"Predicted Mask\")\n", "plt.suptitle(\"Examples of Input Image, Label, and Prediction\")\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "iPV7RMA9TjPC" }, "source": [ "# Key Takeaways\n", "In this tutorial we learned how to train a network to automatically detect and create cutouts of cars from images! \n", "\n", "## Specific concepts that will we covered:\n", "In the process, we hopefully built some practical experience and developed intuition around the following concepts\n", "* [**Functional API**](https://keras.io/getting-started/functional-api-guide/) - we implemented UNet with the Functional API. Functional API gives a lego-like API that allows us to build pretty much any network. \n", "* **Custom Losses and Metrics** - We implemented custom metrics that allow us to see exactly what we need during training time. In addition, we wrote a custom loss function that is specifically suited to our task. \n", "* **Save and load our model** - We saved our best model that we encountered according to our specified metric. When we wanted to perform inference with out best model, we loaded it from disk. Note that saving the model capture more than just the weights of the model: by default, it saves the model architecture, weights, as well as information about the training process such as the state of the optimizer, etc. " ] } ], "metadata": { "accelerator": "GPU", "colab": { "collapsed_sections": [], "name": "Image Segmentation", "private_outputs": true, "provenance": [], "version": "0.3.2" }, "kernelspec": { "display_name": "Python [default]", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.6.4" } }, "nbformat": 4, "nbformat_minor": 1 }