{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.7.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":45867,"databundleVersionId":6924515,"sourceType":"competition"},{"sourceId":65833,"databundleVersionId":7245759,"sourceType":"competition"},{"sourceId":27923,"databundleVersionId":3495119,"sourceType":"competition"},{"sourceId":34639,"databundleVersionId":3268702,"sourceType":"competition"}],"dockerImageVersionId":30301,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"
\n\n

\nPaper Implementation - UNet
\n (implementation in Pytorch)\n

\n
","metadata":{}},{"cell_type":"markdown","source":"
\n\n

Please upvote if you learned from it.\n

\n
","metadata":{}},{"cell_type":"markdown","source":"# Paper covered\n### https://arxiv.org/abs/1505.04597
\n### U-Net: Convolutional Networks for Biomedical Image Segmentation","metadata":{}},{"cell_type":"markdown","source":"\"Unet\"","metadata":{}},{"cell_type":"markdown","source":"### About U-Net\n**U-Net** is an architecture for semantic segmentation. It consists of a contracting path and an expansive path. \nThe contracting path follows the typical architecture of a convolutional network. \n* It consists of the repeated application of two 3x3 convolutions (unpadded convolutions), each followed by a rectified linear unit (ReLU) and a 2x2 max pooling operation with stride 2 for downsampling. \n* At each downsampling step we double the number of feature channels. \n* Every step in the expansive path consists of an upsampling of the feature map followed by a 2x2 convolution (“up-convolution”) that halves the number of feature channels, a concatenation with the correspondingly cropped feature map from the contracting path, and two 3x3 convolutions, each followed by a ReLU. \n* The cropping is necessary due to the loss of border pixels in every convolution. \n* At the final layer a 1x1 convolution is used to map each 64-component feature vector to the desired number of classes. \n* In total the network has 23 convolutional layers.","metadata":{}},{"cell_type":"markdown","source":"### Key-Features\n* **U-Net learns segmentation in an end-to-end setting**.
You input a raw image and get a segmentation map as the output.\n\n* **U-Net is able to precisely localize and distinguish borders.**
Performs classification on every pixel so that the input and output share the same size.\n\n* **U-Net uses very few annotated images**.
Data augmentation with elastic deformations reduces the number of annotated images required for training.","metadata":{}},{"cell_type":"markdown","source":"Modules used:\n\n* `os`: A module that provides a way to interact with the operating system, allowing tasks such as reading or writing files and directories.\n\n* `numpy as np`: The NumPy library is used for numerical operations in Python, providing support for large, multi-dimensional arrays and matrices, along with mathematical functions to operate on these elements.\n\n* `torch`: The main PyTorch library, which is a popular deep learning framework. It provides tensors for efficient numerical computations and supports automatic differentiation for building and training neural networks.\n\n* `torch.nn as nn`: The neural network module of PyTorch, containing classes and functions for building and training neural networks. The nn module includes predefined layers, loss functions, and utilities for creating neural network architectures.\n\n* `torch.optim as optim`: The optimization module in PyTorch, which includes various optimization algorithms for updating the parameters of a neural network during training.\n\n* `torchvision.transforms`: Transforms provided by torchvision, a PyTorch library for computer vision tasks. These transforms are often used to preprocess and augment images before feeding them into neural networks.\n\n* `torch.nn.functional as F`: A module containing functional operations, including activation functions, loss functions, and other functions commonly used in neural network training.","metadata":{}},{"cell_type":"code","source":"import os\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torchvision.transforms\nimport torch.nn.functional as F","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2023-12-26T20:26:47.045482Z","iopub.execute_input":"2023-12-26T20:26:47.046401Z","iopub.status.idle":"2023-12-26T20:26:47.051436Z","shell.execute_reply.started":"2023-12-26T20:26:47.046365Z","shell.execute_reply":"2023-12-26T20:26:47.050388Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Attempt - 1","metadata":{}},{"cell_type":"markdown","source":"## The Black Dots / Block","metadata":{}},{"cell_type":"code","source":"class Block(nn.Module):\n def __init__(self, in_ch, out_ch):\n super().__init__()\n \n # First convolution layer with a 3x3 kernel\n self.conv1 = nn.Conv2d(in_ch, out_ch, kernel_size=3)\n # Batch normalization for the output of the first convolution\n self.bn1 = nn.BatchNorm2d(out_ch)\n # ReLU activation function\n self.relu = nn.ReLU()\n \n # Second convolution layer with a 3x3 kernel, operates on the output of the first layer\n self.conv2 = nn.Conv2d(out_ch, out_ch, kernel_size=3)\n # Batch normalization for the output of the second convolution\n self.bn2 = nn.BatchNorm2d(out_ch)\n\n def forward(self, x):\n # Pass through the first convolution\n x = self.conv1(x)\n # Batch normalization\n x = self.bn1(x)\n # ReLU activation\n x = self.relu(x)\n \n # Pass through the second convolution, operating on the output of the first convolution\n x = self.conv2(x)\n # Batch normalization for the output of the second convolution\n x = self.bn2(x)\n # ReLU activation\n x = self.relu(x)\n \n return x","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2023-12-26T20:26:47.657975Z","iopub.execute_input":"2023-12-26T20:26:47.658364Z","iopub.status.idle":"2023-12-26T20:26:47.667294Z","shell.execute_reply.started":"2023-12-26T20:26:47.65833Z","shell.execute_reply":"2023-12-26T20:26:47.666276Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Create an instance of the Block class with 1 input channel and 64 output channels\nenc_block = Block(1, 64)\n\n# Generate a random tensor with shape (1, 1, 572, 572)\nx = torch.randn(1, 1, 572, 572)\n\n# Pass the random tensor through the enc_block (instance of Block class)\noutput = enc_block(x)\n\n# Print the shape of the output tensor\nprint(output.shape)","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2023-12-26T20:26:47.784133Z","iopub.execute_input":"2023-12-26T20:26:47.784834Z","iopub.status.idle":"2023-12-26T20:26:48.148661Z","shell.execute_reply.started":"2023-12-26T20:26:47.784797Z","shell.execute_reply":"2023-12-26T20:26:48.147634Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Encoder","metadata":{}},{"cell_type":"code","source":"from pylab import *\nimport pylab as pylab","metadata":{"execution":{"iopub.status.busy":"2023-12-26T20:26:48.780943Z","iopub.execute_input":"2023-12-26T20:26:48.781637Z","iopub.status.idle":"2023-12-26T20:26:48.78647Z","shell.execute_reply.started":"2023-12-26T20:26:48.781601Z","shell.execute_reply":"2023-12-26T20:26:48.785303Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Encoder(nn.Module):\n def __init__(self, chs=(3, 64, 128, 256, 512, 1024)):\n super().__init__()\n\n # Create a list of encoder blocks based on the specified channels\n self.enc_blocks = nn.ModuleList([Block(chs[i], chs[i+1]) for i in range(len(chs) - 1)])\n # Max pooling layer with a kernel size of 2\n self.pool = nn.MaxPool2d(2)\n\n def forward(self, x):\n # List to store intermediate features from each encoder block\n ftrs = []\n \n # Iterate through each encoder block\n for block in self.enc_blocks:\n # Pass the input through the current encoder block\n x = block(x)\n # Append the output of the encoder block to the list of features\n ftrs.append(x)\n # Apply max pooling to reduce spatial dimensions\n x = self.pool(x)\n \n # Return the list of intermediate features\n return ftrs","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2023-12-26T20:26:49.161906Z","iopub.execute_input":"2023-12-26T20:26:49.162304Z","iopub.status.idle":"2023-12-26T20:26:49.171951Z","shell.execute_reply.started":"2023-12-26T20:26:49.162269Z","shell.execute_reply":"2023-12-26T20:26:49.170576Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Create an instance of the Encoder class\nencoder = Encoder()\n\n# Generate a random input tensor with shape (1, 3, 572, 572)\nx = torch.randn(1, 3, 572, 572)\n\n# Apply the encoder to the input tensor, obtaining a list of intermediate features\nftrs = encoder(x)\n\n# Iterate through the list of intermediate features and print their shapes\ncount = -1\nfor ftr in ftrs:\n count += 1\n print(f\"Filter encoder block {count} size: \", ftr.shape)","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2023-12-26T20:26:49.487163Z","iopub.execute_input":"2023-12-26T20:26:49.487597Z","iopub.status.idle":"2023-12-26T20:26:51.054708Z","shell.execute_reply.started":"2023-12-26T20:26:49.48756Z","shell.execute_reply":"2023-12-26T20:26:51.053688Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Decoder","metadata":{}},{"cell_type":"code","source":"class Decoder(nn.Module):\n def __init__(self, chs=(1024, 512, 256, 128, 64)):\n super().__init__()\n\n # Channels for each decoder block\n self.chs = chs\n # List of transposed convolution layers for upsampling\n self.upconvs = nn.ModuleList([nn.ConvTranspose2d(chs[i], chs[i+1], kernel_size=2, stride=2) for i in range(len(chs)-1)])\n # List of decoder blocks\n self.dec_blocks = nn.ModuleList([Block(chs[i], chs[i+1]) for i in range(len(chs)-1)])\n\n def forward(self, x, encoder_features):\n # Iterate through each decoder block\n for i in range(len(self.chs) - 1):\n # Upsample the input using transposed convolution\n x = self.upconvs[i](x)\n # Crop the corresponding encoder feature map to match the size of the upsampled tensor\n enc_ftrs = self.crop(encoder_features[i], x)\n # Concatenate the upsampled tensor with the cropped encoder feature map\n x = torch.cat([x, enc_ftrs], dim=1)\n # Pass the concatenated tensor through the decoder block\n x = self.dec_blocks[i](x)\n \n return x\n\n def crop(self, enc_ftrs, x):\n # Get the height and width of the upsampled tensor\n _, _, H, W = x.shape\n # Use torchvision transforms to center crop the encoder feature map\n enc_ftrs = torchvision.transforms.CenterCrop([H, W])(enc_ftrs)\n return enc_ftrs","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2023-12-26T20:26:51.056423Z","iopub.execute_input":"2023-12-26T20:26:51.056747Z","iopub.status.idle":"2023-12-26T20:26:51.068156Z","shell.execute_reply.started":"2023-12-26T20:26:51.056718Z","shell.execute_reply":"2023-12-26T20:26:51.067152Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Create an instance of the Decoder class\ndecoder = Decoder()\n\n# Generate a random tensor with shape (1, 1024, 28, 28)\nx = torch.randn(1, 1024, 28, 28)\n\n# Pass the tensor through the decoder, using reversed intermediate features (excluding the last element)\noutput = decoder(x, ftrs[::-1][1:])\n\n# Print the shape of the output tensor\nprint(\"Decoder output shape:\", output.shape)","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2023-12-26T20:26:51.069343Z","iopub.execute_input":"2023-12-26T20:26:51.069688Z","iopub.status.idle":"2023-12-26T20:26:52.334537Z","shell.execute_reply.started":"2023-12-26T20:26:51.069659Z","shell.execute_reply":"2023-12-26T20:26:52.333551Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Unet","metadata":{}},{"cell_type":"code","source":"class Unet(nn.Module):\n def __init__(self, enc_chs=(3, 64, 128, 256, 512, 1024), dec_chs=(1024, 512, 256, 128, 64), num_class=1, retain_dim=False):\n super().__init__()\n\n # Encoder module\n self.encoder = Encoder(enc_chs)\n # Decoder module\n self.decoder = Decoder(dec_chs)\n # Final convolutional layer to obtain the output\n self.head = nn.Conv2d(dec_chs[-1], num_class, kernel_size=1)\n # Option to retain the original input dimensions after interpolation\n self.retain_dim = retain_dim\n\n def forward(self, x, out_sz=(572, 572)):\n # Pass the input through the encoder\n enc_ftrs = self.encoder(x)\n \n # Pass the encoder features through the decoder\n out = self.decoder(enc_ftrs[::-1][0], enc_ftrs[::-1][1:])\n \n # Final convolutional layer\n out = self.head(out)\n \n # Interpolate to the specified output size if retain_dim is True\n if self.retain_dim:\n out = F.interpolate(out, out_sz)\n \n return out","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2023-12-26T20:26:52.336798Z","iopub.execute_input":"2023-12-26T20:26:52.337476Z","iopub.status.idle":"2023-12-26T20:26:52.346755Z","shell.execute_reply.started":"2023-12-26T20:26:52.337436Z","shell.execute_reply":"2023-12-26T20:26:52.345911Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Create an instance of the Unet class\nunet = Unet()\n\n# Generate a random tensor with shape (1, 3, 572, 572)\nx = torch.randn(1, 3, 572, 572)\n\n# Pass the tensor through the Unet model and print the shape of the output\noutput = unet(x)\nprint(\"Unet output shape:\", output.shape)","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2023-12-26T20:26:52.347931Z","iopub.execute_input":"2023-12-26T20:26:52.34821Z","iopub.status.idle":"2023-12-26T20:26:55.221021Z","shell.execute_reply.started":"2023-12-26T20:26:52.348184Z","shell.execute_reply":"2023-12-26T20:26:55.220031Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Attempt - 2","metadata":{}},{"cell_type":"markdown","source":"Let's make it such that we can change num_channels and out_channels. We will use this when we are doing application","metadata":{}},{"cell_type":"code","source":"def double_conv(in_c, out_c):\n # Define a sequential container for the double convolution block\n conv = nn.Sequential(\n # First convolutional layer with a 3x3 kernel\n nn.Conv2d(in_c, out_c, kernel_size=3),\n # Batch normalization for the output of the first convolution\n nn.BatchNorm2d(out_c),\n # ReLU activation function\n nn.ReLU(inplace=True),\n # Second convolutional layer with a 3x3 kernel, operates on the output of the first layer\n nn.Conv2d(out_c, out_c, kernel_size=3),\n # Batch normalization for the output of the second convolution\n nn.BatchNorm2d(out_c),\n # ReLU activation function\n nn.ReLU(inplace=True),\n )\n return conv","metadata":{"execution":{"iopub.status.busy":"2023-12-26T20:26:55.223091Z","iopub.execute_input":"2023-12-26T20:26:55.223395Z","iopub.status.idle":"2023-12-26T20:26:55.229646Z","shell.execute_reply.started":"2023-12-26T20:26:55.223367Z","shell.execute_reply":"2023-12-26T20:26:55.228565Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def crop_img(tensor, target_tensor):\n # Get the size of the target tensor\n target_size = target_tensor.size()[2]\n # Get the size of the input tensor\n tensor_size = tensor.size()[2]\n # Calculate the difference in size between the two tensors\n delta = tensor_size - target_size\n # Calculate the crop size for each side\n delta = delta // 2\n # Crop the input tensor\n return tensor[:, :, delta:tensor_size - delta, delta:tensor_size - delta]","metadata":{"execution":{"iopub.status.busy":"2023-12-26T20:26:55.230789Z","iopub.execute_input":"2023-12-26T20:26:55.231081Z","iopub.status.idle":"2023-12-26T20:26:55.245961Z","shell.execute_reply.started":"2023-12-26T20:26:55.231055Z","shell.execute_reply":"2023-12-26T20:26:55.245063Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class UNet(nn.Module):\n def __init__(self, num_channels, num_classes, retain_dim=True):\n super(UNet, self).__init__()\n\n # Encoder\n self.max_pool_2x2 = nn.MaxPool2d(kernel_size=2, stride=2)\n self.down_conv_1 = double_conv(num_channels, out_c=64)\n self.down_conv_2 = double_conv(in_c=64, out_c=128)\n self.down_conv_3 = double_conv(in_c=128, out_c=256)\n self.down_conv_4 = double_conv(in_c=256, out_c=512)\n self.down_conv_5 = double_conv(in_c=512, out_c=1024)\n\n # Decoder\n self.up_trans_1 = nn.ConvTranspose2d(in_channels=1024, out_channels=512, kernel_size=2, stride=2)\n self.up_conv_1 = double_conv(1024, 512)\n self.up_trans_2 = nn.ConvTranspose2d(in_channels=512, out_channels=256, kernel_size=2, stride=2)\n self.up_conv_2 = double_conv(512, 256)\n self.up_trans_3 = nn.ConvTranspose2d(in_channels=256, out_channels=128, kernel_size=2, stride=2)\n self.up_conv_3 = double_conv(256, 128)\n self.up_trans_4 = nn.ConvTranspose2d(in_channels=128, out_channels=64, kernel_size=2, stride=2)\n self.up_conv_4 = double_conv(128, 64)\n\n # Output layer\n self.out = nn.Conv2d(\n in_channels=64,\n out_channels=num_classes, # Number of objects to segment\n kernel_size=1,\n )\n self.retain_dim = retain_dim\n\n def forward(self, image, out_size=(512, 512)):\n # Encoder\n x1 = self.down_conv_1(image)\n x2 = self.max_pool_2x2(x1)\n x3 = self.down_conv_2(x2)\n x4 = self.max_pool_2x2(x3)\n x5 = self.down_conv_3(x4)\n x6 = self.max_pool_2x2(x5)\n x7 = self.down_conv_4(x6)\n x8 = self.max_pool_2x2(x7)\n x9 = self.down_conv_5(x8)\n\n # Decoder\n x = self.up_trans_1(x9)\n y = crop_img(x7, x)\n x = self.up_conv_1(torch.cat([x, y], 1))\n\n x = self.up_trans_2(x)\n y = crop_img(x5, x)\n x = self.up_conv_2(torch.cat([x, y], 1))\n\n x = self.up_trans_3(x)\n y = crop_img(x3, x)\n x = self.up_conv_3(torch.cat([x, y], 1))\n\n x = self.up_trans_4(x)\n y = crop_img(x1, x)\n x = self.up_conv_4(torch.cat([x, y], 1))\n\n # Output\n x = self.out(x)\n x = F.softmax(x, dim=1) # Applying softmax along the channel dimension\n return x","metadata":{"execution":{"iopub.status.busy":"2023-12-26T20:26:55.24724Z","iopub.execute_input":"2023-12-26T20:26:55.247566Z","iopub.status.idle":"2023-12-26T20:26:55.266066Z","shell.execute_reply.started":"2023-12-26T20:26:55.247531Z","shell.execute_reply":"2023-12-26T20:26:55.265135Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Generate a random tensor with shape (1, 3, 572, 572)\nimage = torch.rand((1, 3, 572, 572))\n\n# Create an instance of the UNet model\nmodel = UNet(num_channels=3, num_classes=1)\n\n# Pass the tensor through the UNet model and print the shape of the output\noutput = model(image)\nprint(\"UNet output shape:\", output.shape)","metadata":{"execution":{"iopub.status.busy":"2023-12-26T20:26:55.267246Z","iopub.execute_input":"2023-12-26T20:26:55.267566Z","iopub.status.idle":"2023-12-26T20:26:57.862514Z","shell.execute_reply.started":"2023-12-26T20:26:55.267526Z","shell.execute_reply":"2023-12-26T20:26:57.861382Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### So, now it's working and you can understand the mechanism of U-Net. I'll apply this to different datasets in my upcoming notebooks. Try experimenting with the given baseline also.\n\n#### Experiment with these datasets are image inputs.","metadata":{}},{"cell_type":"markdown","source":"#### Thank you for reading this notebook. I hope you found something useful. If you do find this notebook useful, please upvote. This will motivate me to add more examples to the series. Also comment which models you want.","metadata":{}},{"cell_type":"markdown","source":"
\n\n

Please upvote 🙏\n

\n
","metadata":{}}]}