{"cells":[{"metadata":{},"cell_type":"markdown","source":" This will be a quick and dirty kernel showing how to get started on segmenting nuclei using a neural network in Pytorch.The architecture used is the so-called U-Net, which is very common for image segmentation problems such as this. \n<br>\n<br>\n    You can understand more about unet from [https://arxiv.org/abs/1505.04597]. Also implementation of same using keras is described in my [blog](http://https://medium.com/@arunm8489/an-overview-on-u-net-architecture-d6caabf7caa4) post."},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","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\nfrom albumentations import (HorizontalFlip, ShiftScaleRotate, Normalize, Resize, Compose, GaussNoise)\nimport cv2\nfrom albumentations.pytorch import ToTensor\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\ndevice = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#! mkdir /kaggle/working/data\n# ! unzip /kaggle/input/data-science-bowl-2018/stage1_train.zip -d /kaggle/working/data","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if not os.path.exists('/kaggle/working/data'):\n    os.mkdir('/kaggle/working/data')\n    \nzip_path = '/kaggle/input/data-science-bowl-2018/stage1_train.zip'\nwith zipfile.ZipFile(zip_path, 'r') as zip_ref:\n    zip_ref.extractall('/kaggle/working/data')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Loading Data\n\nNow we will load the data by using custom data loader from pytorch along with small preprocessing such as horizontal fliping and normalization."},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_transforms(mean, std):\n            list_transforms = []\n            \n            list_transforms.extend(\n                    [\n                HorizontalFlip(p=0.5), # only horizontal flip as of now\n                    ])\n            list_transforms.extend(\n                    [\n            Normalize(mean=mean, std=std, p=1),\n            ToTensor(),\n                    ])\n            list_trfms = Compose(list_transforms)\n            return list_trfms\n\n\n\n\n\nclass Nuclie_data(Dataset):\n        def __init__(self,path):\n            self.path = path\n            self.folders = os.listdir(path)\n            self.transforms = get_transforms(0.5, 0.5)\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\n            \n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#loading the data\nbase_dir = '/kaggle/working/data/'\ndata = Nuclie_data(base_dir)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# print out some sample data\nprint(data.__len__())\ndata.__getitem__(0)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Now we will just check the shape of images and masks loaded"},{"metadata":{"trusted":true},"cell_type":"code","source":"for img,msk in data:\n  print(img.shape)\n  print(msk.shape)\n  break","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# some utility functions\ndef mask_convert(mask):\n    mask = mask.clone().cpu().detach().numpy()\n    mask = mask.transpose((1,2,0))\n    std = np.array((0.5))\n    mean = np.array((0.5))\n    mask  = std * mask + mean\n    mask = mask.clip(0,1)\n    mask = np.squeeze(mask)\n    return mask\n\n# converting tensor to image\ndef image_convert(image):\n    image = image.clone().cpu().numpy()\n    image = image.transpose((1,2,0))\n    std = np.array((0.5,0.5,0.5))\n    mean = np.array((0.5,0.5,0.5))\n    image  = std * image + mean\n    image = image.clip(0,1)\n    image = (image * 255).astype(np.uint8)\n    return image\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def plot_img(no_):\n    iter_ = iter(train_loader)\n    images,masks = next(iter_)\n    images = images.to(device)\n    masks = masks.to(device)\n    plt.figure(figsize=(10,6))\n    for idx in range(0,no_):\n         image = image_convert(images[idx])\n         plt.subplot(2,no_,idx+1)\n         plt.title('image')\n         plt.imshow(image)\n    for idx in range(0,no_):\n         mask = mask_convert(masks[idx])\n         plt.subplot(2,no_,idx+no_+1)\n         plt.title('mask')\n         plt.imshow(mask,cmap='gray')\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We will use 580 images for train and 90 images for validation"},{"metadata":{"trusted":true},"cell_type":"code","source":"# splitting to trainset and validation set and loading the data with batch size of 10\n\ntrainset, valset = random_split(data, [580, 90])\n\ntrain_loader = torch.utils.data.DataLoader(dataset=trainset, batch_size=10, shuffle=True)\n\nval_loader = torch.utils.data.DataLoader(dataset=valset, batch_size=10)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# we will try visualizing images and corresponding masks\nplot_img(5)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def double_conv(in_channels, out_channels):\n    return nn.Sequential(\n        nn.Conv2d(in_channels, out_channels, 3, padding=1),\n        nn.ReLU(inplace=True),\n        nn.Conv2d(out_channels, out_channels, 3, padding=1),\n        nn.ReLU(inplace=True)\n    )\n\nclass Unet(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.dblock1 = double_conv(3, 64)\n        self.dblock2 = double_conv(64, 128)\n        self.dblock3 = double_conv(128,256)\n        self.dblock4 = double_conv(256,512)\n\n        self.pool = nn.MaxPool2d(2)\n        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) \n\n        self.dblock5 = double_conv(256 + 512, 256)\n        self.dblock6 = double_conv(128 + 256, 128)\n        self.dblock7 = double_conv(128 + 64, 64)\n\n        self.last_layer = nn.Conv2d(64,1,1)\n        self.sigmoid = nn.Sigmoid()\n\n    def forward(self,x):\n        conv1 = self.dblock1(x)\n        x = self.pool(conv1)\n\n        conv2 = self.dblock2(x)\n        x = self.pool(conv2)\n\n        conv3 = self.dblock3(x)\n        x = self.pool(conv3)\n\n        conv4 = self.dblock4(x)\n        \n        x = self.upsample(conv4)\n\n        x = torch.cat([x, conv3], dim=1)\n\n        x = self.dblock5(x)\n        x = self.upsample(x)\n        x = torch.cat([x, conv2], dim=1)\n        \n        x = self.dblock6(x)\n        x = self.upsample(x)\n        x = torch.cat([x, conv1], dim=1)\n        \n        x = self.dblock7(x)\n\n        out = self.last_layer(x)\n        # out = self.sigmoid(x)\n        return out\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = Unet().to(device)\nmodel","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Loss function and IOU metric\n\n\nThis loss combines Dice loss with the standard binary cross-entropy (BCE) loss that is generally the default for segmentation models. Combined the two methods allows for some diversity in the loss, while benefitting from the stability of BCE.Similarly we use intersection over union as a metric"},{"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        bce_weight = 0.5\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        loss_final = BCE * bce_weight + dice_loss * (1 - bce_weight)\n        return loss_final\n    \n## IOU computation\ndef iou_(y_pred,y):\n    inputs = y_pred.reshape(-1)\n    targets = y.reshape(-1)\n    intersection = (inputs * targets).sum()\n    total = (inputs + targets).sum()\n    union = total - intersection \n    smooth = 1    \n    iou = (intersection + smooth)/(union + smooth)\n    return iou\n\n\n\ndef iou_batch(y_pred,y):\n    '''computes mean iou for a batch of ground truth masks and predicted masks'''\n    ious = []\n    y_pred = F.sigmoid(y_pred)\n    y_pred = y_pred.clone().cpu().detach().numpy()\n    y = y.clone().cpu().detach().numpy() \n    \n    for pred, label in zip(y_pred, y):\n        ious.append(iou_(pred, label))\n    iou = np.nanmean(ious)\n    return iou    ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Now we will write some utility functions to save and load best model on training"},{"metadata":{"trusted":true},"cell_type":"code","source":"#ref https://towardsdatascience.com/how-to-save-and-load-a-model-in-pytorch-with-a-complete-example-c2920e617dee\n\ndef 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    # save checkpoint data to the path given, checkpoint_path\n    torch.save(state, checkpoint_path)\n    # if it is a best model, min validation loss\n    if is_best:\n        # copy that checkpoint file to best path given, best_model_path\n        shutil.copyfile(checkpoint_path, best_model_path)\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":{},"cell_type":"markdown","source":"## Model training"},{"metadata":{"trusted":true},"cell_type":"code","source":"checkpoint_path = '/kaggle/working/chkpoint_'\nbest_model_path = '/kaggle/working/bestmodel.pt'\nepochs = 25\ncriterion = DiceBCELoss()\nlearning_rate = 1e-3\noptimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\nvalid_loss_min = 3.95275\n\n\ntrain_loss,val_loss = [],[]\ntrain_iou,val_iou = [],[]\n\nfor epoch in range(epochs):\n    print('Epoch {}/{}'.format(epoch + 1, epochs))\n    start_time = time.time()\n     \n\n    \n    running_train_loss = []\n    running_train_score = []\n    for image,mask in train_loader: \n            image = image.to(device,dtype=torch.float)\n            mask = mask.to(device,dtype=torch.float)\n            \n            pred_mask = model.forward(image) # forward propogation\n            loss = criterion(pred_mask,mask)\n            score = iou_batch(pred_mask,mask)\n            optimizer.zero_grad() # setting gradient to zero\n            loss.backward()\n            optimizer.step()\n            running_train_loss.append(loss.item())\n            running_train_score.append(score)\n                              \n\n    else:           \n        running_val_loss = []\n        running_val_score = []\n        with torch.no_grad():\n            for image,mask in val_loader:\n                    image = image.to(device,dtype=torch.float)\n                    mask = mask.to(device,dtype=torch.float)                            \n                    pred_mask = model.forward(image)\n                    loss = criterion(pred_mask,mask)\n                    score = iou_batch(pred_mask,mask)\n                    running_val_loss.append(loss.item())\n                    running_val_score.append(score)\n\n                                    \n    \n    epoch_train_loss,epoch_train_score = np.mean(running_train_loss) ,np.mean(running_train_score)\n    print('Train loss : {} iou : {}'.format(epoch_train_loss,epoch_train_score))                       \n    train_loss.append(epoch_train_loss)\n    train_iou.append(epoch_train_score)\n    \n    epoch_val_loss,epoch_val_score = np.mean(running_val_loss),np.mean(running_val_score)\n    print('Validation loss : {} iou : {}'.format(epoch_val_loss,epoch_val_score))                                \n    val_loss.append(epoch_val_loss)\n    val_iou.append(epoch_val_score)\n    \n    # create checkpoint variable and add important data\n    checkpoint = {\n            'epoch': epoch + 1,\n            'valid_loss_min': epoch_val_loss,\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    ## TODO: save the model if validation loss has decreased\n    if epoch_val_loss <= valid_loss_min:\n            print('Validation loss decreased ({:.6f} --> {:.6f}).  Saving model ...'.format(valid_loss_min,epoch_val_loss))\n            # save checkpoint as best model\n            save_ckp(checkpoint, True, checkpoint_path, best_model_path)\n            valid_loss_min = epoch_val_loss\n    \n    time_elapsed = time.time() - start_time\n    print('{:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))\n                        \n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Plotting Loss and IOU Scores"},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.figure(figsize=(20,10))\nplt.subplot(1,2,1)\nplt.plot(train_loss,label='train_loss')\nplt.plot(val_loss,label='val_loss')\nplt.legend()\nplt.title('Loss Plot')\nplt.subplot(1,2,2)\nplt.plot(train_iou,label='train_iou')\nplt.plot(val_iou,label='val_iou')\nplt.legend()\nplt.title('IOU Plot')\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":{},"cell_type":"markdown","source":"## Visualisation of actual and predicted mask"},{"metadata":{"trusted":true},"cell_type":"code","source":"\niter_ = iter(val_loader)\nimage,mask = next(iter_)\nimage = image.to(device,dtype=torch.float)\nmask = mask.to(device,dtype=torch.float)\ny_pred = model.forward(image)\n\n\nplt.figure(figsize=(20,15))\nfor i in range(0,5):\n    plt.subplot(3,5,i+1)\n    plt.title('Actual image')\n    plt.imshow(image_convert(image[i]))\nfor i in range(0,5):\n    plt.subplot(3,5,i+5+1)\n    plt.title('Actual mask')\n    plt.imshow(mask_convert(mask[i]),cmap='gray')\nfor i in range(0,5):\n    plt.subplot(3,5,i+10+1)\n    plt.title('Predicted mask')\n    plt.imshow(mask_convert(y_pred[i]),cmap='gray')\nplt.show()\n    \n    ","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}