{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!mkdir /kaggle/temp/\n# !ls /kaggle/temp/","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-output":true,"trusted":true,"collapsed":true},"cell_type":"code","source":"!unzip ../input/carvana-image-masking-challenge/train.zip -d /kaggle/temp/\n!unzip ../input/carvana-image-masking-challenge/train_masks.zip -d /kaggle/temp/","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms, datasets, models\nfrom PIL import Image","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class CarvanaDataset(Dataset):\n    \"\"\"\n    Loads the Camvid Dataset\n    \n    Splits can be 'train', 'val', test'\n    \"\"\"\n    def __init__(self, root_dir='/kaggle/temp/', split='train', transform=None, \n                 scale=0.5, mkt=None):\n        self.split = split\n        self.scale = scale\n        self.mkts = mkt\n        self.transforms = transform\n        self.image_dir = os.path.join(root_dir, split)\n        self.label_dir = os.path.join(root_dir, split+'_masks')\n        \n        self.image_file_names = os.listdir(self.image_dir)\n        self.label_file_names = [fn[:-4] + '_mask.gif' for fn in self.image_file_names]\n    def __len__(self):\n        return len(self.image_file_names)\n    \n    def __getitem__(self, idx):\n        image_name = self.image_file_names[idx]\n        label_name = self.label_file_names[idx]\n        \n        image = Image.open(os.path.join(self.image_dir, image_name))\n        mask = Image.open(os.path.join(self.label_dir, label_name))\n        w, h = image.size\n        newW, newH = int(self.scale * w), int(self.scale * h)\n        image = image.resize((newW, newH))\n        mask = mask.resize((newW, newH))\n        \n        assert image.size == mask.size, \\\n            f'Image and mask {idx} should be the same size, but are {img.size} and {mask.size}'\n\n        if self.transforms:\n            image = self.transforms(image)\n        if self.mkts:\n            mask = self.mkts(mask)\n        \n        return [image.type(torch.float32), \n                mask.type(torch.float32)]\n        ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"trans = transforms.Compose([\n  transforms.ToTensor(),\n  transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # imagenet\n])\nmask_trans = transforms.Compose([\n  transforms.ToTensor(),\n])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset = CarvanaDataset(scale=0.34, transform=trans, mkt = mask_trans)\nlengths = [int(len(dataset)*0.8), int(len(dataset)*0.2) + 1]\nlengths","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"trainset, valset = torch.utils.data.random_split(dataset, lengths)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image_datasets = {\n  'train': trainset, 'val': valset\n}\n\nbatch_size = 8\n\ndataloaders = {\n  'train': DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=0),\n  'val': DataLoader(valset, batch_size=batch_size, shuffle=True, num_workers=0)\n}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import matplotlib.pyplot as plt","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torchvision.utils\n\ndef reverse_transform(inp):\n  inp = inp.numpy().transpose((1, 2, 0))\n  mean = np.array([0.485, 0.456, 0.406])\n  std = np.array([0.229, 0.224, 0.225])\n  inp = std * inp + mean\n  inp = np.clip(inp, 0, 1)\n  inp = (inp * 255).astype(np.uint8)\n\n  return inp\n\n# Get a batch of training data\ninputs, masks = next(iter(dataloaders['train']))\n\nprint(inputs.shape, masks.shape)\n\nplt.imshow(reverse_transform(inputs[3]))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from torch import nn, optim\nfrom torch.nn import functional as F","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class DoubleConv(nn.Module):\n    \"\"\"(convolution => [BN] => ReLU) * 2\"\"\"\n\n    def __init__(self, in_channels, out_channels, mid_channels=None):\n        super().__init__()\n        if not mid_channels:\n            mid_channels = out_channels\n        self.double_conv = nn.Sequential(\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\nclass Down(nn.Module):\n    \"\"\"Downscaling with maxpool then double conv\"\"\"\n\n    def __init__(self, in_channels, out_channels):\n        super().__init__()\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    \"\"\"Upscaling then double conv\"\"\"\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)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class 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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torch.nn as nn\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint('device', device)\n\nmodel = UNet(3, 1) #6 = num classes\nmodel = model.to(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch.nn as nn\n\ndef dice_loss(pred, target, smooth = 1.):\n    pred = pred.contiguous()\n    target = target.contiguous()    \n\n    intersection = (pred * target).sum(dim=2).sum(dim=2)\n    \n    loss = (1 - ((2. * intersection + smooth) / (pred.sum(dim=2).sum(dim=2) + target.sum(dim=2).sum(dim=2) + smooth)))\n    \n    return loss.mean()\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from collections import defaultdict\nimport torch.nn.functional as F\n\ncheckpoint_path = \"checkpoint.pth\"\n\ndef calc_loss(pred, target, metrics, bce_weight=0.5):\n    bce = F.binary_cross_entropy_with_logits(pred, target)\n\n    pred = torch.sigmoid(pred)\n    dice = dice_loss(pred, target)\n\n    loss = bce * bce_weight + dice * (1 - bce_weight)\n\n    metrics['bce'] += bce.data.cpu().numpy() * target.size(0)\n    metrics['dice'] += dice.data.cpu().numpy() * target.size(0)\n    metrics['loss'] += loss.data.cpu().numpy() * target.size(0)\n\n    return loss\n\ndef print_metrics(metrics, epoch_samples, phase):\n    outputs = []\n    for k in metrics.keys():\n        outputs.append(\"{}: {:4f}\".format(k, metrics[k] / epoch_samples))\n\n    print(\"{}: {}\".format(phase, \", \".join(outputs)))\n\ndef train_model(model, optimizer, scheduler, num_epochs=25):\n    best_loss = 1e10\n\n    for epoch in range(num_epochs):\n        print('Epoch {}/{}'.format(epoch, num_epochs - 1))\n        print('-' * 10)\n\n        since = time.time()\n\n        # Each epoch has a training and validation phase\n        for phase in ['train', 'val']:\n            if phase == 'train':\n                model.train()  # Set model to training mode\n            else:\n                model.eval()   # Set model to evaluate mode\n\n            metrics = defaultdict(float)\n            epoch_samples = 0\n\n            for inputs, labels in dataloaders[phase]:\n                inputs = inputs.to(device)\n                labels = labels.to(device)\n\n                # zero the parameter gradients\n                optimizer.zero_grad()\n\n                # forward\n                # track history if only in train\n                with torch.set_grad_enabled(phase == 'train'):\n                    outputs = model(inputs)\n                    loss = calc_loss(outputs, labels, metrics)\n\n                    # backward + optimize only if in training phase\n                    if phase == 'train':\n                        loss.backward()\n                        optimizer.step()\n\n                # statistics\n                epoch_samples += inputs.size(0)\n\n            print_metrics(metrics, epoch_samples, phase)\n            epoch_loss = metrics['loss'] / epoch_samples\n\n            if phase == 'train':\n              scheduler.step()\n              for param_group in optimizer.param_groups:\n                  print(\"LR\", param_group['lr'])\n\n            # save the model weights\n            if phase == 'val' and epoch_loss < best_loss:\n                print(f\"saving best model to {checkpoint_path}\")\n                best_loss = epoch_loss\n                torch.save(model.state_dict(), checkpoint_path)\n\n        time_elapsed = time.time() - since\n        print('{:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))\n\n    print('Best val loss: {:4f}'.format(best_loss))\n\n    # load best model weights\n    model.load_state_dict(torch.load(checkpoint_path))\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torch.optim as optim\nfrom torch.optim import lr_scheduler\nimport time\n\nnum_class = 1\nmodel = UNet(3, num_class).to(device)\n\n# freeze backbone layers\n# for l in model.base_layers:\n#   for param in l.parameters():\n#     param.requires_grad = False\n\noptimizer_ft = optim.Adam(model.parameters(), lr=0.000002)\n\nexp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=8, gamma=0.1)\n\nmodel = train_model(model, optimizer_ft, exp_lr_scheduler, num_epochs=15)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.save(model.state_dict(), \"model.pth\")","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}