Implement U-Net from Scratch for Image Segmentation
This type of image classification is called semantic image segmentation. It's similar to object detection in that both ask the question: "What objects are in this image and where in the image are those objects located?," but where object detection labels objects with bounding boxes that may include pixels that aren't part of the object, semantic image segmentation allows you to predict a precise mask for each object in the image by labeling each pixel in the image with its corresponding class. The word “semantic” here refers to what's being shown, so for example the “Car” class is indicated below by the dark blue mask, and "Person" is indicated with a red mask:
Figure 1: Example of a segmented image
As you might imagine, region-specific labeling is a pretty crucial consideration for self-driving cars, which require a pixel-perfect understanding of their environment so they can change lanes and avoid other cars, or any number of traffic obstacles that can put peoples' lives in danger.
By the time you finish this notebook, you'll be able to:
Build your own U-Net
Explain the difference between a regular CNN and a U-net
Implement semantic image segmentation on the CARLA self-driving car dataset
Apply sparse categorical cross-entropy for pixelwise prediction
Example of Masked and Unmasked images from the dataset
Preprocess Your Data
Normally, you normalize your image values by dividing them by 255. This sets them between 0 and 1. However, using tf.image.convert_image_dtype with tf.float32 sets them between 0 and 1 for you, so there's no need to further divide them by 255.
U-Net
U-Net, named for its U-shape, was originally created in 2015 for tumor detection, but in the years since has become a very popular choice for other semantic segmentation tasks.
U-Net builds on a previous architecture called the Fully Convolutional Network, or FCN, which replaces the dense layers found in a typical CNN with a transposed convolution layer that upsamples the feature map back to the size of the original input image, while preserving the spatial information. This is necessary because the dense layers destroy spatial information (the "where" of the image), which is an essential part of image segmentation tasks. An added bonus of using transpose convolutions is that the input size no longer needs to be fixed, as it does when dense layers are used.
Unfortunately, the final feature layer of the FCN suffers from information loss due to downsampling too much. It then becomes difficult to upsample after so much information has been lost, causing an output that looks rough.
U-Net improves on the FCN, using a somewhat similar design, but differing in some important ways. Instead of one transposed convolution at the end of the network, it uses a matching number of convolutions for downsampling the input image to a feature map, and transposed convolutions for upsampling those maps back up to the original input image size. It also adds skip connections, to retain information that would otherwise become lost during encoding. Skip connections send information to every upsampling layer in the decoder from the corresponding downsampling layer in the encoder, capturing finer information while also keeping computation low. These help prevent information loss, as well as model overfitting.
Images are first fed through several convolutional layers which reduce height and width, while growing the number of channels.
The contracting path follows a regular CNN architecture, with convolutional layers, their activations, and pooling layers to downsample the image and extract its features. In detail, it consists of the repeated application of two 3 x 3 same padding convolutions, each followed by a rectified linear unit (ReLU) and a 2 x 2 max pooling operation with stride 2 for downsampling. At each downsampling step, the number of feature channels is doubled.
Crop function: This step crops the image from the contracting path and concatenates it to the current image on the expanding path to create a skip connection.
The expanding path performs the opposite operation of the contracting path, growing the image back to its original size, while shrinking the channels gradually.
In detail, each step in the expanding path upsamples the feature map, followed by a 2 x 2 convolution (the transposed convolution). This transposed convolution halves the number of feature channels, while growing the height and width of the image.
Next is a concatenation with the correspondingly cropped feature map from the contracting path, and two 3 x 3 convolutions, each followed by a ReLU. You need to perform cropping to handle the loss of border pixels in every convolution.
Final Feature Mapping Block: In the final layer, a 1x1 convolution is used to map each 64-component feature vector to the desired number of classes. The channel dimensions from the previous layer correspond to the number of filters used, so when you use 1x1 convolutions, you can transform that dimension by choosing an appropriate number of 1x1 filters. When this idea is applied to the last layer, you can reduce the channel dimensions to have one layer per class.
The U-Net network has 23 convolutional layers in total.
Important Note:
The figures shown in the assignment for the U-Net architecture depict the layer dimensions and filter sizes as per the original paper on U-Net with smaller images. However, due to computational constraints for this assignment, you will code only half of those filters. The purpose of showing you the original dimensions is to give you the flavour of the original U-Net architecture. The important takeaway is that you multiply by 2 the number of filters used in the previous step. The notebook includes all of the necessary instructions and hints to help you code the U-Net architecture needed for this assignment.
Encoder (Downsampling Block)
Figure 3: The U-Net Encoder up close
The encoder is a stack of various conv_blocks:
Each conv_block() is composed of 2 Conv2D layers with ReLU activations. We will apply Dropout, and MaxPooling2D to some conv_blocks, as you will verify in the following sections, specifically to the last two blocks of the downsampling.
The function will return two tensors:
next_layer: That will go into the next block.
skip_connection: That will go into the corresponding decoding block.
Note: If max_pooling=True, the next_layer will be the output of the MaxPooling2D layer, but the skip_connection will be the output of the previously applied layer(Conv2D or Dropout, depending on the case). Else, both results will be identical.
Decoder (Upsampling Block)
The decoder, or upsampling block, upsamples the features back to the original image size. At each upsampling level, you'll take the output of the corresponding encoder block and concatenate it before feeding to the next decoder block.
Figure 4: The U-Net Decoder up close
There are two new components in the decoder: up and merge. These are the transpose convolution and the skip connections. In addition, there are two more convolutional layers set to the same parameters as in the encoder.
Here you'll encounter the Conv2DTranspose layer, which performs the inverse of the Conv2D layer. You can read more about it here.
In semantic segmentation, you need as many masks as you have object classes. In the dataset you're using, each pixel in every mask has been assigned a single integer probability that it belongs to a certain class, from 0 to num_classes-1. The correct class is the layer with the higher probability.
This is different from categorical crossentropy, where the labels should be one-hot encoded (just 0s and 1s). Here, you'll use sparse categorical crossentropy as your loss function, to perform pixel-wise multiclass prediction. Sparse categorical cross-entropy is more efficient than other loss functions when you're dealing with lots of classes.
Dataset Handling
Below, define a function that allows you to display both an input image, and its ground truth: the true mask. The true mask is what your trained model output is aiming to get as close to as possible.
True Mask of Input Image
Accuracy of Training
Predictions
Code
import tensorflow as tfimport numpy as npfrom tensorflow.keras.layers import Inputfrom tensorflow.keras.layers import Conv2Dfrom tensorflow.keras.layers import MaxPooling2Dfrom tensorflow.keras.layers import Dropout from tensorflow.keras.layers import Conv2DTransposefrom tensorflow.keras.layers import concatenatefrom test_utils import summary, comparator# Load and Split the Dataimport osimport numpy as np # linear algebraimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)import imageioimport matplotlib.pyplot as plt%matplotlib inlinepath =''image_path = os.path.join(path, './data/CameraRGB/')mask_path = os.path.join(path, './data/CameraMask/')image_list_orig = os.listdir(image_path)image_list = [image_path+i for i in image_list_orig]mask_list = [mask_path+i for i in image_list_orig]# Check out the some of the unmasked and masked images from the datasetN =2img = imageio.imread(image_list[N])mask = imageio.imread(mask_list[N])#mask = np.array([max(mask[i, j]) for i in range(mask.shape[0]) for j in range(mask.shape[1])]).reshape(img.shape[0], img.shape[1])
fig, arr = plt.subplots(1, 2, figsize=(14, 10))arr[0].imshow(img)arr[0].set_title('Image')arr[1].imshow(mask[:, :, 0])arr[1].set_title('Segmentation')# Split Dataset into Unmasked and Masked Imagesimage_list_ds = tf.data.Dataset.list_files(image_list, shuffle=False)mask_list_ds = tf.data.Dataset.list_files(mask_list, shuffle=False)for path inzip(image_list_ds.take(3), mask_list_ds.take(3)):print(path)image_filenames = tf.constant(image_list)masks_filenames = tf.constant(mask_list)dataset = tf.data.Dataset.from_tensor_slices((image_filenames, masks_filenames))for image, mask in dataset.take(1):print(image)print(mask)# Preprocess Datadefprocess_path(image_path,mask_path): img = tf.io.read_file(image_path) img = tf.image.decode_png(img, channels=3) img = tf.image.convert_image_dtype(img, tf.float32) mask = tf.io.read_file(mask_path) mask = tf.image.decode_png(mask, channels=3) mask = tf.math.reduce_max(mask, axis=-1, keepdims=True)return img, maskdefpreprocess(image,mask): input_image = tf.image.resize(image, (96, 128), method='nearest') input_mask = tf.image.resize(mask, (96, 128), method='nearest')return input_image, input_maskimage_ds = dataset.map(process_path)processed_image_ds = image_ds.map(preprocess)# U-Net Modeldefconv_block(inputs=None,n_filters=32,dropout_prob=0,max_pooling=True):""" Convolutional downsampling block Arguments: inputs -- Input tensor n_filters -- Number of filters for the convolutional layers dropout_prob -- Dropout probability max_pooling -- Use MaxPooling2D to reduce the spatial dimensions of the output volume Returns: next_layer, skip_connection -- Next layer and skip connection outputs """### START CODE HERE conv =Conv2D(n_filters, # Number of filters3, # Kernel size activation='relu', padding='same', kernel_initializer='he_normal')(inputs) conv =Conv2D(n_filters, # Number of filters3, # Kernel size activation='relu', padding='same',# set 'kernel_initializer' same as above kernel_initializer='he_normal')(conv)### END CODE HERE# if dropout_prob > 0 add a dropout layer, with the variable dropout_prob as parameterif dropout_prob >0:### START CODE HERE conv =Dropout(rate=dropout_prob)(conv)### END CODE HERE# if max_pooling is True add a MaxPooling2D with 2x2 pool_sizeif max_pooling:### START CODE HERE next_layer =MaxPooling2D(2,2)(conv)### END CODE HEREelse: next_layer = conv skip_connection = convreturn next_layer, skip_connection# Decoderdefupsampling_block(expansive_input,contractive_input,n_filters=32):""" Convolutional upsampling block Arguments: expansive_input -- Input tensor from previous layer contractive_input -- Input tensor from previous skip layer n_filters -- Number of filters for the convolutional layers Returns: conv -- Tensor output """### START CODE HERE up =Conv2DTranspose( n_filters, # number of filters3, # Kernel size strides=2, padding='same')(expansive_input)# Merge the previous output and the contractive_input merge =concatenate([up, contractive_input], axis=3) conv =Conv2D(n_filters, # Number of filters3, # Kernel size activation='relu', padding='same', kernel_initializer='he_normal')(merge) conv =Conv2D(n_filters, # Number of filters3, # Kernel size activation='relu', padding='same',# set 'kernel_initializer' same as above kernel_initializer='he_normal')(conv)### END CODE HEREreturn conv# Build the Modeldefunet_model(input_size=(96,128,3),n_filters=32,n_classes=23):""" Unet model Arguments: input_size -- Input shape n_filters -- Number of filters for the convolutional layers n_classes -- Number of output classes Returns: model -- tf.keras.Model """ inputs =Input(input_size)# Contracting Path (encoding)# Add a conv_block with the inputs of the unet_ model and n_filters### START CODE HERE cblock1 =conv_block(inputs, n_filters)# Chain the first element of the output of each block to be the input of the next conv_block. # Double the number of filters at each new step cblock2 =conv_block(cblock1[0], n_filters*2) cblock3 =conv_block(cblock2[0], n_filters*4) cblock4 =conv_block(cblock3[0], n_filters*8, dropout_prob=0.3)# Include a dropout_prob of 0.3 for this layer# Include a dropout_prob of 0.3 for this layer, and avoid the max_pooling layer cblock5 =conv_block(cblock4[0], n_filters*16, dropout_prob=0.3, max_pooling=False)### END CODE HERE# Expanding Path (decoding)# Add the first upsampling_block.# Use the cblock5[0] as expansive_input and cblock4[1] as contractive_input and n_filters * 8### START CODE HERE ublock6 =upsampling_block(cblock5[0], cblock4[1], n_filters *8)# Chain the output of the previous block as expansive_input and the corresponding contractive block output.# Note that you must use the second element of the contractive block i.e before the maxpooling layer. # At each step, use half the number of filters of the previous block ublock7 =upsampling_block(ublock6, cblock3[1], n_filters *4) ublock8 =upsampling_block(ublock7, cblock2[1], n_filters *2) ublock9 =upsampling_block(ublock8, cblock1[1], n_filters)### END CODE HERE conv9 =Conv2D(n_filters,3, activation='relu', padding='same',# set 'kernel_initializer' same as above exercises kernel_initializer='he_normal')(ublock9)# Add a Conv2D layer with n_classes filter, kernel size of 1 and a 'same' padding### START CODE HERE conv10 =Conv2D(n_classes, 1, padding='same')(conv9)### END CODE HERE model = tf.keras.Model(inputs=inputs, outputs=conv10)return modelimg_height =96img_width =128num_channels =3unet =unet_model((img_height, img_width, num_channels))unet.summary()# Loss Functionunet.compile(optimizer='adam', loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), metrics=['accuracy'])# Dataset Handlingdefdisplay(display_list): plt.figure(figsize=(15, 15)) title = ['Input Image','True Mask','Predicted Mask']for i inrange(len(display_list)): plt.subplot(1, len(display_list), i+1) plt.title(title[i]) plt.imshow(tf.keras.preprocessing.image.array_to_img(display_list[i])) plt.axis('off') plt.show()for image, mask in image_ds.take(1): sample_image, sample_mask = image, maskprint(mask.shape)display([sample_image, sample_mask])for image, mask in processed_image_ds.take(1): sample_image, sample_mask = image, maskprint(mask.shape)display([sample_image, sample_mask])# Train the Model EPOCHS =5VAL_SUBSPLITS =5BUFFER_SIZE =500BATCH_SIZE =32train_dataset = processed_image_ds.cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE)print(processed_image_ds.element_spec)model_history = unet.fit(train_dataset, epochs=EPOCHS)# Create Predicted Masksdefcreate_mask(pred_mask): pred_mask = tf.argmax(pred_mask, axis=-1) pred_mask = pred_mask[..., tf.newaxis]return pred_mask[0]# Plot Model Accuracyplt.plot(model_history.history["accuracy"])# Show Predictionsdefshow_predictions(dataset=None,num=1):""" Displays the first image of each of the num batches """if dataset:for image, mask in dataset.take(num): pred_mask = unet.predict(image)display([image[0], mask[0], create_mask(pred_mask)])else:display([sample_image, sample_mask,create_mask(unet.predict(sample_image[tf.newaxis, ...]))])show_predictions(train_dataset, 6)
