Checking out the Implementation of UNet in TensorFlow platform

This study would provide the exact implementation of the original research paper of UNet, published in the not-so-secret archives of arXiv. Every AI enthusiast know the usages of CNN into their models and how they actually perform with their versatile applications. But this is something different. Here we would be using the CNN layer over another layer. And again. And again. This process is expected to continue till there are 23 layers of implementations done to process the data.

Literature Review:

U-Net architecture as defined on the research paper

The above shown diagram looks quite confusing but the code isn’t and is easily applicable. Just some basic coding experience is required and a lot of patience!

Decoding the Architecture of U-Net

After the image is fed into the system for processing (for this case, the researchers have chosen an image of dimensions 572px * 572px), the image immediately faces a double convolution network. In this double convolution network, two 3X3 convolution networks are implemented with ‘no padding’ conditions. Then, we apply Max Pooling. This procedure reduces the size of the image (after two convolutions) into half. You can say this acts as a image compressing algorithm. A 2X2 Max Pooling is implemented with no strides. The implementation has been done till the samples become 1024. This was the part of down-sampling. Now in the Up-sampling, we directly apply up-sampling of the data along with the convolution network and then we concatenate with the corresponding layer at its level (after convolution, as shown in the above picture). When we reach the implementation level of 2 samples, we then apply the last convolution layer of 1X1 form, and get the final desired output.

Still confused? Maybe the code will help you in this regard.

The Implementation:

Let’s first import all the necessary modules into our python environment. This would be the necessary modules that would define the implementation of the U-Net.

Necessary modules for the U-Net implementation

After the import of all the necessary modules into the environment, we need to define some functions which would be used in the functions otherwise.

Since, function definitions are trending now, hence we would define a function for carrying out all the necessary actions in the environment and hence make the decisions otherwise, by the course of time.

UNet function:

This is the main function of our concern and would be the basis of our entire project. In the project file itself, we would give the input size of the image, 572px X 572px.

The code for the implementation of the Down Sampling

Here, a manual representation has been preferred and implemented. However, it depends upon the sole interest of the coder to implement the same in the form of an integrated function or not. For the sake of easiness of the learner, the manual hard-coded representation has been preferred.

The code for the implementation of Up Sampling

The up sampling ensures that we receive the correct set of values for our images as the output. Here we are required to create a function named as ‘crop_image’. The work of this function is to crop the received image (in the form of tensor) and generate the final output in match with the dimensions of the tensor to which it would be concatenated.

The function as written for cropping the image

Here, we would be devising the ‘tensor’ and try to make the file have the same dimensions as that of our ‘target_tensor’.

The final segment of the code

In this segment, we have applied the last convolution network and hence made the final presentation by compiling and combining the same in the form of a attribute named ‘model’. We finally instantiate the attribute for the display.

Final output after the application of the U-Net (confirmed by the research paper)

Here, we are done. The final dimensions are reported and has been confirmed by the original research paper as well. This shows the complete implementation of U-Net using our favorite TensorFlow.

Usages:

U-Net can be used in various platforms for the functions ranging from edge detection to filtering and many more. You can expect to see the implementation in various robotics and biomedical projects, including brain tumor images segmentation, etc.

Good day, Happy Coding!