Title: U-Net Image Segmentation in Keras - PyImageSearch

URL Source: https://pyimagesearch.com/2022/02/21/u-net-image-segmentation-in-keras/

Published Time: 2022-02-21T14:00:00+00:00

Markdown Content:
In this tutorial, you will learn how to create U-Net, an image segmentation model in TensorFlow 2 / Keras. We will first present a brief introduction on image segmentation, U-Net architecture, and then walk through the code implementation with a Colab notebook.

![Image 1](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%201024%20512'%3E%3C/svg%3E)

**To learn how to implement a U-Net with TensorFlow 2 / Keras, _just keep reading._**

![Image 2](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20128%20128'%3E%3C/svg%3E)

#### Looking for the source code to this post?

[Jump Right To The Downloads Section](#download-the-code)

Image segmentation is a computer vision task that segments an image into multiple areas by assigning a label to every pixel of the image. It provides much more information about an image than object detection, which draws a bounding box around the detected object, or image classification, which assigns a label to the object.

Segmentation is useful and can be used in real-world applications such as medical imaging, clothes segmentation, flooding maps, self-driving cars, etc.

There are two types of image segmentation:

*   **Semantic segmentation:** classify each pixel with a label.
*   **Instance segmentation:** classify each pixel and differentiate each object instance.

U-Net is a semantic segmentation technique originally proposed for medical imaging segmentation. It’s one of the earlier deep learning segmentation models, and the U-Net architecture is also used in many GAN variants such as the Pix2Pix generator.

### **U-Net Architecture**

U-Net was introduced in the paper, [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://arxiv.org/abs/1505.04597). The model architecture is fairly simple: an encoder (for downsampling) and a decoder (for upsampling) with skip connections. As **Figure 1** shows, it shapes like the letter U hence the name U-Net.

![Image 3](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20700%20389'%3E%3C/svg%3E)

**Figure 1:** U-Net architecture (image source: [U-Net paper](https://arxiv.org/abs/1505.04597)).

The gray arrows indicate the skip connections that concatenate the encoder feature map with the decoder, which helps the backward flow of gradients for improved training.

Now that we have a basic understanding of semantic segmentation and the U-Net architecture, let’s implement a U-Net with TensorFlow 2 / Keras. Please follow the tutorial below with this [Colab notebook](https://colab.research.google.com/github/margaretmz/image-segmentation/blob/main/unet_pet_segmentation.ipynb).

### **Setup**

We will be using Colab for model training, so make sure you set “Hardware accelerator” to “GPU under Runtime / change runtime type.” Then import the libraries and packages this project depends on:

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import tensorflow\_datasets as tfds
import matplotlib.pyplot as plt
import numpy as np

### **Dataset**

We will use the [Oxford-IIIT pet dataset](https://www.robots.ox.ac.uk/~vgg/data/pets/), available as part of the [TensorFlow Datasets](https://www.tensorflow.org/datasets/catalog/oxford_iiit_pet) (TFDS). It can be easily loaded with TFDS, and then with a bit of data preprocessing, ready for training segmentation models.

With just one line of code, we can use `tfds` to load the dataset by specifying the name of the dataset, and get the dataset info by setting `with_info=True`:

dataset, info = tfds.load('oxford\_iiit\_pet:3.\*.\*', with\_info=True)

Print the dataset info with `print(info)`, and we will see all kinds of detailed information about the Oxford pet dataset. For example, in **Figure 2**, we can see there are a total of 7349 images with a built-in test/train split.

![Image 4](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20700%20455'%3E%3C/svg%3E)

**Figure 2:** Oxford-IIIT pet dataset info (image by the author).

Let’s first make a few changes to the downloaded data before we start training U-Net with it.

First, we need to resize the images and masks to `128x128`:

def resize(input\_image, input\_mask):
   input\_image = tf.image.resize(input\_image, (128, 128), method="nearest")
   input\_mask = tf.image.resize(input\_mask, (128, 128), method="nearest")

   return input\_image, input\_mask

We then create a function to augment the dataset by flipping them horizontally:

def augment(input\_image, input\_mask):
   if tf.random.uniform(()) > 0.5:
       # Random flipping of the image and mask
       input\_image = tf.image.flip\_left\_right(input\_image)
       input\_mask = tf.image.flip\_left\_right(input\_mask)

   return input\_image, input\_mask

We create a function to normalize the dataset by scaling the images to the range of `[-1, 1]` and decreasing the image mask by `1`:

def normalize(input\_image, input\_mask):
   input\_image = tf.cast(input\_image, tf.float32) / 255.0
   input\_mask -= 1
   return input\_image, input\_mask

We create two functions to preprocess the training and test datasets with a slight difference between the two – we only perform image augmentation on the training dataset.

def load\_image\_train(datapoint):
   input\_image = datapoint\["image"\]
   input\_mask = datapoint\["segmentation\_mask"\]
   input\_image, input\_mask = resize(input\_image, input\_mask)
   input\_image, input\_mask = augment(input\_image, input\_mask)
   input\_image, input\_mask = normalize(input\_image, input\_mask)

   return input\_image, input\_mask

def load\_image\_test(datapoint):
   input\_image = datapoint\["image"\]
   input\_mask = datapoint\["segmentation\_mask"\]
   input\_image, input\_mask = resize(input\_image, input\_mask)
   input\_image, input\_mask = normalize(input\_image, input\_mask)

   return input\_image, input\_mask

Now we are ready to build an input pipeline with `tf.data` by using the `map()` function:

train\_dataset = dataset\["train"\].map(load\_image\_train, num\_parallel\_calls=tf.data.AUTOTUNE)
test\_dataset = dataset\["test"\].map(load\_image\_test, num\_parallel\_calls=tf.data.AUTOTUNE)

If we execute `print(train_dataset)`, we will notice that the image is in the shape of `128x128x3` of `tf.float32` while the image mask is in the shape of `128x128x1` with the data type of `tf.uint8`.

We define a batch size of `64` and a buffer size of `1000` for creating batches of training and test datasets. With the original TFDS dataset, there are 3680 training samples and 3669 test samples, which are further split into validation/test sets. We will use the `train_batches` and the `validation_batches` for training the U-Net model. After the training finishes, we will then use the `test_batches` to test the model predictions.

BATCH\_SIZE = 64
BUFFER\_SIZE = 1000
train\_batches = train\_dataset.cache().shuffle(BUFFER\_SIZE).batch(BATCH\_SIZE).repeat()
train\_batches = train\_batches.prefetch(buffer\_size=tf.data.experimental.AUTOTUNE)
validation\_batches = test\_dataset.take(3000).batch(BATCH\_SIZE)
test\_batches = test\_dataset.skip(3000).take(669).batch(BATCH\_SIZE)

Now the datasets are ready for training. Let’s visualize a random sample image and its mask from the training dataset, to get an idea of how the data looks.​​

def display(display\_list):
 plt.figure(figsize=(15, 15))

 title = \["Input Image", "True Mask", "Predicted Mask"\]

 for i in range(len(display\_list)):
   plt.subplot(1, len(display\_list), i+1)
   plt.title(title\[i\])
   plt.imshow(tf.keras.utils.array\_to\_img(display\_list\[i\]))
   plt.axis("off")
 plt.show()

sample\_batch = next(iter(train\_batches))
random\_index = np.random.choice(sample\_batch\[0\].shape\[0\])
sample\_image, sample\_mask = sample\_batch\[0\]\[random\_index\], sample\_batch\[1\]\[random\_index\]
display(\[sample\_image, sample\_mask\])

The sample input image of a cat is in the shape of `128x128x3`. The true mask has three segments: the green background; the purple foreground object, in this case, a cat; and the yellow outline. **Figure 3** shows both the original input image and the true mask image.

![Image 5](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20700%20337'%3E%3C/svg%3E)

**Figure 3:** Sample input image and mask (image source: Oxford-IIIT pet TFDS).

### **Model Architecture**

Now that we have the data ready for training, let’s define the U-Net model architecture. As mentioned earlier, the U-Net is shaped like a letter U with an encoder, decoder, and the skip connections between them. So we will create a few building blocks to make the U-Net model.

**Building blocks**

First, we create a function `double_conv_block` with layers `Conv2D-ReLU-Conv2D-ReLU`, which we will use in both the encoder (or the contracting path) and the bottleneck of the U-Net.

def double\_conv\_block(x, n\_filters):

   # Conv2D then ReLU activation
   x = layers.Conv2D(n\_filters, 3, padding = "same", activation = "relu", kernel\_initializer = "he\_normal")(x)
   # Conv2D then ReLU activation
   x = layers.Conv2D(n\_filters, 3, padding = "same", activation = "relu", kernel\_initializer = "he\_normal")(x)

   return x

Then we define a `downsample_block` function for downsampling or feature extraction to be used in the encoder.

def downsample\_block(x, n\_filters):
   f = double\_conv\_block(x, n\_filters)
   p = layers.MaxPool2D(2)(f)
   p = layers.Dropout(0.3)(p)

   return f, p

Finally, we define an upsampling function `upsample_block` for the decoder (or expanding path) of the U-Net.

def upsample\_block(x, conv\_features, n\_filters):
   # upsample
   x = layers.Conv2DTranspose(n\_filters, 3, 2, padding="same")(x)
   # concatenate
   x = layers.concatenate(\[x, conv\_features\])
   # dropout
   x = layers.Dropout(0.3)(x)
   # Conv2D twice with ReLU activation
   x = double\_conv\_block(x, n\_filters)

   return x

### **U-Net Model**

There are three options for making a Keras model, as well explained in [Adrian’s blog](https://pyimagesearch.com/2019/10/28/3-ways-to-create-a-keras-model-with-tensorflow-2-0-sequential-functional-and-model-subclassing/) and the Keras documentation:

1.  **`Sequential` API:** easiest and beginner-friendly, stacking the layers sequentially.
2.  **Functional API:** more flexible and allows non-linear topology, shared layers, and multiple inputs or multi-outputs.
3.  **`Model` subclassing:** most flexible and best for complex models that need custom training loops.

U-Net has a fairly simple architecture; however, to create the skip connections between the encoder and decoder, we will need to concatenate some layers. So the Keras Functional API is most appropriate for this purpose.

First, we create a `build_unet_model` function, specify the inputs, encoder layers, bottleneck, decoder layers, and finally the output layer with `Conv2D` with activation of `softmax`. Note the input image shape is `128x128x3`. The output has three channels corresponding to the three classes that the model will classify each pixel for: background, foreground object, and object outline.

 # inputs
   inputs = layers.Input(shape=(128,128,3))

   # encoder: contracting path - downsample
   # 1 - downsample
   f1, p1 = downsample\_block(inputs, 64)
   # 2 - downsample
   f2, p2 = downsample\_block(p1, 128)
   # 3 - downsample
   f3, p3 = downsample\_block(p2, 256)
   # 4 - downsample
   f4, p4 = downsample\_block(p3, 512)

   # 5 - bottleneck
   bottleneck = double\_conv\_block(p4, 1024)

   # decoder: expanding path - upsample
   # 6 - upsample
   u6 = upsample\_block(bottleneck, f4, 512)
   # 7 - upsample
   u7 = upsample\_block(u6, f3, 256)
   # 8 - upsample
   u8 = upsample\_block(u7, f2, 128)
   # 9 - upsample
   u9 = upsample\_block(u8, f1, 64)

   # outputs
   outputs = layers.Conv2D(3, 1, padding="same", activation = "softmax")(u9)

   # unet model with Keras Functional API
   unet\_model = tf.keras.Model(inputs, outputs, name="U-Net")

   return unet\_model

We call the `build_unet_model` function to create the model `unet_model`:

unet\_model = build\_unet\_model()

And we can visualize the model architecture with `model.summary()` to see each detail of the model. And we can use a Keras utils function called `plot_model` to generate a more visual diagram, including the skip connections. The generated image generated in Colab is rotated 90 degrees so that you can see U-shaped architecture in **Figure 4** (see the details better in the Colab notebook):

![Image 6](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20700%20215'%3E%3C/svg%3E)

**Figure 4:** U-Net architecture generated by `plot_model` (image by the author).

### **Compile and Train U-Net**

To compile `unet_model`, we specify the optimizer, the loss function, and the accuracy metrics to track during training:

unet\_model.compile(optimizer=tf.keras.optimizers.Adam(),
                  loss="sparse\_categorical\_crossentropy",
                  metrics="accuracy")

We train the `unet_model` by calling `model.fit()` and training it for 20 epochs.

NUM\_EPOCHS = 20

TRAIN\_LENGTH = info.splits\["train"\].num\_examples
STEPS\_PER\_EPOCH = TRAIN\_LENGTH // BATCH\_SIZE

VAL\_SUBSPLITS = 5
TEST\_LENTH = info.splits\["test"\].num\_examples
VALIDATION\_STEPS = TEST\_LENTH // BATCH\_SIZE // VAL\_SUBSPLITS

model\_history = unet\_model.fit(train\_batches,
                              epochs=NUM\_EPOCHS,
                              steps\_per\_epoch=STEPS\_PER\_EPOCH,
                              validation\_steps=VALIDATION\_STEPS,
                              validation\_data=test\_batches)

After training for 20 epochs, we get a training accuracy and a validation accuracy of `~0.88`. The learning curve during training indicates that the model is doing well on both the training dataset and validation set, which indicates the model is generalizing well without much overfitting (**Figure 5).**

![Image 7](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20700%20347'%3E%3C/svg%3E)

**Figure 5:** U-Net learning curve (image by the author).

### **Prediction**

Now that we have completed training the `unet_model`, let’s use it to make predictions on a few sample images of the test dataset.

def create\_mask(pred\_mask):
 pred\_mask = tf.argmax(pred\_mask, axis=-1)
 pred\_mask = pred\_mask\[..., tf.newaxis\]
 return pred\_mask\[0\]

def show\_predictions(dataset=None, num=1):
 if dataset:
   for image, mask in dataset.take(num):
     pred\_mask = unet\_model.predict(image)
     display(\[image\[0\], mask\[0\], create\_mask(pred\_mask)\])
 else:
   display(\[sample\_image, sample\_mask,
            create\_mask(model.predict(sample\_image\[tf.newaxis, ...\]))\])

count = 0
for i in test\_batches:
   count +=1
print("number of batches:", count)

See **Figure 6** for the input images, the true masks, and the masks predicted by the U-Net model we trained.

![Image 8](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20500%20500'%3E%3C/svg%3E)

**Figure 6:** U-Net model predictions on test samples (image by the author).

### What's next? We recommend [PyImageSearch University](https://pyimagesearch.com/pyimagesearch-university/?utm_source=blogPost&utm_medium=bottomBanner&utm_campaign=What%27s%20next%3F%20I%20recommend).

**Course information:**  
84 total classes • 114+ hours of on-demand code walkthrough videos • Last updated: February 2024  
★★★★★ 4.84 (128 Ratings) • 16,000+ Students Enrolled

**I strongly believe that if you had the right teacher you could _master_ computer vision and deep learning.**

Do you think learning computer vision and deep learning has to be time-consuming, overwhelming, and complicated? Or has to involve complex mathematics and equations? Or requires a degree in computer science?

That’s _not_ the case.

All you need to master computer vision and deep learning is for someone to explain things to you in _simple, intuitive_ terms. _And that’s exactly what I do_. My mission is to change education and how complex Artificial Intelligence topics are taught.

If you're serious about learning computer vision, your next stop should be PyImageSearch University, the most comprehensive computer vision, deep learning, and OpenCV course online today. Here you’ll learn how to _successfully_ and _confidently_ apply computer vision to your work, research, and projects. Join me in computer vision mastery.

**Inside PyImageSearch University you'll find:**

*   ✓ **84 courses** on essential computer vision, deep learning, and OpenCV topics
*   ✓ **84 Certificates** of Completion
*   ✓ **114+ hours** of on-demand video
*   ✓ **Brand new courses released _regularly_**, ensuring you can keep up with state-of-the-art techniques
*   ✓ **Pre-configured Jupyter Notebooks in Google Colab**
*   ✓ Run all code examples in your web browser — works on Windows, macOS, and Linux (no dev environment configuration required!)
*   ✓ Access to **centralized code repos for _all_ 536+ tutorials** on PyImageSearch
*   ✓ **Easy one-click downloads** for code, datasets, pre-trained models, etc.
*   ✓ **Access** on mobile, laptop, desktop, etc.

[Click here to join PyImageSearch University](https://pyimagesearch.com/pyimagesearch-university/?utm_source=blogPost&utm_medium=bottomBanner&utm_campaign=What%27s%20next%3F%20I%20recommend)

**Summary**
-----------

In this post, you have learned how to load the Oxford-IIIT pet data with the TensorFlow dataset and how to train an image segmentation U-Net model from scratch. We created the U-Net with Keras Functional API and visualized the U-shaped architecture with skip connections. This post has been inspired by the official TensorFlow.org [image segmentation tutorial](https://www.tensorflow.org/tutorials/images/segmentation) and the [U-Net tutorial](https://keras.io/examples/vision/oxford_pets_image_segmentation/) on Keras.io, which uses `keras.utils.Sequence` for loading the data and has an Xception-style U-Net architecture. U-Net is a great start for learning semantic segmentation on images. To learn more about this topic, read segmentation papers on modern models such as [DeepLab V3](https://arxiv.org/abs/1706.05587), [HRNet](https://arxiv.org/abs/1908.07919), [U2-Net](https://arxiv.org/abs/2005.09007), etc., among many other papers.

### **Citation Information**

**Maynard-Reid, M.** “U-Net Image Segmentation in Keras,” _PyImageSearch_, 2022, [https://pyimg.co/6m5br](https://pyimg.co/6m5br)

@article{Maynard-Reid\_2022\_U-Net,
  author = {Margaret Maynard-Reid},
  title = {{U-Net} Image Segmentation in Keras},
  journal = {PyImageSearch},
  year = {2022},
  note = {https://pyimg.co/6m5br},
}

![Image 9: Featured Image](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20500%20276'%3E%3C/svg%3E)

### Unleash the potential of computer vision with Roboflow - Free!

*   Step into the realm of the future by [signing up or logging into your Roboflow account](https://roboflow.com/?ref=pyimagesearch). Unlock a wealth of innovative dataset libraries and revolutionize your computer vision operations.
*   Jumpstart your journey by choosing from our broad array of datasets, or benefit from PyimageSearch’s comprehensive library, crafted to cater to a wide range of requirements.
*   Transfer your data to Roboflow in any of the 40+ compatible formats. Leverage cutting-edge model architectures for training, and deploy seamlessly across diverse platforms, including API, NVIDIA, browser, iOS, and beyond. Integrate our platform effortlessly with your applications or your favorite third-party tools.
*   Equip yourself with the ability to train a potent computer vision model in a mere afternoon. With a few images, you can import data from any source via API, annotate images using our superior cloud-hosted tool, kickstart model training with a single click, and deploy the model via a hosted API endpoint. Tailor your process by opting for a code-centric approach, leveraging our intuitive, cloud-based UI, or combining both to fit your unique needs.
*   Embark on your journey today with absolutely no credit card required. Step into the future with Roboflow.

[Join Roboflow Now](https://roboflow.com/?ref=pyimagesearch)

**To download the source code to this post (and be notified when future tutorials are published here on PyImageSearch), _simply enter your email address in the form below!_**

![Image 10](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20410%20567'%3E%3C/svg%3E)

#### Download the Source Code and FREE 17-page Resource Guide

Enter your email address below to get a .zip of the code and a **FREE 17-page Resource Guide on Computer Vision, OpenCV, and Deep Learning.** Inside you'll find my hand-picked tutorials, books, courses, and libraries to help you master CV and DL!