{"cells":[{"metadata":{},"cell_type":"markdown","source":"## Import all necessary modules"},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport time\nimport copy\nfrom collections import defaultdict\nimport torch\nimport shutil\nimport pandas as pd\nfrom skimage import io, transform\nimport numpy as np\nfrom PIL import Image\nimport matplotlib.pyplot as plt\nfrom torch.utils.data import Dataset, DataLoader, random_split\nfrom torchvision import transforms, utils\nfrom torch import nn\nimport albumentations as A\nfrom albumentations.pytorch import ToTensor\n#import tqdm as tqdm\nfrom tqdm import tqdm as tqdm\n\nfrom albumentations import (HorizontalFlip, ShiftScaleRotate, Normalize, Resize, Compose, GaussNoise)\nimport cv2\n\nfrom torch.autograd import Variable\nfrom torch.nn import Linear, ReLU, CrossEntropyLoss, Sequential, Conv2d, MaxPool2d, Module, Softmax, BatchNorm2d, Dropout\nfrom torch.optim import Adam, SGD\nimport torch.nn.functional as F\nfrom PIL import Image\nfrom torch import nn\nimport zipfile\n\nimport random\n\ndevice = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Extracting Train zip file"},{"metadata":{"trusted":true},"cell_type":"code","source":"from zipfile import ZipFile  \nfile_name = \"../input/data-science-bowl-2018/stage1_train.zip\" \nwith ZipFile(file_name, 'r') as zip: \n    print('Extracting the files') \n    zip.extractall(\"stage1_train\") \n    print('Done!')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Extracting test files"},{"metadata":{"trusted":true},"cell_type":"code","source":"from zipfile import ZipFile \nfile_name = \"../input/data-science-bowl-2018/stage1_test.zip\"\nwith ZipFile(file_name, 'r') as zip: \n    print('Extracting the files') \n    zip.extractall(\"stage1_test\") \n    print('Done!') ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"TRAIN_PATH = 'stage1_train/'","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Load Datasets"},{"metadata":{"trusted":true},"cell_type":"code","source":"#Albumentation\ndef get_train_transform():\n   return A.Compose(\n       [\n        A.Resize(256, 256),\n        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n        A.HorizontalFlip(p=0.25),\n        A.VerticalFlip(p=0.25),\n        ToTensor()\n        ])\n\n#Dataset Loader\nclass LoadDataSet(Dataset):\n        def __init__(self,path, transform=None):\n            self.path = path\n            self.folders = os.listdir(path)\n            self.transforms = get_train_transform()\n        \n        def __len__(self):\n            return len(self.folders)\n              \n        \n        def __getitem__(self,idx):\n            image_folder = os.path.join(self.path,self.folders[idx],'images/')\n            mask_folder = os.path.join(self.path,self.folders[idx],'masks/')\n            image_path = os.path.join(image_folder,os.listdir(image_folder)[0])\n            \n            img = io.imread(image_path)[:,:,:3].astype('float32')\n            img = transform.resize(img,(128,128))\n            \n            mask = self.get_mask(mask_folder, 128, 128 ).astype('float32')\n\n            augmented = self.transforms(image=img, mask=mask)\n            img = augmented['image']\n            mask = augmented['mask']\n            mask = mask[0].permute(2, 0, 1)\n            return (img,mask) \n\n\n        def get_mask(self,mask_folder,IMG_HEIGHT, IMG_WIDTH):\n            mask = np.zeros((IMG_HEIGHT, IMG_WIDTH, 1), dtype=np.bool)\n            for mask_ in os.listdir(mask_folder):\n                    mask_ = io.imread(os.path.join(mask_folder,mask_))\n                    mask_ = transform.resize(mask_, (IMG_HEIGHT, IMG_WIDTH))\n                    mask_ = np.expand_dims(mask_,axis=-1)\n                    mask = np.maximum(mask, mask_)\n              \n            return mask","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dataset = LoadDataSet(TRAIN_PATH, transform=get_train_transform())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#Print the shape of image and mask\nimage, mask = train_dataset.__getitem__(0)\nprint(image.shape)\nprint(mask.shape)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#Print total number of unique images.\ntrain_dataset.__len__()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"mask","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def format_image(img):\n    img = np.array(np.transpose(img, (1,2,0)))\n    mean=np.array((0.485, 0.456, 0.406))\n    std=np.array((0.229, 0.224, 0.225))\n    img  = std * img + mean\n    img = img*255\n    img = img.astype(np.uint8)\n    return img\ndef format_mask(mask):\n    mask = np.squeeze(np.transpose(mask, (1,2,0)))\n    return mask","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Visualize images as well as masks\ndef visualize_dataset(n_images, predict=None):\n  \"\"\"\n  Function to visualize images and masks\n  \"\"\"\n  images = random.sample(range(0, 670), n_images)\n  figure, ax = plt.subplots(nrows=len(images), ncols=2, figsize=(5, 8))\n  print(images)\n  for i in range(0, len(images)):\n    img_no = images[i]\n    image, mask = train_dataset.__getitem__(img_no)\n    image = format_image(image)\n    mask = format_mask(mask)\n    ax[i, 0].imshow(image)\n    ax[i, 1].imshow(mask, interpolation=\"nearest\", cmap=\"gray\")\n    ax[i, 0].set_title(\"Ground Truth Image\")\n    ax[i, 1].set_title(\"Mask\")\n    ax[i, 0].set_axis_off()\n    ax[i, 1].set_axis_off()\n  plt.tight_layout()\n  plt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"visualize_dataset(3)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"## Split train and validation set of split ratio 0.25.\n## ie, 75% of images in train and left 25% of data in valid\nsplit_ratio = 0.25\ntrain_size=int(np.round(train_dataset.__len__()*(1 - split_ratio),0))\nvalid_size=int(np.round(train_dataset.__len__()*split_ratio,0))\n\ntrain_data, valid_data = random_split(train_dataset, [train_size, valid_size])\n\ntrain_loader = DataLoader(dataset=train_data, batch_size=10, shuffle=True)\n\nval_loader = DataLoader(dataset=valid_data, batch_size=10)\n\nprint(\"Length of train and valid datas: {}, {}\".format(len(train_data), len(valid_data)))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"class double_conv(nn.Module):\n    '''\n    Double Convolution layer with both 2 BN and Activation Layer in between\n    Conv2d==>BN==>Activation==>Conv2d==>BN==>Activation\n    '''\n    def __init__(self, in_channel, out_channel):\n        super(double_conv, self).__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(in_channel, out_channel, 3, padding=1),\n            nn.ReLU(inplace=True),\n            nn.BatchNorm2d(out_channel),\n            nn.Conv2d(out_channel, out_channel, 3, padding=1),\n            nn.ReLU(inplace=True),\n            nn.BatchNorm2d(out_channel)\n        )\n    def forward(self, x):\n        x = self.conv(x)\n        return x\n\nclass down_conv(nn.Module):\n  '''\n  A maxpool layer followed by a Double Convolution.\n  MaxPool2d==>double_conv.\n  '''\n  def __init__(self, in_channel, out_channel):\n    super(down_conv, self).__init__()\n    self.down = nn.Sequential(\n        nn.MaxPool2d(2),\n        double_conv(in_channel, out_channel)\n    )\n  def forward(self, x):\n    x = self.down(x)\n    return x\n\nclass up_sample(nn.Module):\n  def __init__(self, in_channel, out_channel):\n    super(up_sample, self).__init__()\n    self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\n    self.double_conv = double_conv(in_channel, out_channel)\n\n  def forward(self, x1, x2):\n      x1 = self.up(x1)\n      x = torch.cat([x1, x2], dim=1)\n      x = self.double_conv(x)\n      return x\n\nclass UNet(nn.Module):\n  '''Main Unet Model'''\n  def __init__(self, in_channel, out_channel):\n    super(UNet, self).__init__()\n    ## DownSampling Block\n    self.down_block1 = double_conv(in_channel, 16)\n    self.down_block2 = down_conv(16, 32)\n    self.down_block3 = down_conv(32, 64)\n    self.down_block4 = down_conv(64, 128)\n    self.down_block5 = down_conv(128, 256)\n    self.down_block6 = down_conv(256, 512)\n    self.down_block7 = down_conv(512, 1024)\n    ## UpSampling Block\n    self.up_block1 = up_sample(1024+512, 512)\n    self.up_block2 = up_sample(512+256, 256)\n    self.up_block3 = up_sample(256+128, 128)\n    self.up_block4 = up_sample(128+64, 64)\n    self.up_block5 = up_sample(64+32, 32)\n    self.up_block6 = up_sample(32+16, 16)\n    self.up_block7 = nn.Conv2d(16, out_channel, 1)\n\n\n  def forward(self, x):\n    #Down\n    x1 = self.down_block1(x)\n    x2 = self.down_block2(x1)\n    x3 = self.down_block3(x2)\n    x4 = self.down_block4(x3)\n    x5 = self.down_block5(x4)\n    x6 = self.down_block6(x5)\n    x7 = self.down_block7(x6)\n    #Up\n    x8 = self.up_block1(x7, x6)\n    x9 = self.up_block2(x8, x5)\n    x10 = self.up_block3(x9, x4)\n    x11 = self.up_block4(x10, x3)\n    x12 = self.up_block5(x11, x2)\n    x13 = self.up_block6(x12, x1)\n    x14 = self.up_block7(x13)\n    out = torch.sigmoid(x14)\n    return out","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Loss Function"},{"metadata":{"trusted":true},"cell_type":"code","source":"class DiceBCELoss(nn.Module):\n    def __init__(self, weight=None, size_average=True):\n        super(DiceBCELoss, self).__init__()\n\n    def forward(self, inputs, targets, smooth=1):\n        \n        #comment out if your model contains a sigmoid or equivalent activation layer\n        inputs = F.sigmoid(inputs)       \n        \n        #flatten label and prediction tensors\n        inputs = inputs.view(-1)\n        targets = targets.view(-1)\n        \n        intersection = (inputs * targets).sum()                            \n        dice_loss = 1 - (2.*intersection + smooth)/(inputs.sum() + targets.sum() + smooth)  \n        BCE = F.binary_cross_entropy(inputs, targets, reduction='mean')\n        Dice_BCE = BCE + dice_loss\n        \n        return Dice_BCE\n\n\nclass DiceLoss(nn.Module):\n    def __init__(self, weight=None, size_average=True):\n        super(DiceLoss, self).__init__()\n\n    def forward(self, inputs, targets, smooth=1):\n        \n        #comment out if your model contains a sigmoid or equivalent activation layer\n        #inputs = F.sigmoid(inputs)       \n        \n        #flatten label and prediction tensors\n        inputs = inputs.view(-1)\n        targets = targets.view(-1)\n        \n        intersection = (inputs * targets).sum()                            \n        dice = (2.*intersection + smooth)/(inputs.sum() + targets.sum() + smooth)  \n        \n        return 1 - dice\n\n\nclass IoU(nn.Module):\n    def __init__(self, weight=None, size_average=True):\n        super(IoU, self).__init__()\n\n    def forward(self, inputs, targets, smooth=1):\n        \n        #comment out if your model contains a sigmoid or equivalent activation layer\n        #inputs = F.sigmoid(inputs)       \n        \n        #flatten label and prediction tensors\n        inputs = inputs.view(-1)\n        targets = targets.view(-1)\n        \n        #intersection is equivalent to True Positive count\n        #union is the mutually inclusive area of all labels & predictions \n        intersection = (inputs * targets).sum()\n        total = (inputs + targets).sum()\n        union = total - intersection \n        \n        IoU = (intersection + smooth)/(union + smooth)\n                \n        return IoU","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = UNet(3,1).cuda()\noptimizer = torch.optim.Adam(model.parameters(),lr = 1e-3)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def save_ckp(state, is_best, checkpoint_path, best_model_path):\n    \"\"\"\n    state: checkpoint we want to save\n    is_best: is this the best checkpoint; min validation loss\n    checkpoint_path: path to save checkpoint\n    best_model_path: path to save best model\n    \"\"\"\n    f_path = checkpoint_path\n    # save checkpoint data to the path given, checkpoint_path\n    torch.save(state, f_path)\n    # if it is a best model, min validation loss\n    if is_best:\n        best_fpath = best_model_path\n        # copy that checkpoint file to best path given, best_model_path\n        shutil.copyfile(f_path, best_fpath)\n\ndef load_ckp(checkpoint_fpath, model, optimizer):\n    \"\"\"\n    checkpoint_path: path to save checkpoint\n    model: model that we want to load checkpoint parameters into       \n    optimizer: optimizer we defined in previous training\n    \"\"\"\n    # load check point\n    checkpoint = torch.load(checkpoint_fpath)\n    # initialize state_dict from checkpoint to model\n    model.load_state_dict(checkpoint['state_dict'])\n    # initialize optimizer from checkpoint to optimizer\n    optimizer.load_state_dict(checkpoint['optimizer'])\n    # initialize valid_loss_min from checkpoint to valid_loss_min\n    valid_loss_min = checkpoint['valid_loss_min']\n    # return model, optimizer, epoch value, min validation loss \n    return model, optimizer, checkpoint['epoch'], valid_loss_min.item()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if not os.path.exists(\"model\"):\n    os.makedirs(\"model\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Train"},{"metadata":{"trusted":true},"cell_type":"code","source":"#from engine import evaluate\ncriterion = DiceLoss()\naccuracy_metric = IoU()\nnum_epochs=20\nvalid_loss_min = np.Inf\n\ncheckpoint_path = 'model/chkpoint_'\nbest_model_path = 'model/bestmodel.pt'\n\ntotal_train_loss = []\ntotal_train_score = []\ntotal_valid_loss = []\ntotal_valid_score = []\n\nlosses_value = 0\nfor epoch in range(num_epochs):\n  \n    train_loss = []\n    train_score = []\n    valid_loss = []\n    valid_score = []\n    #<-----------Training Loop---------------------------->\n    pbar = tqdm(train_loader, desc = 'description')\n    for x_train, y_train in pbar:\n      x_train = torch.autograd.Variable(x_train).cuda()\n      y_train = torch.autograd.Variable(y_train).cuda()\n      optimizer.zero_grad()\n      output = model(x_train)\n      #Loss\n      loss = criterion(output, y_train)\n      losses_value = loss.item()\n      #Score\n      score = accuracy_metric(output,y_train)\n      loss.backward()\n      optimizer.step()\n      train_loss.append(losses_value)\n      train_score.append(score.item())\n      #train_score.append(score)\n      pbar.set_description(f\"Epoch: {epoch+1}, loss: {losses_value}, IoU: {score}\")\n\n    #<---------------Validation Loop---------------------->\n    with torch.no_grad():\n      for image,mask in val_loader:\n        image = torch.autograd.Variable(image).cuda()\n        mask = torch.autograd.Variable(mask).cuda()\n        output = model(image)\n        ## Compute Loss Value.\n        loss = criterion(output, mask)\n        losses_value = loss.item()\n        ## Compute Accuracy Score\n        score = accuracy_metric(output,mask)\n        valid_loss.append(losses_value)\n        valid_score.append(score.item())\n\n    total_train_loss.append(np.mean(train_loss))\n    total_train_score.append(np.mean(train_score))\n    total_valid_loss.append(np.mean(valid_loss))\n    total_valid_score.append(np.mean(valid_score))\n    print(f\"\\n###############Train Loss: {total_train_loss[-1]}, Train IOU: {total_train_score[-1]}###############\")\n    print(f\"###############Valid Loss: {total_valid_loss[-1]}, Valid IOU: {total_valid_score[-1]}###############\")\n\n    #Save best model Checkpoint\n    # create checkpoint variable and add important data\n    checkpoint = {\n        'epoch': epoch + 1,\n        'valid_loss_min': total_valid_loss[-1],\n        'state_dict': model.state_dict(),\n        'optimizer': optimizer.state_dict(),\n    }\n    \n    # save checkpoint\n    save_ckp(checkpoint, False, checkpoint_path, best_model_path)\n    \n    ## TODO: save the model if validation loss has decreased\n    if total_valid_loss[-1] <= valid_loss_min:\n        print('Validation loss decreased ({:.6f} --> {:.6f}).  Saving model ...'.format(valid_loss_min,total_valid_loss[-1]))\n        # save checkpoint as best model\n        save_ckp(checkpoint, True, checkpoint_path, best_model_path)\n        valid_loss_min = total_valid_loss[-1]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import seaborn as sns\n\nplt.figure(1)\nplt.figure(figsize=(15,5))\nsns.set_style(style=\"darkgrid\")\nplt.subplot(1, 2, 1)\nsns.lineplot(x=range(1,num_epochs+1), y=total_train_loss, label=\"Train Loss\")\nsns.lineplot(x=range(1,num_epochs+1), y=total_valid_loss, label=\"Valid Loss\")\nplt.title(\"Loss\")\nplt.xlabel(\"epochs\")\nplt.ylabel(\"DiceLoss\")\n\nplt.subplot(1, 2, 2)\nsns.lineplot(x=range(1,num_epochs+1), y=total_train_score, label=\"Train Score\")\nsns.lineplot(x=range(1,num_epochs+1), y=total_valid_score, label=\"Valid Score\")\nplt.title(\"Score (IoU)\")\nplt.xlabel(\"epochs\")\nplt.ylabel(\"IoU\")\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#loading the saved model\nmodel, optimizer, start_epoch, valid_loss_min = load_ckp(checkpoint_path, model, optimizer)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def visualize_predict(model, n_images):\n  #model = model.eval()\n  figure, ax = plt.subplots(nrows=n_images, ncols=3, figsize=(15, 18))\n  with torch.no_grad():\n    for data,mask in val_loader:\n        data = torch.autograd.Variable(data, volatile=True).cuda()\n        mask = torch.autograd.Variable(mask, volatile=True).cuda()\n        o = model(data)\n        break\n  for img_no in range(0, n_images):\n    tm=o[img_no][0].data.cpu().numpy()\n    img = data[img_no].data.cpu()\n    msk = mask[img_no].data.cpu()\n    img = format_image(img)\n    msk = format_mask(msk)\n    ax[img_no, 0].imshow(img)\n    ax[img_no, 1].imshow(msk, interpolation=\"nearest\", cmap=\"gray\")\n    ax[img_no, 2].imshow(tm, interpolation=\"nearest\", cmap=\"gray\")\n    ax[img_no, 0].set_title(\"Ground Truth Image\")\n    ax[img_no, 1].set_title(\"Ground Truth Mask\")\n    ax[img_no, 2].set_title(\"Predicted Mask\")\n    ax[img_no, 0].set_axis_off()\n    ax[img_no, 1].set_axis_off()\n    ax[img_no, 2].set_axis_off()\n  plt.tight_layout()\n  plt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"visualize_predict(model, 6)","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}