{"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":"#Unet segmentation with Carvana Image Masking dataset","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport os , glob\nimport matplotlib.pyplot as plt\nfrom PIL import Image\nimport cv2\nfrom tqdm import tqdm\n\nimport torch\nfrom torch import Tensor\nfrom torch.autograd import Function\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\n\nfrom keras.preprocessing import image\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\nfrom torch.utils.data import Dataset , DataLoader\nfrom torchvision import transforms , utils , datasets\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')","metadata":{"execution":{"iopub.status.busy":"2021-07-23T00:23:49.886052Z","iopub.execute_input":"2021-07-23T00:23:49.88638Z","iopub.status.idle":"2021-07-23T00:23:49.896752Z","shell.execute_reply.started":"2021-07-23T00:23:49.886352Z","shell.execute_reply":"2021-07-23T00:23:49.893394Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Unzipping the files ","metadata":{}},{"cell_type":"code","source":"!unzip '/kaggle/input/carvana-image-masking-challenge/train_hq.zip' -d './'\n!unzip '/kaggle/input/carvana-image-masking-challenge/train_masks.zip' -d './'\n!unzip '/kaggle/input/carvana-image-masking-challenge/train.zip' -d './'\n!unzip '/kaggle/input/carvana-image-masking-challenge/train_masks.csv.zip' -d './'","metadata":{"execution":{"iopub.status.busy":"2021-07-23T00:23:49.898371Z","iopub.execute_input":"2021-07-23T00:23:49.89888Z","iopub.status.idle":"2021-07-23T00:24:17.656394Z","shell.execute_reply.started":"2021-07-23T00:23:49.898842Z","shell.execute_reply":"2021-07-23T00:24:17.655283Z"},"collapsed":true,"jupyter":{"outputs_hidden":true},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset","metadata":{}},{"cell_type":"code","source":"#Simple Sanity check with csv file and images\nprint('Train Size : ',len(glob.glob('./train/*')))\nprint('Mask Size : ',len(glob.glob('./train_masks/*')))\ndata = pd.read_csv('train_masks.csv')\ndata","metadata":{"execution":{"iopub.status.busy":"2021-07-23T00:24:17.658759Z","iopub.execute_input":"2021-07-23T00:24:17.659071Z","iopub.status.idle":"2021-07-23T00:24:18.121678Z","shell.execute_reply.started":"2021-07-23T00:24:17.65904Z","shell.execute_reply":"2021-07-23T00:24:18.120582Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fig=plt.figure(figsize=(15, 15))\nimg=cv2.imread('./train/d46244bc42ed_04.jpg')\nmask=Image.open('./train_masks/d46244bc42ed_04_mask.gif') #Image from PIL \nprint(img.shape)\nfiles=[img,mask]\nfor i in range(len(files)):\n    plt.subplot(1, 2 , i+1)\n    plt.imshow(files[i])","metadata":{"execution":{"iopub.status.busy":"2021-07-23T00:24:18.123467Z","iopub.execute_input":"2021-07-23T00:24:18.123815Z","iopub.status.idle":"2021-07-23T00:24:19.179115Z","shell.execute_reply.started":"2021-07-23T00:24:18.123779Z","shell.execute_reply":"2021-07-23T00:24:19.178168Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CaravanDataset(Dataset) :\n    \n    def __init__(self,root_dir,img_list,transform_img=None,transform_m=None) :\n        \n        self.root_dir = root_dir\n        self.image_list = img_list\n        self.transform_img = transform_img\n        self.transform_m = transform_m\n    \n    def __len__(self) :\n        return len(self.image_list)\n    \n    def __getitem__(self,idx) :\n        \n        img_name = self.image_list[idx]\n        img_path = self.root_dir +'/'+ str(img_name)\n        \n        mask_name = img_name.split('.')[0] + '_mask.gif'\n        mask_path = self.root_dir +'_masks/'+ str(mask_name)\n        \n        img = cv2.imread(img_path)\n        mask = Image.open(mask_path)    #Object of PIL \n        mask = image.img_to_array(mask)\n        \n\n        \n        if self.transform_img :\n            img = self.transform_img(image=img)['image']\n            \n        if self.transform_m :\n            mask = self.transform_m(image=mask)['image']\n        \n        return img,mask\n    \n#Transformations to be performed on the image and mask\ntransform_img = A.Compose([\n    A.Resize(600,600),\n    A.Normalize(\n    mean=[0.485, 0.456, 0.406],\n    std=[0.229, 0.224, 0.225]),\n    ToTensorV2()])\n\ntransform_mask = A.Compose([\n    A.Resize(600,600),ToTensorV2()])\n\nimages = (pd.read_csv('./train_masks.csv')['img'])\ndataset = CaravanDataset('./train',images,transform_img,transform_mask)","metadata":{"execution":{"iopub.status.busy":"2021-07-23T00:24:19.180385Z","iopub.execute_input":"2021-07-23T00:24:19.180804Z","iopub.status.idle":"2021-07-23T00:24:19.562588Z","shell.execute_reply.started":"2021-07-23T00:24:19.180764Z","shell.execute_reply":"2021-07-23T00:24:19.561719Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"im , mk = dataset[6]\nmyfiles=[im,mk]\nprint(mk.shape)\nprint(im.shape)","metadata":{"execution":{"iopub.status.busy":"2021-07-23T00:24:19.56412Z","iopub.execute_input":"2021-07-23T00:24:19.564527Z","iopub.status.idle":"2021-07-23T00:24:19.650158Z","shell.execute_reply.started":"2021-07-23T00:24:19.564478Z","shell.execute_reply":"2021-07-23T00:24:19.649131Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_size = int(len(images)*0.8)\nval_size = len(images)-train_size\ntrain_set, val_set = torch.utils.data.random_split(dataset, [train_size, val_size]) # We sample 10% of the images as a validation dataset\nprint(train_size,val_size)\n\n\ntrain_dataloader = torch.utils.data.DataLoader(train_set,\n                                          batch_size=4,\n                                          shuffle=True)\n\nvalid_dataloader = torch.utils.data.DataLoader(val_set,\n                                          batch_size=4,\n                                          shuffle=True)\n","metadata":{"execution":{"iopub.status.busy":"2021-07-23T00:24:19.651573Z","iopub.execute_input":"2021-07-23T00:24:19.651923Z","iopub.status.idle":"2021-07-23T00:24:19.673539Z","shell.execute_reply.started":"2021-07-23T00:24:19.651886Z","shell.execute_reply":"2021-07-23T00:24:19.672491Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":" # **Building the UNet Network**\n\nref : https://github.com/milesial/Pytorch-UNet","metadata":{}},{"cell_type":"code","source":"class DoubleConv(nn.Module):\n\n    def __init__(self, in_channels, out_channels, mid_channels=None):\n        \n        super().__init__()\n        \n        if not mid_channels:\n            mid_channels = out_channels\n        \n        self.double_conv = nn.Sequential(\n            \n            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1),\n            nn.BatchNorm2d(mid_channels),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True)\n        )\n\n    def forward(self, x):\n        return self.double_conv(x)\n    \n    \n    \nclass Down(nn.Module):\n\n    def __init__(self, in_channels, out_channels):\n        super().__init__()\n        \n        self.maxpool_conv = nn.Sequential(\n            nn.MaxPool2d(2),\n            DoubleConv(in_channels, out_channels)\n        )\n\n    def forward(self, x):\n        return self.maxpool_conv(x)\n    \n    \nclass Up(nn.Module):\n\n    def __init__(self, in_channels, out_channels, bilinear=True):\n        super().__init__()\n\n        # if bilinear, use the normal convolutions to reduce the number of channels\n        if bilinear:\n            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\n            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)\n        else:\n            self.up = nn.ConvTranspose2d(in_channels , in_channels // 2, kernel_size=2, stride=2)\n            self.conv = DoubleConv(in_channels, out_channels)\n\n\n    def forward(self, x1, x2):\n        x1 = self.up(x1)\n        # input is CHW\n        diffY = x2.size()[2] - x1.size()[2]\n        diffX = x2.size()[3] - x1.size()[3]\n\n        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,\n                        diffY // 2, diffY - diffY // 2])\n        # if you have padding issues, see\n        # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a\n        # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd\n        x = torch.cat([x2, x1], dim=1)\n        return self.conv(x)\n\n\nclass OutConv(nn.Module):\n    def __init__(self, in_channels, out_channels):\n        super(OutConv, self).__init__()\n        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)\n\n    def forward(self, x):\n        return self.conv(x)\n    \n    \n    \nclass UNet(nn.Module):\n    def __init__(self, n_channels, n_classes, bilinear=True):\n        super(UNet, self).__init__()\n        self.n_channels = n_channels\n        self.n_classes = n_classes\n        self.bilinear = bilinear\n\n        self.inc = DoubleConv(n_channels, 64)\n        self.down1 = Down(64, 128)\n        self.down2 = Down(128, 256)\n        self.down3 = Down(256, 512)\n        factor = 2 if bilinear else 1\n        self.down4 = Down(512, 1024 // factor)\n        self.up1 = Up(1024, 512 // factor, bilinear)\n        self.up2 = Up(512, 256 // factor, bilinear)\n        self.up3 = Up(256, 128 // factor, bilinear)\n        self.up4 = Up(128, 64, bilinear)\n        self.outc = OutConv(64, n_classes)\n\n    def forward(self, x):\n        x1 = self.inc(x)\n        x2 = self.down1(x1)\n        x3 = self.down2(x2)\n        x4 = self.down3(x3)\n        x5 = self.down4(x4)\n        x = self.up1(x5, x4)\n        x = self.up2(x, x3)\n        x = self.up3(x, x2)\n        x = self.up4(x, x1)\n        logits = self.outc(x)\n        return logits\n\ndef dice_calc(gt,pred) :\n    pred = torch.sigmoid(pred)\n    pred = ((pred) >= .5).float()\n    dice_score = (2 * (pred * gt).sum()) / ((pred + gt).sum() + 1e-8)\n    \n    return dice_score","metadata":{"execution":{"iopub.status.busy":"2021-07-23T00:24:19.676068Z","iopub.execute_input":"2021-07-23T00:24:19.676454Z","iopub.status.idle":"2021-07-23T00:24:19.701349Z","shell.execute_reply.started":"2021-07-23T00:24:19.676414Z","shell.execute_reply":"2021-07-23T00:24:19.700173Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"                            \nnet = UNet(n_channels=3, n_classes=1, bilinear=True)\nnet.to(device=device)\n\noptimizer = optim.RMSprop(net.parameters(), lr=0.0001, weight_decay=1e-8, momentum=0.9)\nscheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min' if net.n_classes > 1 else 'max', patience=2)\n\nif net.n_classes > 1:\n    criterion = nn.CrossEntropyLoss()\nelse:\n    criterion = nn.BCEWithLogitsLoss()\n","metadata":{"execution":{"iopub.status.busy":"2021-07-23T00:24:19.703002Z","iopub.execute_input":"2021-07-23T00:24:19.703408Z","iopub.status.idle":"2021-07-23T00:24:26.39638Z","shell.execute_reply.started":"2021-07-23T00:24:19.703367Z","shell.execute_reply":"2021-07-23T00:24:26.395557Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train(epoch,epochs,tloader) :\n    net.train()\n    \n    tloader.set_description(f'EPOCH {epoch}')\n    epoch_loss = 0\n    dice_score = 0\n    num_correct = 0    \n    \n    for images , masks in tloader :\n \n        optimizer.zero_grad()\n        images = images.to(device, dtype=torch.float32)\n        masks  = masks.to(device, dtype=torch.float32)\n        mask_pred = net(images)\n\n        loss = criterion(mask_pred,masks)\n        epoch_loss += loss.item()      \n        loss.backward()\n        optimizer.step()\n    \n        running_DS = dice_calc(masks,mask_pred)\n        dice_score += running_DS\n                \n        tloader.set_postfix(loss=loss.item(),accuracy=(running_DS.item()))\n    print(' Train Dice Score Epoch : ',dice_score/len(train_dataloader))\n    \ndef validation(vloader) :\n    net.eval()\n    vloader.set_description('Validation')\n    \n    n_val = len(valid_dataloader)\n    total = 0\n    dice_score = 0\n    num_correct = 0\n    \n    with torch.no_grad():\n        for images ,masks in vloader :\n\n            images = images.to(device)\n            masks  = masks.to(device)\n                \n            mask_pred = net(images)\n            \n            loss = criterion(mask_pred,masks)\n            \n            running_DS = dice_calc(masks,mask_pred)\n            dice_score += running_DS\n            \n            vloader.set_postfix(loss=loss.item(),accuracy=(running_DS.item()))\n        \n    print('Validation Dice Score Epoch : ',dice_score/len(valid_dataloader))\n            ","metadata":{"execution":{"iopub.status.busy":"2021-07-23T00:25:06.070025Z","iopub.execute_input":"2021-07-23T00:25:06.070345Z","iopub.status.idle":"2021-07-23T00:25:06.080234Z","shell.execute_reply.started":"2021-07-23T00:25:06.070315Z","shell.execute_reply":"2021-07-23T00:25:06.07915Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"epochs = 3\n\nfor epoch in range(epochs) :\n    print(epoch+1,'/',epochs)\n    with tqdm(train_dataloader,unit='batch') as tloader : \n        train(epoch,epochs,tloader)\n    \n    with tqdm(valid_dataloader,unit='batch') as vloader:\n        validation(vloader)\n\n","metadata":{"execution":{"iopub.status.busy":"2021-07-23T00:25:09.05847Z","iopub.execute_input":"2021-07-23T00:25:09.058949Z","iopub.status.idle":"2021-07-23T00:31:00.487506Z","shell.execute_reply.started":"2021-07-23T00:25:09.058897Z","shell.execute_reply":"2021-07-23T00:31:00.484617Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Sample Result","metadata":{}},{"cell_type":"code","source":"with torch.no_grad():\n    for images ,masks in valid_dataloader :\n        images = images.to(device)\n        masks  = masks.to(device)\n\n        mask_pred = net(images)\n\n        img = mask_pred.cpu().numpy()\n        masks = masks.cpu().numpy()\n        masks_2 = (masks > 0.5).astype(int)\n        \n        fig, axes = plt.subplots(1, 3, figsize=(15, 15))\n        \n        axes[0].imshow(masks[0][0])\n        axes[0].set_title('Ground Truth Mask')\n        \n        axes[1].imshow(img[0][0])\n        axes[1].set_title('Prababilistic Mask')\n        \n        axes[2].imshow(masks_2[0][0])\n        axes[2].set_title('Probabilistic Mask threshold')\n        break","metadata":{"execution":{"iopub.status.busy":"2021-07-23T00:34:50.602092Z","iopub.execute_input":"2021-07-23T00:34:50.602434Z","iopub.status.idle":"2021-07-23T00:34:51.595341Z","shell.execute_reply.started":"2021-07-23T00:34:50.602404Z","shell.execute_reply":"2021-07-23T00:34:51.59452Z"},"trusted":true},"execution_count":null,"outputs":[]}]}