Updated to the UNet framework

Extended the implementation be more generic to support variations to the network in terms of the inputs and filter used at each convolution blocks/steps.

Author: sauravmishra1710

UNet - Biomedical_Segmentation.ipynb

@@ -70,8 +70,7 @@
    "outputs": [],
    "source": [
     "# import the UNet design framework module.\n",
-    "from Unet import Unet\n",
-    "unet = Unet()"
+    "from Unet import Unet"
    ]
   },
   {
@@ -80,6 +79,13 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "# Initialize the Unet with the default parameters. \n",
+    "# The default params are the one that were used in the original paper.\n",
+    "# Input shape - (572, 572, 1), \n",
+    "# filters 64, 128, 256, 512, 1024 at each convolutional block and \n",
+    "# unpadded convolutions.\n",
+    "unet = Unet(input_shape = (572, 572, 1), filters = [64, 128, 256, 512, 1024], padding = \"valid\")\n",
+    "\n",
     "# call the build netowrk API to build the network.\n",
     "model = unet.Build_UNetwork()"
    ]
@@ -182,18 +188,20 @@
    ],
    "source": [
     "# compile & summarize the model\n",
-    "unet.CompileAndSummarizeModel(model = model)"
+    "if model is not None:\n",
+    "    unet.CompileAndSummarizeModel(model = model)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {
     "scrolled": false
    },
    "outputs": [],
    "source": [
-    "unet.plotModel(model, dpi = 112)"
+    "if model is not None:\n",
+    "    unet.plotModel(model, dpi = 112)"
    ]
   },
   {

Unet.py

@@ -21,57 +21,62 @@ class Unet():
         https://arxiv.org/pdf/1505.04597
     """
 
-    def __init__(self):
-        pass
-    
-    def Build_UNetwork(self, input_shape = (572, 572, 1), filters = [64, 128, 256, 512, 1024]):
-        
+    def __init__(self, input_shape, filters, padding):
         """
-        Builds the Unet Model network.
+        
+        Initialize the Unet framework and the model parameters - input_shape, 
+        filters and padding type. 
         
         Args:
             input_shape: The shape of the input to the network. A tuple comprising of (img_height, img_width, channels).
-                         Default shape is (572, 572, 1).
+                         Original paper implementation is (572, 572, 1).
             filters: a collection of filters denoting the number of components to be used at each blocks along the 
-                     contracting and expansive paths. The default number of filters along the contracting and expansive paths are
-                     [64, 128, 256, 512, 1024].
+                     contracting and expansive paths. The original paper implementation for number of filters along the 
+                     contracting and expansive paths are [64, 128, 256, 512, 1024].
+            padding: the padding type to be used during the convolution step. The original paper used unpadded convolutions 
+                     which is of type "valid".
          
         **Remarks: The default values are as per the implementation in the original paper @ https://arxiv.org/pdf/1505.04597
-           
+        
+        """
+        self.input_shape = input_shape
+        self.filters = filters
+        self.padding = padding
+    
+    def Build_UNetwork(self):
+        
+        """
+        Builds the Unet Model network.
+        
+        Args:
+            None
+         
         Return:
             The Unet Model.
             
-            **Note - If the total number of filters are not sufficient to implement each block along the contracting 
-                     and expansive path, then the return value is None.
-        
         """
 
-        if len(filters) != 5:
-            print("There are not sufficient filters to implement each block of the UNet model.\nRecheck the filters.")
-            return None
 
-        UnetInput = Input(input_shape)
+        UnetInput = Input(self.input_shape)
 
         # the contracting path. 
         # the last item in the filetrs collection points to the number of filters in the bottleneck block.
-        # so we loop till the 4th item.
-        for num_filters in filters[:4]:
-            conv1, pool1 = UnetUtils.contracting_block(input_layer = UnetInput, filters = num_filters)
-            conv2, pool2 = UnetUtils.contracting_block(input_layer = pool1, filters = num_filters)
-            conv3, pool3 = UnetUtils.contracting_block(input_layer = pool2, filters = num_filters)
-            conv4, pool4 = UnetUtils.contracting_block(input_layer = pool3, filters = num_filters)
+        conv1, pool1 = UnetUtils.contracting_block(input_layer = UnetInput, filters = self.filters[0], padding = self.padding)
+        conv2, pool2 = UnetUtils.contracting_block(input_layer = pool1, filters = self.filters[1], padding = self.padding)
+        conv3, pool3 = UnetUtils.contracting_block(input_layer = pool2, filters = self.filters[2], padding = self.padding)
+        conv4, pool4 = UnetUtils.contracting_block(input_layer = pool3, filters = self.filters[3], padding = self.padding)
 
         # bottleneck block connecting the contracting and the expansive paths.
-        bottleNeck = UnetUtils.bottleneck_block(pool4, filters = filters[-1])
+        bottleNeck = UnetUtils.bottleneck_block(pool4, filters = self.filters[4], padding = self.padding)
 
-        # the expansive path.essentially we loop the reversed filter list leaving out the last item.
-        for num_filters in reversed(filters[:-1]):
-            upConv1 = UnetUtils.expansive_block(bottleNeck, conv4, filters = num_filters) 
-            upConv2 = UnetUtils.expansive_block(upConv1, conv3, filters = num_filters) 
-            upConv3 = UnetUtils.expansive_block(upConv2, conv2, filters = num_filters) 
-            upConv4 = UnetUtils.expansive_block(upConv3, conv1, filters = num_filters) 
+        # the expansive path.
+        upConv1 = UnetUtils.expansive_block(bottleNeck, conv4, filters = self.filters[3], padding = self.padding) 
+        upConv2 = UnetUtils.expansive_block(upConv1, conv3, filters = self.filters[2], padding = self.padding) 
+        upConv3 = UnetUtils.expansive_block(upConv2, conv2, filters = self.filters[1], padding = self.padding) 
+        upConv4 = UnetUtils.expansive_block(upConv3, conv1, filters = self.filters[0], padding = self.padding) 
 
-        UnetOutput = Conv2D(1, (1, 1), padding = "valid", activation = "sigmoid")(upConv4)
+        UnetOutput = Conv2D(1, (1, 1), padding = self.padding, activation = tf.math.sigmoid)(upConv4)
+        
         model = Model(UnetInput, UnetOutput, name = "UNet")
 
         return model

UnetUtils.py

@@ -21,7 +21,7 @@ class UnetUtils():
     def __init__(self):
         pass
 
-    def contracting_block(self, input_layer, filters, kernel_size = 3, padding = "valid"):
+    def contracting_block(self, input_layer, filters, padding, kernel_size = 3):
 
         """ 
         UNet Contracting block
@@ -56,7 +56,7 @@ def contracting_block(self, input_layer, filters, kernel_size = 3, padding = "va
 
         return conv, pool
 
-    def bottleneck_block(self, input_layer, filters, kernel_size = 3, padding = "valid", strides = 1):
+    def bottleneck_block(self, input_layer, filters, padding, kernel_size = 3, strides = 1):
 
         """ 
         UNet bottleneck block
@@ -90,7 +90,7 @@ def bottleneck_block(self, input_layer, filters, kernel_size = 3, padding = "val
 
         return conv
 
-    def expansive_block(self, input_layer, skip_conn_layer, filters, kernel_size = 3, padding = "valid", strides = 1):
+    def expansive_block(self, input_layer, skip_conn_layer, filters, padding, kernel_size = 3, strides = 1):
 
         """ 
         UNet expansive (upsample) block.
@@ -123,12 +123,15 @@ def expansive_block(self, input_layer, skip_conn_layer, filters, kernel_size = 3
                                     strides = 2, 
                                     padding = padding)(input_layer)
 
-        # crop the spurce feature map so that the skip connection can be established.
-        # Cropping is necessary due to the loss of border pixels in every convolution.
-        cropped = self.crop_tensor(skip_conn_layer, transConv)
-        
+        # crop the source feature map so that the skip connection can be established.
+        # the original paper implemented unpadded convolutions. So cropping is necessary 
+        # due to the loss of border pixels in every convolution.
         # establish the skip connections.
-        concat = Concatenate()([transConv, cropped])
+        if padding == "valid":
+            cropped = self.crop_tensor(skip_conn_layer, transConv)
+            concat = Concatenate()([transConv, cropped])
+        else:
+            concat = Concatenate()([transConv, skip_conn_layer])
 
         # two 3x3 convolutions, each followed by a ReLU
         up_conv = Conv2D(filters = filters,