transfer-learning-DiceBCE loss.ipynb

{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import warnings\nwarnings.simplefilter(action='ignore', category=FutureWarning)\n\nimport torch\nimport time\nimport os\nimport numpy as np\nfrom pathlib import Path\nfrom PIL import Image\nimport helper\nimport matplotlib.pyplot as plt\nfrom matplotlib import pyplot\nfrom matplotlib.image import imread\nfrom torchvision import datasets\nfrom torchvision import datasets, transforms, models\nfrom torch import nn, optim, Tensor\nimport torch.nn.functional as F\nimport torch.nn as nn\nfrom torch.utils.data.sampler import SubsetRandomSampler\nfrom keras.preprocessing.image import load_img\nfrom keras.preprocessing.image import img_to_array\nfrom torch.utils.data import DataLoader, Dataset, TensorDataset\nimport random\nseed=42\ntorch.manual_seed(14)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2021-05-27T20:03:09.058265Z","iopub.execute_input":"2021-05-27T20:03:09.058763Z","iopub.status.idle":"2021-05-27T20:03:16.704306Z","shell.execute_reply.started":"2021-05-27T20:03:09.058654Z","shell.execute_reply":"2021-05-27T20:03:16.703477Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint('Using device:', device)\nprint()","metadata":{"execution":{"iopub.status.busy":"2021-05-27T20:03:16.705875Z","iopub.execute_input":"2021-05-27T20:03:16.706154Z","iopub.status.idle":"2021-05-27T20:03:16.713557Z","shell.execute_reply.started":"2021-05-27T20:03:16.706127Z","shell.execute_reply":"2021-05-27T20:03:16.712152Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class RGBCloudDataset (Dataset):\n    def __init__(self, red_dir, blue_dir, green_dir, gt_dir):\n        \n\n        # Listing subdirectories\n        # Loop through the files in red folder  \n        # and combine, into a dictionary, the other bands\n        \n        self.files = [self.combine_files(f, green_dir, blue_dir,gt_dir) \n                      for f in red_dir.iterdir() if not f.is_dir()]\n        \n        random.seed (seed)\n        self.files = random.sample (self.files, k= 10000)   \n        \n    def combine_files(self, red_file: Path, green_dir, blue_dir,  gt_dir):\n        \n        files = {'red': red_file, \n                 'green':green_dir/red_file.name.replace('red', 'green'),\n                 'blue': blue_dir/red_file.name.replace('red', 'blue'), \n                 'gt': gt_dir/red_file.name.replace('red', 'gt')}\n\n        return files\n    \n    \n    \n    def OpenAsArray(self, idx, invert=False):\n        \n        raw_rgb=np.stack([np.array(Image.open(self.files[idx]['red'])),\n                          np.array(Image.open(self.files[idx]['green'])),\n                          np.array(Image.open(self.files[idx]['blue']))], axis = 2)\n     \n\n        if invert:\n            raw_rgb = raw_rgb.transpose((2, 0, 1))\n    \n    \n        return (raw_rgb / np.iinfo(raw_rgb.dtype).max)\n    \n    \n    \n    \n    def OpenMask(self, idx, add_dims=False):\n        \n        raw_mask=np.array(Image.open(self.files[idx]['gt']))\n        raw_mask = np.where(raw_mask==255, 1, 0)\n        \n        \n        return np.expand_dims(raw_mask, 0) if add_dims else raw_mask\n\n\n\n        \n    def __len__(self):\n        \n        return len(self.files)\n    \n    \n    \n    def __getitem__(self, idx):\n        \n        x = torch.tensor(self.OpenAsArray(idx, invert=True), dtype=torch.float32)\n        y = torch.tensor(self.OpenMask(idx, add_dims=False), dtype=torch.int64)\n        \n        return x, y\n    \n    \n    \n    def open_as_pil(self, idx):\n        \n        arr = 256 * self.OpenAsArray(idx)\n        \n        return Image.fromarray(arr.astype(np.uint8), 'RGB')  \n    \n    \n    \n    def __repr__(self):\n        \n        s = 'Dataset class with {} files'.format(self.__len__())\n\n        return s","metadata":{"execution":{"iopub.status.busy":"2021-05-27T20:03:16.71538Z","iopub.execute_input":"2021-05-27T20:03:16.715713Z","iopub.status.idle":"2021-05-27T20:03:16.732686Z","shell.execute_reply.started":"2021-05-27T20:03:16.715682Z","shell.execute_reply":"2021-05-27T20:03:16.731207Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"base_path = Path('../input/95cloud-cloud-segmentation-on-satellite-images/95-cloud_training_only_additional_to38-cloud')\nred_dir   = base_path/'train_red_additional_to38cloud'\nblue_dir  = base_path/'train_blue_additional_to38cloud'\ngreen_dir = base_path/'train_green_additional_to38cloud'\nnir_dir   = base_path/'train_nir_additional_to38cloud'\ngt_dir    = base_path/'train_gt_additional_to38cloud'\n\nRGBdata = RGBCloudDataset(red_dir, blue_dir, green_dir, gt_dir)  ","metadata":{"execution":{"iopub.status.busy":"2021-05-27T20:03:22.744131Z","iopub.execute_input":"2021-05-27T20:03:22.744613Z","iopub.status.idle":"2021-05-27T20:04:24.128101Z","shell.execute_reply.started":"2021-05-27T20:03:22.744565Z","shell.execute_reply":"2021-05-27T20:04:24.126746Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import seaborn as sns\nfig, ax = plt.subplots(1,4, figsize = (10,9))\nplt.axis('off')\nred=np.array(Image.open(\"../input/95cloud-cloud-segmentation-on-satellite-images/95-cloud_training_only_additional_to38-cloud/train_red_additional_to38cloud/red_patch_100_5_by_12_LC08_L1TP_047011_20160920_20170221_01_T1.TIF\"))\ngreen=np.array(Image.open(\"../input/95cloud-cloud-segmentation-on-satellite-images/95-cloud_training_only_additional_to38-cloud/train_green_additional_to38cloud/green_patch_100_5_by_12_LC08_L1TP_047011_20160920_20170221_01_T1.TIF\"))\nblue=np.array(Image.open(\"../input/95cloud-cloud-segmentation-on-satellite-images/95-cloud_training_only_additional_to38-cloud/train_blue_additional_to38cloud/blue_patch_100_5_by_12_LC08_L1TP_047011_20160920_20170221_01_T1.TIF\"))\nnir=np.array(Image.open(\"../input/95cloud-cloud-segmentation-on-satellite-images/95-cloud_training_only_additional_to38-cloud/train_nir_additional_to38cloud/nir_patch_100_5_by_12_LC08_L1TP_047011_20160920_20170221_01_T1.TIF\"))\nplt.axis('off')\nax[0].imshow(red, cmap='gray')\nax[1].imshow(green, cmap='gray')\nax[2].imshow(blue, cmap='gray')\nax[3].imshow(nir, cmap='gray')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import seaborn as sns\nfig, ax = plt.subplots(1,1, figsize = (5,5))\nnir=np.array(Image.open(\"../input/95cloud-cloud-segmentation-on-satellite-images/95-cloud_training_only_additional_to38-cloud/train_nir_additional_to38cloud/nir_patch_100_5_by_12_LC08_L1TP_047011_20160920_20170221_01_T1.TIF\"))\nax.imshow(nir, cmap='gray')\nplt.axis('off')\nplt.savefig(\"nir_channel.png\", bbox_inches='tight')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"x, y =RGBdata[1000]\nx.shape, y. shape","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import seaborn as sns\nfig, ax = plt.subplots(1,2, figsize = (10,9))\nax[0].imshow(RGBdata.OpenAsArray(89))\nax[1].imshow(RGBdata.OpenMask(89))\nplt.axis('off')\nplt.savefig(\"Original_samples7.png\", bbox_inches='tight')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# splitting the data into train, validation, and test datasets\n\ntrain_size = int(0.75 * len(RGBdata))\nvalid_size = int(0.15 * len(RGBdata))\ntest_size  = len(RGBdata) - train_size - valid_size\nremaining_size = len(RGBdata) - train_size \n\nRGBtrain_dataset, RGBremaining_dataset = torch.utils.data.random_split(RGBdata, [train_size, remaining_size])\nRGBvalid_dataset, RGBtest_dataset      = torch.utils.data.random_split(RGBremaining_dataset, [valid_size, test_size])\n\n\nprint('\\t\\t\\tDataset')\nprint(\"Train data: \\t\\t{}\".format(len(RGBtrain_dataset)),\n      \"\\nValidation data: \\t{}\".format(len(RGBvalid_dataset)),\n     \"\\nTest data: \\t\\t{}\".format(len(RGBtest_dataset)))\n\n\n\nRGBtrain_loader = DataLoader(RGBtrain_dataset, batch_size=32, shuffle=True, num_workers=2)\nRGBvalid_loader = DataLoader(RGBvalid_dataset, batch_size=32, shuffle=True, num_workers=2)\nRGBtest_loader  = DataLoader(RGBtest_dataset , batch_size=32, shuffle=False, num_workers=2)\n\ndata_iter = iter(RGBvalid_loader)\nrgb_img, mask = next(data_iter)\n\nprint('\\n')\nprint('Raw RGB image shape on batch size = {}'.format(rgb_img.size()))\nprint('Cloud Mask shape on batch size    = {}'.format(mask.size()))","metadata":{"execution":{"iopub.status.busy":"2021-05-27T20:04:24.129847Z","iopub.execute_input":"2021-05-27T20:04:24.130145Z","iopub.status.idle":"2021-05-27T20:04:27.200335Z","shell.execute_reply.started":"2021-05-27T20:04:24.130116Z","shell.execute_reply":"2021-05-27T20:04:27.198843Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def batch_to_img(xb, idx):\n    img = np.array(xb[idx,0:3])\n    return img.transpose((1,2,0))\n\ndef predb_to_mask(predb, idx):\n    p = torch.functional.F.softmax(predb[idx], 0)\n    return p.argmax(0).cpu()","metadata":{"execution":{"iopub.status.busy":"2021-05-27T20:04:27.203251Z","iopub.execute_input":"2021-05-27T20:04:27.20361Z","iopub.status.idle":"2021-05-27T20:04:27.21101Z","shell.execute_reply.started":"2021-05-27T20:04:27.203568Z","shell.execute_reply":"2021-05-27T20:04:27.209819Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ResNet34(nn.Module):\n    def __init__(self):\n        super().__init__()\n        # Use a pretrained model\n        network = models.resnet34(pretrained=True)\n        # Replace the classifier with an upsampler\n        self.modified_network = nn.Sequential(*list((*list(network.children())[:-2],\n                                                     nn.Conv2d(512,512, kernel_size = 1),\n                                                     nn.BatchNorm2d(512),\n                                                     nn.ReLU(),\n                                                     nn.Conv2d(512,2, kernel_size = 1),\n                                                     nn.BatchNorm2d(2),\n                                                     nn.ReLU(),\n                                                     nn.Upsample(size=(384,384), mode='bilinear', align_corners=False))))\n    \n    \n    \n    \n    def forward(self, xb):\n        return self.modified_network(xb)\n    \n    def freeze(self):\n        # To freeze the CONV layers\n        for param in self.modified_network[:8].parameters():\n            param.require_grad = False\n        for param in self.modified_network[8].parameters():\n            param.require_grad = True\n        for param in self.modified_network[9].parameters():\n            param.require_grad = True\n    \n    def unfreeze(self):\n        # Unfreeze all layers\n        for param in self.modified_network.parameters():\n            param.require_grad = True","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ResNet34_np(nn.Module):\n    def __init__(self):\n        super().__init__()\n        # Use a pretrained model\n        network = models.resnet34(pretrained=False)\n        # Replace the classifier\n        self.modified_network = nn.Sequential(*list((*list(network.children())[:-2],\n                                                     nn.Conv2d(512,512, kernel_size = 1, bias=False),\n                                                     nn.BatchNorm2d(512),\n                                                     nn.ReLU(),\n                                                     nn.Conv2d(512,2, kernel_size = 1, bias=False),\n                                                     nn.BatchNorm2d(2),\n                                                     nn.ReLU(),\n                                                     nn.Upsample(size=(384,384), mode='bilinear', align_corners=False))))\n    \n    \n    \n    def forward(self, xb):\n        return self.modified_network(xb)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Vgg16(nn.Module):\n    def __init__(self):\n        super().__init__()\n        # Use a pretrained model\n        network = models.vgg16(pretrained=True)\n        # Replace the classifier\n        self.modified_network = nn.Sequential(*list((*list(network.children())[:-2],\n                                                     nn.Conv2d(512,512, kernel_size = 1),\n                                                     nn.BatchNorm2d(512),\n                                                     nn.ReLU(),\n                                                     nn.Conv2d(512,2, kernel_size = 1),\n                                                     nn.BatchNorm2d(2),\n                                                     nn.ReLU(),\n                                                     nn.Upsample(size=(384,384), mode='bilinear', align_corners=False)))) \n\n    \n    \n    \n    def forward(self, xb):\n        return self.modified_network(xb)\n    \n    def freeze(self):\n        # To freeze the CONV layers\n        for param in self.modified_network[0].parameters():\n            param.require_grad = False\n        for param in self.modified_network[1:].parameters():\n            param.require_grad = True\n    \n    def unfreeze(self):\n        # Unfreeze all layers\n        for param in self.modified_network.parameters():\n            param.require_grad = True","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Vgg16_np(nn.Module):\n    def __init__(self):\n        super().__init__()\n        # Use a pretrained model\n        network = models.vgg16(pretrained=False)\n        # Replace the classifier with an upsampling block\n        self.modified_network = nn.Sequential(*list((*list(network.children())[:-2],\n                                                     nn.Conv2d(512,512, kernel_size = 1),\n                                                     nn.BatchNorm2d(512),\n                                                     nn.ReLU(),\n                                                     nn.Conv2d(512,2, kernel_size = 1),\n                                                     nn.BatchNorm2d(2),\n                                                     nn.ReLU(),\n                                                     nn.Upsample(size=(384,384), mode='bilinear', align_corners=False)))) \n    \n    \n    \n    def forward(self, xb):\n        return self.modified_network(xb)","metadata":{"execution":{"iopub.status.busy":"2021-05-27T20:04:27.212928Z","iopub.execute_input":"2021-05-27T20:04:27.213217Z","iopub.status.idle":"2021-05-27T20:04:27.226338Z","shell.execute_reply.started":"2021-05-27T20:04:27.213191Z","shell.execute_reply":"2021-05-27T20:04:27.225554Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"PATH = \"../input/untrained-vgg16-dicebce/Untrained_Vgg16_DiceBCE_lr4.pt\"\n\n\nmodel = Vgg16_np()\nmap_location=torch.device('cpu')\nmodel.load_state_dict(torch.load(PATH,map_location=torch.device('cpu')),strict=False)\nmodel.eval()","metadata":{"execution":{"iopub.status.busy":"2021-05-27T20:04:27.228231Z","iopub.execute_input":"2021-05-27T20:04:27.228684Z","iopub.status.idle":"2021-05-27T20:04:31.51491Z","shell.execute_reply.started":"2021-05-27T20:04:27.228638Z","shell.execute_reply":"2021-05-27T20:04:31.513857Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"xb, yb = next(iter(RGBtest_loader))\n\nwith torch.no_grad():\n    predb = model(xb.to(device))\n\npredb.shape","metadata":{"execution":{"iopub.status.busy":"2021-05-27T20:04:31.516665Z","iopub.execute_input":"2021-05-27T20:04:31.517099Z","iopub.status.idle":"2021-05-27T20:05:05.872026Z","shell.execute_reply.started":"2021-05-27T20:04:31.517051Z","shell.execute_reply":"2021-05-27T20:05:05.871185Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"bs = 32\nfig, ax = plt.subplots(bs,3, figsize=(15,bs*5))\nfor i in range(bs):\n    ax[i,0].imshow(batch_to_img(xb,i))\n    ax[i,1].imshow(yb[i])\n    ax[i,2].imshow(predb_to_mask(predb, i))","metadata":{"execution":{"iopub.status.busy":"2021-05-27T20:05:05.873571Z","iopub.execute_input":"2021-05-27T20:05:05.874053Z","iopub.status.idle":"2021-05-27T20:05:22.15824Z","shell.execute_reply.started":"2021-05-27T20:05:05.874017Z","shell.execute_reply":"2021-05-27T20:05:22.156766Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fig.savefig('./bs32_untrained_vgg16_DiceBCE.png') ","metadata":{"execution":{"iopub.status.busy":"2021-05-27T20:07:57.826868Z","iopub.execute_input":"2021-05-27T20:07:57.827321Z","iopub.status.idle":"2021-05-27T20:08:02.14929Z","shell.execute_reply.started":"2021-05-27T20:07:57.827286Z","shell.execute_reply":"2021-05-27T20:08:02.148212Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Evaluation Metrics\n\nclass Evaluation_Metrics(nn.Module):\n    def __init__(self):\n        super(Evaluation_Metrics, self).__init__()\n\n    def forward(self, prediction, gt, smooth=1):\n        \n        pred = torch.round(prediction.softmax(dim=1)[:, 1])\n\n        # true positives, false positives, true negatives, false negatives\n        TP = torch.sum(pred * gt)\n        FP = torch.sum(pred * (1-gt))\n        TN = torch.sum((1-pred) * (1-gt))\n        FN = torch.sum((1-pred) * gt)\n    \n    \n        # Dice_Score/F1_Score\n        Dice_Score = (2 * TP + smooth)/(2*TP + FP + FN  + smooth)\n    \n        # Jaccard Coefficient: Intersection over Union\n        IoU = (TP + smooth)/(TP + FP + FN + smooth)\n    \n        # Recall\n        Recall= (TP + smooth)/(TP + FN + smooth)\n    \n        # Precision\n        Precision = (TP + smooth)/(TP + FP + smooth)\n\n    \n        return {'Dice_Score/F1_Score':Dice_Score, 'IoU':IoU, 'Recall': Recall, 'Precision': Precision}\n\n\n# Dice Loss function\nclass DiceLoss(nn.Module):\n    def __init__(self):\n        super(DiceLoss, self).__init__()\n\n    def forward(self,prediction, gt, smooth=1):\n        \n        # Softmax to get probabilities beteen 0 and 1\n        # Transform the prediction tensor of shape (N, C, H, W) --> tensor of shape (N, H, W)\n        pred = prediction.softmax(dim=1)[:, 1]\n        \n        #flatten label and prediction tensors\n        pred = pred.contiguous().view(-1)\n        gt = gt.contiguous().view(-1).to(torch.float32)\n        \n        intersection = (pred * gt).sum()                            \n        dice_loss =(2.*intersection + smooth)/(pred.sum() + gt.sum() + smooth)  \n        \n        return -torch.log(dice_loss)\n    \n\n#Binary cross-entropy (BCE)-Dice loss function\nclass DiceBCELoss(nn.Module):\n    def __init__(self):\n        super(DiceBCELoss, self).__init__()\n\n    def forward(self, prediction, gt, smooth=1):\n        \n        # Softmax to get probabilities beteen 0 and 1\n        # Transform the prediction tensor of shape (N, C, H, W) --> tensor of shape (N, H, W)\n        pred = prediction.softmax(dim=1)[:, 1]\n        \n        #flatten label and prediction tensors\n        pred = pred.contiguous().view(-1)\n        gt = gt.contiguous().view(-1).to(torch.float32)\n        \n        intersection = (pred * gt).sum()                            \n        dice_loss = 1 - (2.*intersection + smooth)/(pred.sum() + gt.sum() + smooth)  \n        BCE = F.binary_cross_entropy(pred, gt, reduction='mean')\n        Dice_BCE = BCE + dice_loss\n        \n        return Dice_BCE","metadata":{"execution":{"iopub.status.busy":"2021-05-27T18:52:42.636661Z","iopub.execute_input":"2021-05-27T18:52:42.637096Z","iopub.status.idle":"2021-05-27T18:52:42.656075Z","shell.execute_reply.started":"2021-05-27T18:52:42.637048Z","shell.execute_reply":"2021-05-27T18:52:42.65501Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Set the evaluation metrics and the loss function\nEvaluation = Evaluation_Metrics()\nloss_fn    = DiceBCELoss()\nlogDice_loss = DiceLoss()","metadata":{"execution":{"iopub.status.busy":"2021-05-27T18:52:42.657617Z","iopub.execute_input":"2021-05-27T18:52:42.658028Z","iopub.status.idle":"2021-05-27T18:52:42.673384Z","shell.execute_reply.started":"2021-05-27T18:52:42.657992Z","shell.execute_reply":"2021-05-27T18:52:42.671928Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_loss = 0.0\ntest_DSC  = 0.0\ntest_IoU  = 0.0\ntest_recall  = 0.0\ntest_precision  = 0.0\n\n\n# iterate over test data\nfor x, y in RGBtest_loader:\n    x = x.to(device, dtype=torch.float)\n    y = y.to(device)\n    \n    \n    with torch.no_grad():\n        model.eval()\n        outputs = model(x)\n        loss    = logDice_loss(outputs, y)\n                \n        # stats - whatever is the phase  \n        acc = Evaluation(outputs, y)\n        test_DSC        += acc['Dice_Score/F1_Score'].item()* RGBtest_loader.batch_size\n        test_IoU        += acc['IoU'].item()* RGBtest_loader.batch_size\n        test_recall     += acc['Recall'].item()* RGBtest_loader.batch_size\n        test_precision  += acc['Precision'].item()* RGBtest_loader.batch_size    \n        test_loss += loss.item() * RGBtest_loader.batch_size\n                    \nepoch_loss       = test_loss/ len(RGBtest_loader.dataset)\nepoch_DSC        = test_DSC / len(RGBtest_loader.dataset)\nepoch_IoU        = test_IoU / len(RGBtest_loader.dataset)\nepoch_Recall     = test_recall / len(RGBtest_loader.dataset)\nepoch_Precision  = test_precision / len(RGBtest_loader.dataset)\n            \n            \nprint('Dice Loss: {:.3f}\\tDice Coefficient: {:.3f}\\tJaccard Coefficient: {:.3f}\\tPrecision: {:.3f}\\tRecall: {:.3f}'.format(epoch_loss, epoch_DSC, epoch_IoU, epoch_Precision, epoch_Recall))\nprint()","metadata":{"execution":{"iopub.status.busy":"2021-05-27T18:53:04.387169Z","iopub.execute_input":"2021-05-27T18:53:04.387592Z","iopub.status.idle":"2021-05-27T19:05:52.04447Z","shell.execute_reply.started":"2021-05-27T18:53:04.387532Z","shell.execute_reply":"2021-05-27T19:05:52.043481Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"network = models.resnet34(pretrained=True)\nnetwork ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## **Model: Pretrained Vgg16**","metadata":{}},{"cell_type":"code","source":"class Vgg16(nn.Module):\n    def __init__(self):\n        super().__init__()\n        # Use a pretrained model\n        network = models.vgg16(pretrained=True)\n        # Replace the classifier\n        self.modified_network = nn.Sequential(*list((*list(network.children())[:-2],\n                                                     nn.Conv2d(512,512, kernel_size = 1),\n                                                     nn.BatchNorm2d(512),\n                                                     nn.ReLU(),\n                                                     nn.Conv2d(512,2, kernel_size = 1),\n                                                     nn.BatchNorm2d(2),\n                                                     nn.ReLU(),\n                                                     nn.Upsample(size=(384,384), mode='bilinear', align_corners=False)))) \n\n    \n    \n    \n    def forward(self, xb):\n        return self.modified_network(xb)\n    \n    def freeze(self):\n        # To freeze the CONV layers\n        for param in self.modified_network[0].parameters():\n            param.require_grad = False\n        for param in self.modified_network[1:].parameters():\n            param.require_grad = True\n    \n    def unfreeze(self):\n        # Unfreeze all layers\n        for param in self.modified_network.parameters():\n            param.require_grad = True","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Evaluation Metrics\n\nclass Evaluation_Metrics(nn.Module):\n    def __init__(self):\n        super(Evaluation_Metrics, self).__init__()\n\n    def forward(self, prediction, gt, smooth=1):\n        \n        pred = torch.round(prediction.softmax(dim=1)[:, 1])\n\n        # true positives, false positives, true negatives, false negatives\n        TP = torch.sum(pred * gt)\n        FP = torch.sum(pred * (1-gt))\n        TN = torch.sum((1-pred) * (1-gt))\n        FN = torch.sum((1-pred) * gt)\n    \n    \n        # Dice_Score/F1_Score\n        Dice_Score = (2 * TP + smooth)/(2*TP + FP + FN  + smooth)\n    \n        # Jaccard Coefficient: Intersection over Union\n        IoU = (TP + smooth)/(TP + FP + FN + smooth)\n    \n        # Recall\n        Recall= (TP + smooth)/(TP + FN + smooth)\n    \n        # Precision\n        Precision = (TP + smooth)/(TP + FP + smooth)\n\n    \n        return {'Dice_Score/F1_Score':Dice_Score, 'IoU':IoU, 'Recall': Recall, 'Precision': Precision}\n\n\n# Dice Loss function\nclass DiceBCELoss(nn.Module):\n    def __init__(self):\n        super(DiceBCELoss, self).__init__()\n\n    def forward(self, prediction, gt, smooth=1):\n        \n        # Softmax to get probabilities beteen 0 and 1\n        # Transform the prediction tensor of shape (N, C, H, W) --> tensor of shape (N, H, W)\n        pred = prediction.softmax(dim=1)[:, 1]\n        \n        #flatten label and prediction tensors\n        pred = pred.contiguous().view(-1)\n        gt = gt.contiguous().view(-1).to(torch.float32)\n        \n        intersection = (pred * gt).sum()                            \n        dice_loss = 1 - (2.*intersection + smooth)/(pred.sum() + gt.sum() + smooth)  \n        BCE = F.binary_cross_entropy(pred, gt, reduction='mean')\n        Dice_BCE = BCE + dice_loss\n        \n        return Dice_BCE","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Use GPU if it is available\nmodel = Vgg16().to(device)\n\n#checking if we have the correct dimensions for our (3,384,384) images\n#summary(Vgg16, (3, 384, 384))   \n\n# checking if the network works, using one batch of the training set\nrgb_img, mask = next(iter(RGBtrain_loader))\nrgb_img, mask = rgb_img.to(device, dtype=torch.float), mask.to(device) \noutput = model (rgb_img).to(device)\noutput.shape, mask.shape ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Set the evaluation metrics and the loss function\nEvaluation = Evaluation_Metrics()\nloss_fn    = DiceBCELoss()\n\n# Set the optimizer \noptimizer = optim.Adam(model.parameters(), lr=0.0001)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def plot_DSC(train_DSC, valid_DSC):\n    fig = plt.figure(figsize=(10, 10))\n    plt.plot(train_DSC, '-bx')\n    plt.plot(valid_DSC, '-rx')\n    plt.xlabel('epoch')\n    plt.ylabel('Dice Coefficient')\n    plt.title('Dice Coefficient vs. No. of epochs');\n    \ndef plot_losses(train_losses, valid_losses):\n    fig = plt.figure(figsize=(10, 10))\n    plt.plot(train_losses, '-bx')\n    plt.plot(valid_losses, '-rx')\n    plt.xlabel('epoch')\n    plt.ylabel('DiceBCE loss')\n    plt.legend(['Training', 'Validation'])\n    plt.title('DiceBCE Loss vs. No. of epochs');","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train(model, train_dl, valid_dl, loss_fn, optimizer, Evaluation, epochs, save_path):\n    \n    start = time.time()\n    min_valid_epoch_loss = np.inf\n    best_DSC = 0.0\n    train_losses, valid_losses = [], []\n    train_DSC, train_IoU, train_recall, train_precision = [], [], [], [] \n    valid_DSC, valid_IoU, valid_recall, valid_precision = [], [], [], [] \n\n    for e in range(epochs):\n        print('Epoch  {}/{}'.format(e, epochs-1))\n        print('-' * 10)\n        \n        # Each epoch has a training and validation phase\n        for phase in ['train', 'valid']:\n            if phase == 'train':\n                model.train()  # Set model to train mode\n                dataloader = train_dl\n            else:\n                model.eval()  # Set model to evaluate mode\n                dataloader = valid_dl\n\n            running_loss = 0.0\n            running_DSC  = 0.0\n            running_IoU  = 0.0\n            running_recall  = 0.0\n            running_precision  = 0.0\n\n\n            # iterate over data\n            for x, y in dataloader:\n                x = x.to(device, dtype=torch.float)\n                y = y.to(device)\n                \n\n                # forward pass \n                if phase == 'train':\n                    # zero the gradients\n                    optimizer.zero_grad()\n                    outputs = model(x)\n                    loss    = loss_fn(outputs, y)\n                    # backward + optimize only if in training phase\n                    # the backward pass frees the graph memory, so there is no \n                    # need for torch.no_grad in this training pass\n                    loss.backward()\n                    optimizer.step()\n                    # scheduler.step()\n\n                else:\n                    with torch.no_grad():\n                        outputs = model(x)\n                        loss    = loss_fn(outputs, y)\n\n                # stats - whatever is the phase  \n                acc = Evaluation(outputs, y)\n                running_DSC        += acc['Dice_Score/F1_Score'].item()* dataloader.batch_size\n                running_IoU        += acc['IoU'].item()* dataloader.batch_size\n                running_recall     += acc['Recall'].item()* dataloader.batch_size\n                running_precision  += acc['Precision'].item()* dataloader.batch_size    \n                \n                running_loss += loss.item() * dataloader.batch_size\n                    \n            epoch_loss       = running_loss/ len(dataloader.dataset)\n            epoch_DSC        = running_DSC / len(dataloader.dataset)\n            epoch_IoU        = running_IoU / len(dataloader.dataset)\n            epoch_Recall     = running_recall / len(dataloader.dataset)\n            epoch_Precision  = running_precision / len(dataloader.dataset)\n            \n\n            print('{}\\nDice Loss: {:.3f}\\tDice Coefficient: {:.3f}\\tJaccard Coefficient: {:.3f}\\tPrecision: {:.3f}\\tRecall: {:.3f}'.format( phase, epoch_loss, epoch_DSC, epoch_IoU, epoch_Precision, epoch_Recall))\n            print()\n\n            train_losses.append(epoch_loss), train_DSC.append(epoch_DSC),train_IoU.append(epoch_IoU),train_recall.append(epoch_Recall),train_precision.append(epoch_Precision) if phase=='train' else valid_losses.append(epoch_loss), valid_DSC.append(epoch_DSC),valid_IoU.append(epoch_IoU),valid_recall.append(epoch_Recall),valid_precision.append(epoch_Precision)\n\n            \n             \n             # save model if validation loss has decreased\n            if phase == 'valid':\n                if epoch_loss <= min_valid_epoch_loss:\n                    print('Validation Dice loss decreased ({:.3f} --> {:.3f}).  Saving model ...'.format(\n                        min_valid_epoch_loss, epoch_loss)) \n                    print('Best Validation Dice Score ({:.3f} --> {:.3f}).'.format(\n                        best_DSC, epoch_DSC))\n                    torch.save(model.state_dict(), save_path)\n                    min_valid_epoch_loss = epoch_loss\n                    best_DSC = epoch_DSC\n            \n\n    time_elapsed = time.time() - start\n    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))    \n    \n    \n    return train_losses, valid_losses, train_DSC, train_IoU, train_recall, train_precision,valid_DSC, valid_IoU, valid_recall, valid_precision, best_DSC","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model.freeze()\nPretrained_Vgg16 = './Pretrained_Vgg16_DiceScore_lr4_bs32.pt'\ntrain_losses, valid_losses, train_DSC, train_IoU, train_recall, train_precision,valid_DSC, valid_IoU, valid_recall, valid_precision, best_DSC = train(model, RGBtrain_loader, RGBvalid_loader, loss_fn, optimizer, Evaluation, epochs=40, save_path = Pretrained_Vgg16)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_DSC(train_DSC, valid_DSC)\nplt.savefig('./Dice_plot_Pretrained_Vgg16_bs32.png')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_losses(train_losses, valid_losses)\nplt.savefig('./Losses_plot_Pretrained_Vgg16_bs32.png')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Dice Loss function\nclass DiceLoss(nn.Module):\n    def __init__(self):\n        super(DiceLoss, self).__init__()\n\n    def forward(self,prediction, gt, smooth=1):\n        \n        # Softmax to get probabilities beteen 0 and 1\n        # Transform the prediction tensor of shape (N, C, H, W) --> tensor of shape (N, H, W)\n        pred = prediction.softmax(dim=1)[:, 1]\n        \n        #flatten label and prediction tensors\n        pred = pred.contiguous().view(-1)\n        gt = gt.contiguous().view(-1).to(torch.float32)\n        \n        intersection = (pred * gt).sum()                            \n        dice_loss =(2.*intersection + smooth)/(pred.sum() + gt.sum() + smooth)  \n        \n        return -torch.log(dice_loss)\n\nlogDice_loss = DiceLoss()\n\n\ndef plot_DSC(train_DSC, valid_DSC):\n    fig = plt.figure(figsize=(10, 10))\n    plt.plot(train_DSC, '-bx')\n    plt.plot(valid_DSC, '-rx')\n    plt.xlabel('epoch')\n    plt.ylabel('Dice Coefficient')\n    plt.title('Dice Coefficient vs. No. of epochs');\n    \ndef plot_losses(train_losses, valid_losses):\n    fig = plt.figure(figsize=(10, 10))\n    plt.plot(train_losses, '-bx')\n    plt.plot(valid_losses, '-rx')\n    plt.xlabel('epoch')\n    plt.ylabel('- Log Dice loss')\n    plt.legend(['Training', 'Validation'])\n    plt.title('- Log Dice Loss vs. No. of epochs');","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model2 = Vgg16().to(device)\nmodel2.freeze()\n# Set the optimizer \noptimizer = optim.Adam(model2.parameters(), lr=0.0001) \npretrained_vgg16_path = './Vgg16_logdice_lr4_bs32.pt'  \ntrain_losses, valid_losses, train_DSC, train_IoU, train_recall, train_precision,valid_DSC, valid_IoU, valid_recall, valid_precision, best_DSC= train(model2, RGBtrain_loader, RGBvalid_loader, logDice_loss, optimizer,Evaluation, epochs=40, save_path = pretrained_vgg16_path)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_DSC(train_DSC, valid_DSC)\nplt.savefig('./Dice_plot_Vgg16_logdice_lr4_32bs.png')   ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_losses(train_losses, valid_losses)\nplt.savefig('./Losses_plot_Vgg16_logdice_lr4_32bs.png')  ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## **Model: Untrained Vgg16**","metadata":{}},{"cell_type":"code","source":"class Vgg16_np(nn.Module):\n    def __init__(self):\n        super().__init__()\n        # Use a pretrained model\n        network = models.vgg16(pretrained=False)\n        # Replace the classifier\n        self.modified_network = nn.Sequential(*list((*list(network.children())[:-2],\n                                                     nn.Conv2d(512,512, kernel_size = 1),\n                                                     nn.BatchNorm2d(512),\n                                                     nn.ReLU(),\n                                                     nn.Conv2d(512,2, kernel_size = 1),\n                                                     nn.BatchNorm2d(2),\n                                                     nn.ReLU(),\n                                                     nn.Upsample(size=(384,384), mode='bilinear', align_corners=False)))) \n    \n    \n    \n    def forward(self, xb):\n        return self.modified_network(xb)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Use GPU if it is available \nmodel_np = Vgg16_np().to(device)\n\n# Set the optimizer \noptimizer = optim.Adam(model.parameters(), lr=0.001)\n\nUntrained_Vgg16 = './Untrained_Vgg16_DiceScore.pt'\ntrain_losses, valid_losses, train_DSC, train_IoU, train_recall, train_precision,valid_DSC, valid_IoU, valid_recall, valid_precision, best_DSC = train(model_np,RGBtrain_loader,RGBvalid_loader, loss_fn, optimizer, Evaluation, epochs=25, save_path = Untrained_Vgg16)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_DSC(train_DSC, valid_DSC)\nplt.savefig('./Dice_plot_Untrained_Vgg16.png')","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_losses(train_losses, valid_losses)\nplt.savefig('./Losses_plot_Untrained_Vgg16.png')","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## **Model: Pretrained Resnet34**","metadata":{}},{"cell_type":"code","source":"class ResNet34(nn.Module):\n    def __init__(self):\n        super().__init__()\n        # Use a pretrained model\n        network = models.resnet34(pretrained=True)\n        # Replace the classifier\n        self.modified_network = nn.Sequential(*list((*list(network.children())[:-2],\n                                                     nn.Conv2d(512,2, kernel_size = 1),\n                                                     nn.Upsample(size=(384,384), mode='bilinear', align_corners=False)))) \n    \n    \n    \n    def forward(self, xb):\n        return self.modified_network(xb)\n    \n    def freeze(self):\n        # To freeze the CONV layers\n        for param in self.modified_network[:8].parameters():\n            param.require_grad = False\n        for param in self.modified_network[8].parameters():\n            param.require_grad = True\n        for param in self.modified_network[9].parameters():\n            param.require_grad = True\n    \n    def unfreeze(self):\n        # Unfreeze all layers\n        for param in self.modified_network.parameters():\n            param.require_grad = True","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Use GPU of available\nResNet_model2 = ResNet34().to(device)\n\n# checking if the network works, using one batch of the training set\nimg, mask = next(iter(RGBtrain_loader))\nimg, mask = rgb_img.to(device, dtype=torch.float), mask.to(device) \noutput = ResNet_model (img).to(device)\noutput.shape, mask.shape ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Set up cutom optimizer with weight decay\noptimizer = optim.Adam(ResNet_model2.parameters(), lr=0.0001)\n\nResNet_model2.freeze()\nPretrained_ResNet34 = './Pretrained_ResNet34_DiceScore_lr4_bs32.pt'\ntrain_losses, valid_losses, train_DSC, train_IoU, train_recall, train_precision,valid_DSC, valid_IoU, valid_recall, valid_precision, best_DSC = train(ResNet_model2,RGBtrain_loader, RGBvalid_loader, loss_fn, optimizer, Evaluation, epochs=40, save_path = Pretrained_ResNet34) ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def plot_DSC(train_DSC, valid_DSC):\n    fig = plt.figure(figsize=(10, 10))\n    plt.plot(train_DSC, '-bx')\n    plt.plot(valid_DSC, '-rx')\n    plt.xlabel('epoch')\n    plt.ylabel('Dice Coefficient')\n    plt.title('Dice Coefficient vs. No. of epochs');\n    \ndef plot_losses(train_losses, valid_losses):\n    fig = plt.figure(figsize=(10, 10))\n    plt.plot(train_losses, '-bx')\n    plt.plot(valid_losses, '-rx')\n    plt.xlabel('epoch')\n    plt.ylabel('DiceBCE loss')\n    plt.legend(['Training', 'Validation'])\n    plt.title('DiceBCE Loss vs. No. of epochs');","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_DSC(train_DSC, valid_DSC)\nplt.savefig('./Dice_plot_Pretrained_ResNet34_lr4_bs32.png')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_losses(train_losses, valid_losses)\nplt.savefig('./Losses_plot_Pretrained_ResNet34_lr4_bs32.png')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model3 = ResNet34().to(device)\nmodel3.freeze()\n# Set the optimizer \noptimizer = optim.Adam(model3.parameters(), lr=0.0001) \npretrained_resnet34_path = './Resnet34_logdice_lr4_bs32.pt'  \ntrain_losses, valid_losses, train_DSC, train_IoU, train_recall, train_precision,valid_DSC, valid_IoU, valid_recall, valid_precision, best_DSC= train(model3, RGBtrain_loader, RGBvalid_loader, logDice_loss, optimizer,Evaluation, epochs=40, save_path = pretrained_resnet34_path)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_DSC(train_DSC, valid_DSC)\nplt.savefig('./Dice_plot_resnet34_logdice_lr4_32bs.png')   ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_losses(train_losses, valid_losses)\nplt.savefig('./Losses_plot_Pretrained_ResNet34_logdice_lr4_bs32.png')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## **Model: Untrained Resnet34**","metadata":{}},{"cell_type":"code","source":"class ResNet34_np(nn.Module):\n    def __init__(self):\n        super().__init__()\n        # Use a pretrained model\n        self.network = models.resnet34(pretrained=False)\n        # Replace the classifier\n        self.network = nn.Sequential(*list((*list(self.network.children())[:-2],\n                                                     nn.Conv2d(512,2, kernel_size = 1),\n                                                     nn.Upsample(size=(384,384), mode='bilinear', align_corners=False)))) \n    \n    \n    \n    def forward(self, xb):\n        return self.network(xb)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Use GPU if available\nResnet_np = ResNet34_np().to(device)\n\n# Set up cutom optimizer with weight decay\noptimizer = optim.Adam(Resnet_np.parameters(), lr=0.001)\n\nUntrained_ResNet34 = './Untrained_ResNet34_DiceScore.pt'\ntrain_losses, valid_losses, train_DSC, train_IoU, train_recall, train_precision,valid_DSC, valid_IoU, valid_recall, valid_precision, best_DSC = train(Resnet_np,RGBtrain_loader, RGBvalid_loader, loss_fn, optimizer, Evaluation, epochs=25, save_path = Untrained_ResNet34)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_accuracies(train_DSC, valid_DSC)\nplt.savefig('./Dice_plot_Untrained_ResNet34.png')","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_losses(train_losses, valid_losses)\nplt.savefig('./Losses_plot_Untrained_ResNet34.png')","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**A simple Unet**","metadata":{}},{"cell_type":"code","source":"class CloudDataset (Dataset):\n    def __init__(self, red_dir, blue_dir, green_dir, nir_dir, gt_dir):\n        \n\n        # Listing subdirectories\n        # Loop through the files in red folder  \n        # and combine, into a dictionary, the other bands\n        \n        self.files = [self.combine_files(f, green_dir, blue_dir, nir_dir, gt_dir) \n                      for f in red_dir.iterdir() if not f.is_dir()]\n        \n                \n        \n    def combine_files(self, red_file: Path, green_dir, blue_dir, nir_dir, gt_dir):\n        \n        files = {'red': red_file, \n                 'green':green_dir/red_file.name.replace('red', 'green'),\n                 'blue': blue_dir/red_file.name.replace('red', 'blue'), \n                 'nir': nir_dir/red_file.name.replace('red', 'nir'),\n                 'gt': gt_dir/red_file.name.replace('red', 'gt')}\n\n        return files\n    \n    \n    \n    def OpenAsArray(self, idx, invert=False, include_nir=False):\n        \n        raw_rgb=np.stack([np.array(Image.open(self.files[idx]['red'])),\n                          np.array(Image.open(self.files[idx]['green'])),\n                          np.array(Image.open(self.files[idx]['blue']))], axis = 2)\n     \n     \n        if include_nir:\n            nir = np.expand_dims(np.array(Image.open(self.files[idx]['nir'])), axis = 2)\n            raw_rgb = np.concatenate([raw_rgb, nir], axis = 2) \n                          \n        if invert:\n            raw_rgb = raw_rgb.transpose((2, 0, 1))\n    \n    \n        return (raw_rgb / np.iinfo(raw_rgb.dtype).max)\n    \n    \n    \n    \n    def OpenMask(self, idx, add_dims=False):\n        \n        raw_mask=np.array(Image.open(self.files[idx]['gt']))\n        raw_mask = np.where(raw_mask==255, 1, 0)\n        \n        \n        return np.expand_dims(raw_mask, 0) if add_dims else raw_mask\n\n\n\n        \n    def __len__(self):\n        \n        return len(self.files)\n    \n    \n    \n    def __getitem__(self, idx):\n        \n        x = torch.tensor(self.OpenAsArray(idx, invert=True, include_nir=True), dtype=torch.float32)\n        y = torch.tensor(self.OpenMask(idx, add_dims=False), dtype=torch.int64)\n        \n        return x, y\n    \n    \n    \n    def open_as_pil(self, idx):\n        \n        arr = 256 * self.OpenAsArray(idx)\n        \n        return Image.fromarray(arr.astype(np.uint8), 'RGB')  \n    \n    \n    \n    def __repr__(self):\n        \n        s = 'Dataset class with {} files'.format(self.__len__())\n\n        return s","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"data = CloudDataset(red_dir, blue_dir, green_dir, nir_dir, gt_dir) \n\n# splitting the data into train, validation, and test datasets\n\ntrain_size = int(0.75 * len(data))\nvalid_size = int(0.15 * len(data))\ntest_size  = len(data) - train_size - valid_size\nremaining_size = len(data) - train_size \n\ntrain_dataset, remaining_dataset = torch.utils.data.random_split(data, [train_size, remaining_size])\nvalid_dataset, test_dataset      = torch.utils.data.random_split(remaining_dataset, [valid_size, test_size])\n\n\nprint('\\t\\t\\tDataset')\nprint(\"Train data: \\t\\t{}\".format(len(train_dataset)),\n      \"\\nValidation data: \\t{}\".format(len(valid_dataset)),\n     \"\\nTest data: \\t\\t{}\".format(len(test_dataset)))\n\n\n\ntrain_loader = DataLoader(train_dataset, batch_size=12, shuffle=True, num_workers=2)\nvalid_loader = DataLoader(valid_dataset, batch_size=12, shuffle=True, num_workers=2)\ntest_loader  = DataLoader(test_dataset , batch_size=12, shuffle=True, num_workers=2)\n\ndata_iter = iter(valid_loader)\nrgb_img, mask = next(data_iter)\n\nprint('\\n')\nprint('Raw RGB image shape on batch size = {}'.format(rgb_img.size()))\nprint('Cloud Mask shape on batch size    = {}'.format(mask.size()))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class UNet(nn.Module):\n    def __init__(self, in_channels, out_channels):\n        super(UNet, self).__init__()\n            \n        # downsampling part\n        self.DownConv1 = self.ContractBlock(in_channels, 32, 7, 3)\n        self.DownConv2 = self.ContractBlock(32, 64, 3, 1)\n        self.DownConv3 = self.ContractBlock(64, 128, 3, 1)\n            \n        # upsampling part\n        self.UpConv3 = self.ExpandBlock(128, 64, 3, 1)\n        self.UpConv2 = self.ExpandBlock(64*2, 32, 3, 1)\n        self.UpConv1 = self.ExpandBlock(32*2, out_channels, 3, 1)\n        \n    def __call__(self, x):\n         \n        DownConv1 = self.DownConv1(x)\n        DownConv2 = self.DownConv2(DownConv1) \n        DownConv3 = self.DownConv3(DownConv2)   \n        UpConv3   = self.UpConv3 (DownConv3)\n        UpConv2   = self.UpConv2 (torch.cat([UpConv3, DownConv2], 1))\n        UpConv1   = self.UpConv1 (torch.cat([UpConv2, DownConv1], 1))\n        \n        return UpConv1\n        \n        \n    def ContractBlock(self, in_channels, out_channels, kernel_size, padding):\n        \n        contract = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=1, padding=padding),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(),\n        \n            nn.Conv2d(out_channels, out_channels, kernel_size=kernel_size, stride=1, padding=padding),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(),\n        \n            nn.MaxPool2d(kernel_size=3, stride=2, padding=1))\n    \n        return contract\n\n\n\n    def ExpandBlock(self, in_channels, out_channels, kernel_size, padding):\n        \n        expand = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=1, padding=padding),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(),\n        \n            nn.Conv2d(out_channels, out_channels, kernel_size=kernel_size, stride=1, padding=padding),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(),\n        \n            nn.ConvTranspose2d(out_channels, out_channels, kernel_size=3, stride=2, padding=1, output_padding=1) )\n    \n        return expand","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# checking if the network works, using one batch of the training set\n\nUnet_model = UNet(4, 2)\n\ndata_iter = iter(train_loader)\nimg, mask = next(data_iter)\n\noutput = Unet_model (img)\noutput.shape","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Use GPU if available\nUnet_model.to(device)\n\n# Set up cutom optimizer with weight decay\noptimizer = optim.Adam(Unet_model.parameters(), lr=0.001)\n\nUnet_path = './Unet_DiceScore.pt'\ntrain_losses, valid_losses, train_DSC, train_IoU, train_recall, train_precision,valid_DSC, valid_IoU, valid_recall, valid_precision, best_DSC = train(Unet_model,train_loader, valid_loader, loss_fn, optimizer, Evaluation, epochs=30, save_path = Unet_path)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_accuracies(train_DSC, valid_DSC)\nplt.savefig('./Dice_plot_Unet.png')","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_losses(train_losses, valid_losses)\nplt.savefig('./Losses_plot_Unet.png')","metadata":{},"execution_count":null,"outputs":[]}]}