U-Net & It's Advantages in Biomedical Image Segmentations
Introduction:
Image segmentation is one of the vastly used domain in machine learning. In image segmentation, the machine has to partition the image into various segments, each of them representing a different entity.
Convolutional Neural Networks (CNN) have shown decent results for image segmentation in the past, however, in CNN's we learn various important feature maps from the image by converting it into a vector. Contrary to the traditional CNN's, in medical image segmentation we not only need to learn important features from the image but we also need the same feature mappings to be converted back in an image. This is the complete recipe behind U-Nets. Therefore, in U-Nets we first contract the image in a vector to learn feature maps, in the second step we expand the same vector to a segmented image; while maintaining the structural integrity of an image to avoid any distortions. Let's understand the model in a little bit more detail and step by step implementation:
U Net Model:
The U-Net was first designed for?biomedical image segmentation and?demonstrated great results on the task of cell tracking.?The cool thing about the U-Net,?is that it can achieve relatively good results,?even with hundreds of examples.
The U-Net architecture gets its name from its U-like shape.?The U-Net consists of two paths:
The contracting path is?a typical convolutional network?like used in image classification.?It consists of repeated application of?convolution and pooling operations.?The convolution operation here?is called a down convolution.?The key here, is that in the contracting path,?our feature maps gets spatially smaller,?which is why it's called a contraction
2.?An expanding path.
The expanding path in some ways is doing?the opposite of the contracting path.?It's taking our small feature maps?through a series of up-sampling?and up-convolution steps to?get back to the original size of the image.?It also concatenates the up-sample representations at?each step with the corresponding feature maps?at the contraction pathway.?
Finally, at the last step,?the architecture outputs the probability?of e.g. a tumor for every pixel in the image.?
U-Net Implementation Using Keras library:
Import the elements you'll need to build your U-Net
import keras
from keras import backend as K
from keras.engine import Input, Model
from keras.layers import Conv3D, MaxPooling3D, UpSampling3D, Activation, BatchNormalization, PReLU, Deconvolution3D
from keras.optimizers import Adam
from keras.layers.merge import concatenate
# Set the image shape to have the channels in the first dimension
K.set_image_data_format("channels_first")
The "depth" of your U-Net is equal to the number of down-convolutions we will use. We will use a U-Net depth of 2, meaning you'll have 2 down-convolutions in our network.
1.1. Input Layers and its Depth:
In this article i will be briefly explaining how to do 3D image segmentation, which is to say that, in addition to "height" and "width", your input layer will also have a "length". "length" here to describe the third spatial dimension of the input so as not to confuse it with the depth of the network as defined above.
The shape of the input layer is?(num_channels, height, width, length), where?num_channels?you can think of like color channels in an image( for example red, blue, green etc) ,?height,?width?and?length?are just the size of the input.
For the assignment, the values will be:
input_layer = Input(shape=(4, 160, 160, 16)
2. Downward Contracting Path
In the downward path in our network (the left side of the U-Net). The?(height, width, length)?of the input gets smaller as we move down this path, and the number of channels increases.
2.1 Depth 0
By "depth 0" here, we're referring to the depth of the first down-convolution in the U-net. The number of filters is specified for each depth and for each layer within that depth.
The formula to use for calculating the number of filters is:
so at depth of i= 0 , we will have a total filter count of 32.
2.2 Layer 0
There are 2 convolutional layers for each depth. Let's define a Conv3D tensor with 32 filters. Kernel size , padding, strides concepts are the same that usually we follow for creating a convolution neural network layer.
down_depth_0_layer_0 = Conv3D(filters=32,
? ? ? ? ? ? ? ? ? ? ? ? ? ? ? kernel_size=(3,3,3),
? ? ? ? ? ? ? ? ? ? ? ? ? ? ? padding='same',
? ? ? ? ? ? ? ? ? ? ? ? ? ? ? strides=(1,1,1)
? ? ? ? ? ? ? ? ? ? ? ? ? ? ? )(input_layer)
down_depth_0_layer_0?
Notice that with 32 filters, the result you get above is a tensor with 32 channels. Run the next cell to add a relu activation to the first convolutional layer
down_depth_0_layer_0 = Activation('relu')(down_depth_0_layer_0)
down_depth_0_layer_0)
2.3 Depth 0, Layer 1
For layer 1 of depth 0, the formula for calculating the number of filters is:
where i is the current depth. Notice that additional multiplication by 2 is not there for layer 0.
Using the above formula the total number of filters = 32* (2^0) *2 = 64 filters.
# Create a Conv3D layer with 64 filters and add relu activation
down_depth_0_layer_1 = Conv3D(filters=64,?
? ? ? ? ? ? ? ? kernel_size=(3,3,3),
? ? ? ? ? ? ? ? padding='same',
? ? ? ? ? ? ? ? strides=(1,1,1)
? ? ? ? ? ? ? ?)(down_depth_0_layer_0)
down_depth_0_layer_1 = Activation('relu')(down_depth_0_layer_1)
down_depth_0_layer_1
2.4 Max Pooling
We have a max pooling operation after each down-convolution within the U-Net architecture.
In general we add the max pooling operation after each down-convolution, except for the last layer; so no maxpooling at (depth -1) down-convolution.
In our case:
a)The overall depth of the U-Net we are creating is 2
b) So the bottom of our U is is at a depth index of 2-1 = 1
c) so far we only have defined depth = 0 down convolutions, therefore we will add max pooling next.
Now let's add a max pooling operation to our U-Net.
down_depth_0_layer_pool = MaxPooling3D(pool_size=(2,2,2))(down_depth_0_layer_1)
down_depth_0_layer_pool
2.4.1 Depth 1, Layer 0
At depth 1, layer 0, the formula for calculating the number of filters = 32* (2^1) = 64. Now we will add a Conv3D layer to our network with a relu activation.
down_depth_1_layer_0 = Conv3D(filters=64,
? ? ? ? ? ? ? ? kernel_size=(3,3,3),
? ? ? ? ? ? ? ? padding='same',
? ? ? ? ? ? ? ? strides=(1,1,1)
? ? ? ? ? ? ? ?)(down_depth_0_layer_pool)
down_depth_1_layer_0 = Activation('relu')(down_depth_1_layer_0)
down_depth_1_layer_0?
2.4.1 Depth 1, Layer 1
For layer 1 of depth 1 the formula you'll use for number of filters is:
filters = 32* (2^i) * 2
where i is the current depth. Notice that *2 is not there for layer 0.
So at depth =1
filters = 32* (2^1) *2 = 128
Run the next cell to add another Conv3D with 128 filters to your network.
领英推荐
down_depth_1_layer_1 = Conv3D(filters=128,
? ? ? ? ? ? ? ? kernel_size=(3,3,3),
? ? ? ? ? ? ? ? padding='same',
? ? ? ? ? ? ? ? strides=(1,1,1)
? ? ? ? ? ? ? ?)(down_depth_1_layer_0)
down_depth_1_layer_1 = Activation('relu')(down_depth_1_layer_1)
down_depth_1_layer_1?
No max pooling at depth 1 (the bottom of the U)
When you get to the "bottom" of the U-net, you don't need to apply max pooling after the convolutions.
3. Expanding Upward Path
Now we will work on the expanding path of the U-Net, (going up on the right side, when viewing the diagram). The image's (height, width, length) all get larger in the expanding path.
3.1 Depth 0, Up Sampling Layer 0
We will use a pool size of (2,2,2) for up sampling. This is the default value for tf.keras.layers.UpSampling3D
a) As input to the upsampling at depth 0, you'll use the last layer of the down sampling which is down_depth_1_layer_1.
up_depth_0_layer_0 = UpSampling3D(size=(2,2,2))(down_depth_1_layer_1)
up_depth_0_layer_0
3.2 Concatenating Layers at the Same Depth ( Depth 0)
As you can see the above given diagram, each of the down-convolutions is also connected directly to the concatenation operations in the upsampling portion of the network. The key point to note here is that we are conactination operation on layers that are at the same depth of 0. So the upsampling layers at depth 0 (up_depth_0_layer_0) will be concatenated with the final downsampling layer at depth 0 which is ( down_depth_0_layer_1) in our case.
We need to keep the following pointers in mind while doing concatenations :
Print the shape of both of the layers to double check:
# Print the shape of layers to concatenate
print(up_depth_0_layer_0)
print()
print(down_depth_0_layer_1)
Result of the above print statement:
Tensor("up_sampling3d_1/concat_2:0", shape=(?, 128, 160, 160, 16), dtype=float32
Tensor("activation_2/Relu:0", shape=(?, 64, 160, 160, 16), dtype=float32))
As discussed both layers have the same height, width and length of ( 160,160,16). Let concatenate both of these layers by running the code mentioned below:
# Add a concatenation along axis
up_depth_1_concat = concatenate([up_depth_0_layer_0,
? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ?down_depth_0_layer_1],
? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? axis=1)
up_depth_1_concat1
up_depth_1_concat1 shape on printing is ( 192,160,160,16). Notice that the upsampling layer had 128 channels, and the down-convolution layer had 64 channels so that when concatenated, the result has 128 + 64 = 192 channels
3.3 Up- Convolution Depth 0 Layer 1
There are 2 convolutional layers for each depth as we had 2 layers in down-convolutions.
How to determine number of filters (number of neurons) = The number of filters for this layer will be set to the number of channels in the down-convolution's layer 1 at the same depth of 0 (down_depth_0_layer_1).
print(f"number of filters: {down_depth_0_layer_1._keras_shape[1]}")
number of filters = 64
# Add a Conv3D up-convolution with 64 filters to your network
up_depth_1_layer_1 = Conv3D(filters=64,?
? ? ? ? ? ? ? ? ? ? ? ? ? ? kernel_size=(3,3,3),
? ? ? ? ? ? ? ? ? ? ? ? ? ? padding='same',
? ? ? ? ? ? ? ? ? ? ? ? ? ? strides=(1,1,1)
? ? ? ? ? ? ? ? ? ? ? ? ? ?)(up_depth_1_concat)
up_depth_1_layer_1 = Activation('relu')(up_depth_1_layer_1)
up_depth_1_layer_1
on printing up_depth_1_layer_1 shape , we see that the shape is ( 64,160,160,16).
3.4 Up- Convolution Depth 0 Layer 2
The number of filters in this up-convolution will be equal to the number of filters in the down-convolution depth 0 layer 1. We always take the number of filters at final layer of the same depth. In our case we have upconvolution at depth 0, so we are using the number of filters from final layers of depth 0 at upconvolution which was down_depth_0_layer_1 in our case. Let's reprint and see the results
print(f"number of filters: {down_depth_0_layer_1._keras_shape[1]}")
number of filters = 64
Now we add a Conv3D up-convolution with 64 filters to our network.
# Add a Conv3D up-convolution with 64 filters to your network
up_depth_1_layer_2 = Conv3D(filters=64,?
? ? ? ? ? ? ? ? ? ? ? ? ? ? kernel_size=(3,3,3),
? ? ? ? ? ? ? ? ? ? ? ? ? ? padding='same',
? ? ? ? ? ? ? ? ? ? ? ? ? ? strides=(1,1,1)
? ? ? ? ? ? ? ? ? ? ? ? ? ?)(up_depth_1_layer_1)
up_depth_1_layer_2 = Activation('relu')(up_depth_1_layer_2)
up_depth_1_layer_2
4. Final Convolution
For the final convolution, number of filters will to be equal to the number of classes in your input data
For example the 3 classes can be such as:
Let's add the final Conv3D layer to our U-Net
# Add a final Conv3D with 3 filters to your network
final_conv = Conv3D(filters=3, #3 categories?
? ? ? ? ? ? ? ? ? ? kernel_size=(1,1,1),
? ? ? ? ? ? ? ? ? ? padding='valid',
? ? ? ? ? ? ? ? ? ? strides=(1,1,1)
? ? ? ? ? ? ? ? ? ? )(up_depth_1_layer_2)
final_conv.
4.1 Activation of the Final Convolution
Let's add sigmoid activation function to our final convolution.
final_activation = Activation('sigmoid')(final_conv)
final_activation
4.2 Create & Compile the U-Net Model
Let's set up a pre-built Keras categorical cross-entropy loss function. However, in actual real-life practice, you implement better loss functions to evaluate model performances such as Soft-Dice Loss functions.
# Define and compile your mode
model = Model(inputs=input_layer, outputs=final_activation)
model.compile(optimizer=Adam(lr=0.00001),
? ? ? ? ? ? ? loss='categorical_crossentropy',
? ? ? ? ? ? ? metrics=['categorical_accuracy']
? ? ? ? ? ? ?)l
Print the summary of your model we created.
# Print out a summary of the model we created
model.summary()
Output:
Model: "model_1
__________________________________________________________________________________________________
Layer (type) Output Shape Param # Connected to
==================================================================================================
input_2 (InputLayer) (None, 4, 160, 160, 0
__________________________________________________________________________________________________
conv3d_1 (Conv3D) (None, 32, 160, 160, 3488 input_2[0][0]
__________________________________________________________________________________________________
activation_1 (Activation) (None, 32, 160, 160, 0 conv3d_1[0][0]
__________________________________________________________________________________________________
conv3d_2 (Conv3D) (None, 64, 160, 160, 55360 activation_1[0][0]
__________________________________________________________________________________________________
activation_2 (Activation) (None, 64, 160, 160, 0 conv3d_2[0][0]
__________________________________________________________________________________________________
max_pooling3d_1 (MaxPooling3D) (None, 64, 80, 80, 8 0 activation_2[0][0]
__________________________________________________________________________________________________
conv3d_3 (Conv3D) (None, 64, 80, 80, 8 110656 max_pooling3d_1[0][0]
__________________________________________________________________________________________________
activation_3 (Activation) (None, 64, 80, 80, 8 0 conv3d_3[0][0]
__________________________________________________________________________________________________
conv3d_4 (Conv3D) (None, 128, 80, 80, 221312 activation_3[0][0]
__________________________________________________________________________________________________
activation_4 (Activation) (None, 128, 80, 80, 0 conv3d_4[0][0]
__________________________________________________________________________________________________
up_sampling3d_1 (UpSampling3D) (None, 128, 160, 160 0 activation_4[0][0]
__________________________________________________________________________________________________
concatenate_2 (Concatenate) (None, 192, 160, 160 0 up_sampling3d_1[0][0]
activation_2[0][0]
__________________________________________________________________________________________________
conv3d_5 (Conv3D) (None, 64, 160, 160, 331840 concatenate_2[0][0]
__________________________________________________________________________________________________
activation_5 (Activation) (None, 64, 160, 160, 0 conv3d_5[0][0]
__________________________________________________________________________________________________
conv3d_6 (Conv3D) (None, 64, 160, 160, 110656 activation_5[0][0]
__________________________________________________________________________________________________
activation_6 (Activation) (None, 64, 160, 160, 0 conv3d_6[0][0]
__________________________________________________________________________________________________
conv3d_7 (Conv3D) (None, 3, 160, 160, 195 activation_6[0][0]
__________________________________________________________________________________________________
activation_7 (Activation) (None, 3, 160, 160, 0 conv3d_7[0][0]
==================================================================================================
Total params: 833,507
Trainable params: 833,507
Non-trainable params: 0"
Remarks:
We have successfully introduced, explained and created our very own U-Net Model architecture of depth 2. Let me know if this was useful and i hope in your own medical image segmentation analysis, you will have an understanding on how to make your own U-net models. In future , i will use 3-U Net model to do a project on Brain Tumor Segmentation using MRI Scans data.
Annexure:
2. To go deeper, you can read the original research paper:
https://arxiv.org/pdf/1505.04597.pdf