{"cells":[{"metadata":{},"cell_type":"markdown","source":"Ref: https://www.youtube.com/watch?v=u1loyDCoGbE\nhttps://arxiv.org/pdf/1505.04597.pdf\n\nhttps://www.youtube.com/watch?v=96_oGE8WyPg","execution_count":null},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import torch\nimport torch.nn as nn\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def double_conv(in_c, out_c):\n    conv = nn.Sequential(\n        nn.Conv2d(in_c, out_c, kernel_size = 3),\n        nn.ReLU(inplace = True),\n        nn.Conv2d(out_c, out_c, kernel_size = 3),\n        nn.ReLU(inplace = True)   \n    )\n    return conv\n\ndef crop_img(tensor, target_tensor):\n    target_size = target_tensor.size()[2]\n    tensor_size = tensor.size()[2]\n    delta = tensor_size - target_size\n    delta = delta // 2\n    return tensor[:, :, delta: tensor_size - delta, delta: tensor_size - delta]\n    ","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"class UNet(nn.Module):\n    def __init__(self):\n        super(UNet, self).__init__()\n\n        #define max-pool\n        self.max_pool_2x2 = nn.MaxPool2d(kernel_size = 2, stride = 2)\n        \n        #we need 5 of these\n        self.down_conv_1 = double_conv(1, 64)\n        self.down_conv_2 = double_conv(64, 128)\n        self.down_conv_3 = double_conv(128, 256)\n        self.down_conv_4 = double_conv(256, 512)\n        self.down_conv_5 = double_conv(512, 1024)\n        \n        self.up_trans_1 = nn.ConvTranspose2d(in_channels = 1024, \n                                             out_channels = 512,\n                                             kernel_size = 2, \n                                            stride = 2)\n        self.up_conv_1 = double_conv(1024, 512)\n        self.up_trans_2 = nn.ConvTranspose2d(in_channels = 512, \n                                             out_channels = 256,\n                                             kernel_size = 2, \n                                            stride = 2)\n        self.up_conv_2 = double_conv(512, 256)\n        self.up_trans_3 = nn.ConvTranspose2d(in_channels = 256, \n                                             out_channels = 128,\n                                             kernel_size = 2, \n                                            stride = 2)\n        self.up_conv_3 = double_conv(256, 128)\n        self.up_trans_4 = nn.ConvTranspose2d(in_channels = 128, \n                                             out_channels = 64,\n                                             kernel_size = 2, \n                                            stride = 2)\n        self.up_conv_4 = double_conv(128, 64)\n        \n        self.out = nn.Conv2d(\n            in_channels = 64,\n            out_channels = 2,\n            kernel_size = 1\n        )\n        \n        \n    def forward(self, image):\n        #bs, c, h, w\n        #encoder\n        x1 = self.down_conv_1(image) #\n        print(\"After down conv: \" + str(x1.size()))\n        x2 = self.max_pool_2x2(x1)\n        \n        x3 = self.down_conv_2(x2) #\n        x4 = self.max_pool_2x2(x3)\n        \n        x5 = self.down_conv_3(x4) #\n        x6 = self.max_pool_2x2(x5)\n        \n        x7 = self.down_conv_4(x6) #\n        x8 = self.max_pool_2x2(x7)\n        \n        x9 = self.down_conv_5(x8) \n        \n        print(\"after the last layer passed: \" + str(x9.size()))\n        \n        #decoder\n        x = self.up_trans_1(x9)\n        y = crop_img(x7, x)\n        print(\"the transposed conv part starting: \" + str(x.size()))\n        #x7 + x\n        x = self.up_conv_1(torch.cat([x, y], 1))\n        print(\"size of x7: \" + str(x7.size()))\n        print(\"size after cropping: \" + str(y.size()))\n        #paper says to crop it\n        \n        print(\"size after concat: \" + str(x.size()))\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        print(\"size after the FINAL layer: \" + str(x.size()))\n        \n        #add the output channel\n        x = self.out(x)\n        print(\"size of output: \" + str(x.size()))\n        #2 channel - 1 foreground, 1 background      ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if __name__ == \"__main__\":\n    #torch.rand(batch_size, channel, im_width, im_height)\n    image = torch.rand (1, 1, 572, 572)\n    model = UNet()\n    print(model(image))","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}