{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load in \n\nimport os\nimport sys\nimport random\nimport warnings\n\nimport numpy as np\nimport pandas as pd\nimport gc\nimport matplotlib.pyplot as plt\n\nfrom tqdm import tqdm\nfrom itertools import chain\nfrom sklearn.model_selection import train_test_split\nfrom skimage.io import imread, imshow, imread_collection, concatenate_images\nfrom skimage.transform import resize\nfrom skimage.morphology import label\n\n\n# Input data files are available in the \"../input/\" directory.\n# For example, running this (by clicking run or pressing Shift+Enter) will list the files in the input directory\n\n# Set some parameters\nIMG_WIDTH = 128\nIMG_HEIGHT = 128\nIMG_CHANNELS = 3\nwidth_out = 128\nheight_out = 128\nTRAIN_PATH = '../input/stage1_train/'\nTEST_PATH = '../input/stage1_test/'\n\nwarnings.filterwarnings('ignore', category=UserWarning, module='skimage')\nseed = 42\nrandom.seed = seed\nnp.random.seed = seed\n\nprint(os.listdir(\"../input\"))\n\n# Any results you write to the current directory are saved as output.","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","trusted":true},"cell_type":"code","source":"# Get train and test IDs\ntrain_ids = next(os.walk(TRAIN_PATH))[1]\ntest_ids = next(os.walk(TEST_PATH))[1]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Get and resize train images and masks\nx_train = np.zeros((len(train_ids), IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS), dtype=np.uint8)\ny_train = np.zeros((len(train_ids), height_out, width_out, 1), dtype=np.bool)\nprint('Getting and resizing train images and masks ... ')\nsys.stdout.flush()\nfor n, id_ in tqdm(enumerate(train_ids), total=len(train_ids)):\n    path = TRAIN_PATH + id_\n    img = imread(path + '/images/' + id_ + '.png')[:,:,:IMG_CHANNELS]\n    img = resize(img, (IMG_HEIGHT, IMG_WIDTH), mode='constant', preserve_range=True)\n    x_train[n] = img\n    mask = np.zeros((height_out, width_out, 1), dtype=np.bool)\n    for mask_file in next(os.walk(path + '/masks/'))[2]:\n        mask_ = imread(path + '/masks/' + mask_file)\n        mask_ = np.expand_dims(resize(mask_, (height_out, width_out), mode='constant', \n                                      preserve_range=True), axis=-1)\n        mask = np.maximum(mask, mask_)\n    y_train[n] = mask\n\n# Get and resize test images\nx_test = np.zeros((len(test_ids), IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS), dtype=np.uint8)\nsizes_test = []\nprint('Getting and resizing test images ... ')\nsys.stdout.flush()\nfor n, id_ in tqdm(enumerate(test_ids), total=len(test_ids)):\n    path = TEST_PATH + id_\n    img = imread(path + '/images/' + id_ + '.png')[:,:,:IMG_CHANNELS]\n    sizes_test.append([img.shape[0], img.shape[1]])\n    img = resize(img, (IMG_HEIGHT, IMG_WIDTH), mode='constant', preserve_range=True)\n    x_test[n] = img\n\nprint('Done!')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"x_train.shape, y_train.shape, x_test.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fig, ax = plt.subplots(nrows=3, ncols=2, figsize=(9,12))\nax[0][0].imshow(x_train[0])\nax[0][1].imshow(np.squeeze(y_train[0]))\nax[1][0].imshow(x_train[1])\nax[1][1].imshow(np.squeeze(y_train[1]))\nax[2][0].imshow(x_train[2])\nax[2][1].imshow(np.squeeze(y_train[2]))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"x_train, x_val, y_train, y_val =  train_test_split(x_train, y_train, test_size=0.25)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\nx_train = x_train.transpose((0,3,1,2))\ny_train = y_train.transpose((0,3,1,2))\nx_val = x_val.transpose((0,3,1,2))\ny_val = y_val.transpose((0,3,1,2))\nx_train.shape, y_train.shape, x_val.shape, y_val.shape","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## U-Net Implementation"},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torchvision\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nfrom torch.autograd import Variable\nfrom tqdm import trange\nfrom time import sleep\nuse_gpu = torch.cuda.is_available()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"batch_size = 64\nepochs = 120\nepoch_lapse = 20\nthreshold = 0.33\nsample_size = None","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class UNet(nn.Module):\n    def contracting_block(self, in_channels, out_channels, kernel_size=3):\n        block = torch.nn.Sequential(\n                    torch.nn.Conv2d(kernel_size=kernel_size, in_channels=in_channels, out_channels=out_channels, padding=1),\n                    torch.nn.ReLU(),\n                    torch.nn.Dropout2d(),\n                    torch.nn.Conv2d(kernel_size=kernel_size, in_channels=out_channels, out_channels=out_channels, padding=1),\n                    torch.nn.ReLU(),\n                    torch.nn.BatchNorm2d(out_channels),\n                )\n        return block\n    \n    def expansive_block(self, in_channels, mid_channel, out_channels, kernel_size=3):\n            block = torch.nn.Sequential(\n                    torch.nn.Conv2d(kernel_size=kernel_size, in_channels=in_channels, out_channels=mid_channel, padding=1),\n                    torch.nn.ReLU(),\n                    torch.nn.Dropout2d(),\n                    torch.nn.Conv2d(kernel_size=kernel_size, in_channels=mid_channel, out_channels=mid_channel, padding=1),\n                    torch.nn.ReLU(),\n                    torch.nn.BatchNorm2d(mid_channel),\n                    torch.nn.ConvTranspose2d(in_channels=mid_channel, out_channels=out_channels, kernel_size=3, stride=2, padding=1, output_padding=1)\n                    )\n            return  block\n    \n    def final_block(self, in_channels, mid_channel, out_channels, kernel_size=3):\n            block = torch.nn.Sequential(\n                    torch.nn.Conv2d(kernel_size=kernel_size, in_channels=in_channels, out_channels=mid_channel, padding=1),\n                    torch.nn.ReLU(),\n                    torch.nn.Dropout2d(),\n                    torch.nn.Conv2d(kernel_size=kernel_size, in_channels=mid_channel, out_channels=mid_channel, padding=1),\n                    torch.nn.ReLU(),\n                    torch.nn.BatchNorm2d(mid_channel),\n                    torch.nn.Conv2d(kernel_size=kernel_size, in_channels=mid_channel, out_channels=out_channels, padding=1),\n                    torch.nn.Sigmoid(),\n                    )\n            return  block\n    \n    def __init__(self, in_channel, out_channel):\n        super(UNet, self).__init__()\n        #Encode\n        self.conv_encode1 = self.contracting_block(in_channels=in_channel, out_channels=64)\n        self.conv_maxpool1 = torch.nn.MaxPool2d(kernel_size=2)\n        self.conv_encode2 = self.contracting_block(64, 128)\n        self.conv_maxpool2 = torch.nn.MaxPool2d(kernel_size=2)\n        self.conv_encode3 = self.contracting_block(128, 256)\n        self.conv_maxpool3 = torch.nn.MaxPool2d(kernel_size=2)\n        self.conv_encode4 = self.contracting_block(256, 512)\n        self.conv_maxpool4 = torch.nn.MaxPool2d(kernel_size=2)\n        # Bottleneck\n        self.bottleneck = torch.nn.Sequential(\n                            torch.nn.Conv2d(kernel_size=3, in_channels=512, out_channels=1024, padding=1),\n                            torch.nn.ReLU(),\n                            torch.nn.BatchNorm2d(1024),\n                            torch.nn.Conv2d(kernel_size=3, in_channels=1024, out_channels=1024, padding=1),\n                            torch.nn.ReLU(),\n                            torch.nn.BatchNorm2d(1024),\n                            torch.nn.ConvTranspose2d(in_channels=1024, out_channels=512, kernel_size=3, stride=2, padding=1, output_padding=1)\n                            )\n        # Decode\n        self.conv_decode4 = self.expansive_block(1024, 512, 256)\n        self.conv_decode3 = self.expansive_block(512, 256, 128)\n        self.conv_decode2 = self.expansive_block(256, 128, 64)\n        self.final_layer = self.final_block(128, 64, out_channel)\n        \n    def crop_and_concat(self, upsampled, bypass, crop=False):\n        if crop:\n            c = (bypass.size()[2] - upsampled.size()[2]) // 2\n            bypass = F.pad(bypass, (-c, -c, -c, -c))\n        return torch.cat((upsampled, bypass), 1)\n    \n    def forward(self, x):\n        # Encode\n        encode_block1 = self.conv_encode1(x)\n        encode_pool1 = self.conv_maxpool1(encode_block1)\n        #print(encode_block1.shape)\n        encode_block2 = self.conv_encode2(encode_pool1)\n        encode_pool2 = self.conv_maxpool2(encode_block2)\n        encode_block3 = self.conv_encode3(encode_pool2)\n        encode_pool3 = self.conv_maxpool3(encode_block3)\n        encode_block4 = self.conv_encode4(encode_pool3)\n        encode_pool4  = self.conv_maxpool4(encode_block4)\n        # Bottleneck\n        bottleneck1 = self.bottleneck(encode_pool4)\n        # Decode\n        #print(x.shape, encode_block1.shape, encode_block2.shape, encode_block3.shape, encode_pool3.shape, bottleneck1.shape)\n        #print('Decode Block 3')\n        #print(bottleneck1.shape, encode_block3.shape)\n        decode_block4 = self.crop_and_concat(bottleneck1, encode_block4, crop=True)\n        cat_layer3 = self.conv_decode4(decode_block4)\n        decode_block3 = self.crop_and_concat(cat_layer3, encode_block3, crop=True)\n        #print(decode_block3.shape)\n        #print('Decode Block 2')\n        cat_layer2 = self.conv_decode3(decode_block3)\n        #print(cat_layer2.shape, encode_block2.shape)\n        decode_block2 = self.crop_and_concat(cat_layer2, encode_block2, crop=True)\n        cat_layer1 = self.conv_decode2(decode_block2)\n        #print(cat_layer1.shape, encode_block1.shape)\n        #print('Final Layer')\n        #print(cat_layer1.shape, encode_block1.shape)\n        decode_block1 = self.crop_and_concat(cat_layer1, encode_block1, crop=True)\n        #print(decode_block1.shape)\n        final_layer = self.final_layer(decode_block1)\n        #print(final_layer.shape)\n        return  final_layer\n        \ndef train_step(inputs, labels, optimizer, criterion, batch_size):\n    optimizer.zero_grad()\n    # forward + backward + optimize\n    outputs = unet(inputs)\n    # outputs.shape =(batch_size, n_classes, img_cols, img_rows)\n    outputs = outputs.permute(0, 2, 3, 1)\n    # outputs.shape =(batch_size, img_cols, img_rows, n_classes) \n    #print(outputs.shape)\n    outputs = outputs.reshape(batch_size*width_out*height_out, 2)\n    labels = labels.reshape(batch_size*width_out*height_out)\n    loss = criterion(outputs, labels)\n    loss.backward()\n    optimizer.step()\n    return loss","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"learning_rate = 0.01\nunet = UNet(in_channel=3,out_channel=2)\nif use_gpu:\n    unet = unet.cuda()\ncriterion = nn.CrossEntropyLoss()\noptimizer = optim.SGD(unet.parameters(), lr = 0.01, momentum=0.99)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_val_loss(x_val, y_val, batch_size=batch_size):\n    epoch_iter = np.ceil(x_val.shape[0] / batch_size).astype(int)\n    for _ in range(epoch_iter):\n        total_loss = 0\n        for i in range(epoch_iter):\n            batch_val_x = torch.from_numpy(x_val[i * batch_size : (i + 1) * batch_size]).float()\n            batch_val_y = torch.from_numpy(y_val[i * batch_size : (i + 1) * batch_size].astype(int)).long()\n            if use_gpu:\n                batch_val_x = batch_train_x.cuda()\n                batch_val_y = batch_train_y.cuda()\n            m = batch_val_x.shape[0]\n            outputs = unet(batch_val_x)\n            outputs = outputs.permute(0, 2, 3, 1)\n            # outputs.shape =(batch_size, img_cols, img_rows, n_classes)\n            #print(outputs.shape)\n            outputs = outputs.reshape(m*width_out*height_out, 2)\n            labels = batch_val_y.reshape(m*width_out*height_out)\n            loss = F.cross_entropy(outputs, labels)\n            total_loss += loss.data\n            gc.collect()\n    return total_loss / epoch_iter","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"epoch_iter = np.ceil(x_train.shape[0] / batch_size).astype(int)\nt = trange(epochs, leave=True)\nfor _ in t:\n    total_loss = 0\n    for i in range(epoch_iter):\n        batch_train_x = torch.from_numpy(x_train[i * batch_size : (i + 1) * batch_size]).float()\n        batch_train_y = torch.from_numpy(y_train[i * batch_size : (i + 1) * batch_size].astype(int)).long()\n        if use_gpu:\n            batch_train_x = batch_train_x.cuda()\n            batch_train_y = batch_train_y.cuda()\n        batch_loss = train_step(batch_train_x , batch_train_y, optimizer, criterion, batch_train_x.shape[0])\n        total_loss += batch_loss\n        gc.collect()\n    if (_+1) % epoch_lapse == 0:\n        val_loss = get_val_loss(x_val, y_val)\n        print(f\"Total loss in epoch {_+1} : {total_loss / epoch_iter} and validation loss : {val_loss}\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"gc.collect()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def plot_examples(datax, datay, num_examples=3):\n    fig, ax = plt.subplots(nrows=num_examples, ncols=3, figsize=(18,4*num_examples))\n    m = datax.shape[0]\n    for row_num in range(num_examples):\n        image_indx = np.random.randint(m)\n        image_arr = unet(torch.from_numpy(datax[image_indx:image_indx+1]).float().cuda()).squeeze(0).detach().cpu().numpy()\n        ax[row_num][0].imshow(np.transpose(datax[image_indx], (1,2,0))[:,:,0])\n        ax[row_num][0].set_title(\"Orignal Image\")\n        ax[row_num][1].imshow(np.squeeze((image_arr > 0.40)[1,:,:].astype(int)))\n        ax[row_num][1].set_title(\"Segmented Image localization\")\n        ax[row_num][2].imshow(np.transpose(datay[image_indx], (1,2,0))[:,:,0])\n        ax[row_num][2].set_title(\"Target image\")\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_examples(x_train, y_train, 12)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_examples(x_val, y_val, 12)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.save(unet.state_dict(), 'unet.pt')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.6.4","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat":4,"nbformat_minor":1}