{"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\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\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":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport pathlib, sys, os, random, time\nimport numba, cv2, gc\n\nimport matplotlib.pyplot as plt\n%matplotlib inline\n\nimport warnings\nwarnings.filterwarnings('ignore')\n\nfrom tqdm.notebook import tqdm\n\nimport albumentations as A","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import rasterio\nfrom rasterio.windows import Window","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.utils.data as D\n\nimport torchvision\nfrom torchvision import transforms as T","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def set_seeds(seed=42):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    torch.backends.cudnn.deterministic = True\n\nset_seeds();","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"DATA_PATH = '../input/hubmap-kidney-segmentation'\nEPOCHES = 5\nBATCH_SIZE = 8\nDEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# used for converting the decoded image to rle mask\ndef rle_encode(im):\n    '''\n    im: numpy array, 1 - mask, 0 - background\n    Returns run length as string formated\n    '''\n    pixels = im.flatten(order = 'F')\n    pixels = np.concatenate([[0], pixels, [0]])\n    runs = np.where(pixels[1:] != pixels[:-1])[0] + 1\n    runs[1::2] -= runs[::2]\n    return ' '.join(str(x) for x in runs)\n\ndef rle_decode(mask_rle, shape=(256, 256)):\n    '''\n    mask_rle: run-length as string formated (start length)\n    shape: (height,width) of array to return \n    Returns numpy array, 1 - mask, 0 - background\n\n    '''\n    s = mask_rle.split()\n    starts, lengths = [np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])]\n    starts -= 1\n    ends = starts + lengths\n    img = np.zeros(shape[0]*shape[1], dtype=np.uint8)\n    for lo, hi in zip(starts, ends):\n        img[lo:hi] = 1\n    return img.reshape(shape, order='F')\n\n@numba.njit()\ndef rle_numba(pixels):\n    size = len(pixels)\n    points = []\n    if pixels[0] == 1: points.append(0)\n    flag = True\n    for i in range(1, size):\n        if pixels[i] != pixels[i-1]:\n            if flag:\n                points.append(i+1)\n                flag = False\n            else:\n                points.append(i+1 - points[-1])\n                flag = True\n    if pixels[-1] == 1: points.append(size-points[-1]+1)    \n    return points\n\ndef rle_numba_encode(image):\n    pixels = image.flatten(order = 'F')\n    points = rle_numba(pixels)\n    return ' '.join(str(x) for x in points)\n\ndef make_grid(shape, window=256, min_overlap=32):\n    \"\"\"\n        Return Array of size (N,4), where N - number of tiles,\n        2nd axis represente slices: x1,x2,y1,y2 \n    \"\"\"\n    x, y = shape\n    nx = x // (window - min_overlap) + 1\n    x1 = np.linspace(0, x, num=nx, endpoint=False, dtype=np.int64)\n    x1[-1] = x - window\n    x2 = (x1 + window).clip(0, x)\n    ny = y // (window - min_overlap) + 1\n    y1 = np.linspace(0, y, num=ny, endpoint=False, dtype=np.int64)\n    y1[-1] = y - window\n    y2 = (y1 + window).clip(0, y)\n    slices = np.zeros((nx,ny, 4), dtype=np.int64)\n    \n    for i in range(nx):\n        for j in range(ny):\n            slices[i,j] = x1[i], x2[i], y1[j], y2[j]    \n    return slices.reshape(nx*ny,4)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"identity = rasterio.Affine(1, 0, 0, 0, 1, 0)\n\nclass HubDataset(D.Dataset):\n\n    def __init__(self, root_dir, transform,\n                 window=256, overlap=32, threshold = 100):\n        self.path = pathlib.Path(root_dir)\n        self.overlap = overlap\n        self.window = window\n        self.transform = transform\n        self.csv = pd.read_csv((self.path / 'train.csv').as_posix(),\n                               index_col=[0])\n        self.threshold = threshold\n        \n        self.x, self.y = [], []\n        self.build_slices()\n        self.len = len(self.x)\n        self.as_tensor = T.Compose([\n            T.ToTensor(),\n            T.Normalize([0.625, 0.448, 0.688],\n                        [0.131, 0.177, 0.101]),\n        ])\n        \n    \n    def build_slices(self):\n        self.masks = []\n        self.files = []\n        self.slices = []\n        for i, filename in enumerate(self.csv.index.values):\n            filepath = (self.path /'train'/(filename+'.tiff')).as_posix()\n            self.files.append(filepath)\n            \n            print('Transform', filename)\n            with rasterio.open(filepath, transform = identity) as dataset:\n                self.masks.append(rle_decode(self.csv.loc[filename, 'encoding'], dataset.shape))\n                slices = make_grid(dataset.shape, window=self.window, min_overlap=self.overlap)\n                \n                for slc in tqdm(slices):\n                    x1,x2,y1,y2 = slc\n                    if self.masks[-1][x1:x2,y1:y2].sum() > self.threshold or np.random.randint(100) > 120:\n                        self.slices.append([i,x1,x2,y1,y2])\n                        \n                        image = dataset.read([1,2,3],\n                            window=Window.from_slices((x1,x2),(y1,y2)))\n                        \n#                         if image.std().mean() < 10:\n#                             continue\n                        \n                        # print(image.std().mean(), self.masks[-1][x1:x2,y1:y2].sum())\n                        image = np.moveaxis(image, 0, -1)\n                        self.x.append(image)\n                        self.y.append(self.masks[-1][x1:x2,y1:y2])\n    \n    # get data operation\n    def __getitem__(self, index):\n        image, mask = self.x[index], self.y[index]\n        augments = self.transform(image=image, mask=mask)\n        return self.as_tensor(augments['image']), augments['mask'][None]\n    \n    def __len__(self):\n        \"\"\"\n        Total number of samples in the dataset\n        \"\"\"\n        return self.len\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"WINDOW=1024\nMIN_OVERLAP=32\nNEW_SIZE=256\n\ntrfm = A.Compose([\n    A.Resize(NEW_SIZE,NEW_SIZE),\n    A.HorizontalFlip(p=0.5),\n    A.VerticalFlip(p=0.5),\n    \n    A.OneOf([\n        A.RandomContrast(),\n        A.RandomGamma(),\n        A.RandomBrightness(),\n        A.ColorJitter(brightness=0.07, contrast=0.07,\n                   saturation=0.1, hue=0.1, always_apply=False, p=0.3),\n        ], p=0.3),\n    A.OneOf([\n        A.ElasticTransform(alpha=120, sigma=120 * 0.05, alpha_affine=120 * 0.03),\n        A.GridDistortion(),\n        A.OpticalDistortion(distort_limit=2, shift_limit=0.5),\n        ], p=0.0),\n    A.ShiftScaleRotate(),\n])\n\nds = HubDataset(DATA_PATH, window=WINDOW, overlap=MIN_OVERLAP, transform=trfm)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image, mask = ds[2]\nplt.figure(figsize=(16,8))\nplt.subplot(121)\nplt.imshow(mask[0], cmap='gray')\nplt.subplot(122)\nplt.imshow(image[0]);\n\n_ = rle_numba_encode(mask[0]) # compile function with numba","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"valid_idx, train_idx = [], []\nfor i in range(len(ds)):\n    if ds.slices[i][0] == 7:\n        valid_idx.append(i)\n    else:\n        train_idx.append(i)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_ds = D.Subset(ds, train_idx)\nvalid_ds = D.Subset(ds, valid_idx)\n\n# define training and validation data loaders\nloader = D.DataLoader(\n    train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)\n\nvloader = D.DataLoader(\n    valid_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class ConvBlock(nn.Module):\n    \"\"\"\n    Helper module that consists of a Conv -> BN -> ReLU\n    \"\"\"\n\n    def __init__(self, in_channels, out_channels, padding=1, kernel_size=3, stride=1, with_nonlinearity=True):\n        super().__init__()\n        self.conv = nn.Conv2d(in_channels, out_channels, padding=padding, kernel_size=kernel_size, stride=stride)\n        self.bn = nn.BatchNorm2d(out_channels)\n        self.relu = nn.ReLU()\n        self.with_nonlinearity = with_nonlinearity\n\n    def forward(self, x):\n        x = self.conv(x)\n        x = self.bn(x)\n        if self.with_nonlinearity:\n            x = self.relu(x)\n        return x\n\n\nclass Bridge(nn.Module):\n    \"\"\"\n    This is the middle layer of the UNet which just consists of some\n    \"\"\"\n\n    def __init__(self, in_channels, out_channels):\n        super().__init__()\n        self.bridge = nn.Sequential(\n            ConvBlock(in_channels, out_channels),\n            ConvBlock(out_channels, out_channels)\n        )\n\n    def forward(self, x):\n        return self.bridge(x)\n\n\nclass UpBlockForUNetWithResNet50(nn.Module):\n    \"\"\"\n    Up block that encapsulates one up-sampling step which consists of Upsample -> ConvBlock -> ConvBlock\n    \"\"\"\n\n    def __init__(self, in_channels, out_channels, up_conv_in_channels=None, up_conv_out_channels=None,\n                 upsampling_method=\"conv_transpose\"):\n        super().__init__()\n\n        if up_conv_in_channels == None:\n            up_conv_in_channels = in_channels\n        if up_conv_out_channels == None:\n            up_conv_out_channels = out_channels\n\n        if upsampling_method == \"conv_transpose\":\n            self.upsample = nn.ConvTranspose2d(up_conv_in_channels, up_conv_out_channels, kernel_size=2, stride=2)\n        elif upsampling_method == \"bilinear\":\n            self.upsample = nn.Sequential(\n                nn.Upsample(mode='bilinear', scale_factor=2),\n                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1)\n            )\n        self.conv_block_1 = ConvBlock(in_channels, out_channels)\n        self.conv_block_2 = ConvBlock(out_channels, out_channels)\n\n    def forward(self, up_x, down_x):\n        \"\"\"\n        :param up_x: this is the output from the previous up block\n        :param down_x: this is the output from the down block\n        :return: upsampled feature map\n        \"\"\"\n        x = self.upsample(up_x)\n        x = torch.cat([x, down_x], 1)\n        x = self.conv_block_1(x)\n        x = self.conv_block_2(x)\n        return x\n#---------- load pretrained weights\nmodel1 = torchvision.models.resnet.resnet50(pretrained=False)\nstate_dict = torch.load(\"../input/pretrainedweight/resnet50-19c8e357.pth\")\nmodel1.load_state_dict(state_dict)\nresnet1 = model1\n\nclass UNetWithResnet50Encoder(nn.Module):\n    DEPTH = 6\n\n    def __init__(self, n_classes=1):\n        super().__init__()\n        resnet = resnet1\n        down_blocks = []\n        up_blocks = []\n        self.input_block = nn.Sequential(*list(resnet.children()))[:3]\n        self.input_pool = list(resnet.children())[3]\n        for bottleneck in list(resnet.children()):\n            if isinstance(bottleneck, nn.Sequential):\n                down_blocks.append(bottleneck)\n        self.down_blocks = nn.ModuleList(down_blocks)\n        self.bridge = Bridge(2048, 2048)\n        up_blocks.append(UpBlockForUNetWithResNet50(2048, 1024))\n        up_blocks.append(UpBlockForUNetWithResNet50(1024, 512))\n        up_blocks.append(UpBlockForUNetWithResNet50(512, 256))\n        up_blocks.append(UpBlockForUNetWithResNet50(in_channels=128 + 64, out_channels=128,\n                                                    up_conv_in_channels=256, up_conv_out_channels=128))\n        up_blocks.append(UpBlockForUNetWithResNet50(in_channels=64 + 3, out_channels=64,\n                                                    up_conv_in_channels=128, up_conv_out_channels=64))\n\n        self.up_blocks = nn.ModuleList(up_blocks)\n\n        self.out = nn.Conv2d(64, n_classes, kernel_size=1, stride=1)\n\n    def forward(self, x, with_output_feature_map=False):\n        pre_pools = dict()\n        pre_pools[f\"layer_0\"] = x\n        x = self.input_block(x)\n        pre_pools[f\"layer_1\"] = x\n        x = self.input_pool(x)\n\n        for i, block in enumerate(self.down_blocks, 2):\n            x = block(x)\n            if i == (UNetWithResnet50Encoder.DEPTH - 1):\n                continue\n            pre_pools[f\"layer_{i}\"] = x\n\n        x = self.bridge(x)\n\n        for i, block in enumerate(self.up_blocks, 1):\n            key = f\"layer_{UNetWithResnet50Encoder.DEPTH - 1 - i}\"\n            x = block(x, pre_pools[key])\n        output_feature_map = x\n        x = self.out(x)\n        del pre_pools\n        if with_output_feature_map:\n            return x, output_feature_map\n        else:\n            return x\n\ndef get_model():\n    # model = torchvision.models.segmentation.fcn_resnet50(False)\n    model = UNetWithResnet50Encoder()\n    # pth = torch.load(\"./pretrain-coco-weights-pytorch/fcn_resnet50_coco-1167a1af.pth\")\n    # for key in [\"aux_classifier.0.weight\", \"aux_classifier.1.weight\", \"aux_classifier.1.bias\", \"aux_classifier.1.running_mean\", \"aux_classifier.1.running_var\", \"aux_classifier.1.num_batches_tracked\", \"aux_classifier.4.weight\", \"aux_classifier.4.bias\"]:\n    #     del pth[key]\n    \n    # model.classifier[4] = nn.Conv2d(512, 1, kernel_size=(1, 1), stride=(1, 1))\n    return model\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"@torch.no_grad()\ndef validation(model, loader, loss_fn):\n    losses = []\n    model.eval()\n    for image, target in loader:\n        image, target = image.to(DEVICE), target.float().to(DEVICE)\n        output = model(image)#['out']\n        loss = loss_fn(output, target)\n        losses.append(loss.item())\n        \n    return np.array(losses).mean()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# deeplabv3_resnet101_coco-586e9e4e.pth  fcn_resnet50_coco-1167a1af.pth\n# deeplabv3_resnet50_coco-cd0a2569.pth   resnet101-5d3b4d8f.pth\n# fcn_resnet101_coco-7ecb50ca.pth        resnet50-19c8e357.pth\n# Copy pretrain weight for model to cache dir\n!mkdir -p /root/.cache/torch/hub/checkpoints/\n!cp ../input/pytorch-pretrained-models/resnet50-19c8e357.pth /root/.cache/torch/hub/checkpoints/\n!cp ../input/pretrain-coco-weights-pytorch/fcn_resnet50_coco-1167a1af.pth /root/.cache/torch/hub/checkpoints/","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = get_model()\nmodel.to(DEVICE);\n\noptimizer = torch.optim.AdamW(model.parameters(),\n                  lr=1e-4, weight_decay=1e-3)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"### Table for results\nheader = r'''\n        Train | Valid\nEpoch |  Loss |  Loss | Time, m\n'''\n#          Epoch         metrics            time\nraw_line = '{:6d}' + '\\u2502{:7.3f}'*2 + '\\u2502{:6.2f}'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class SoftDiceLoss(nn.Module):\n    def __init__(self, smooth=1., dims=(-2,-1)):\n\n        super(SoftDiceLoss, self).__init__()\n        self.smooth = smooth\n        self.dims = dims\n    \n    def forward(self, x, y):\n\n        tp = (x * y).sum(self.dims)\n        fp = (x * (1 - y)).sum(self.dims)\n        fn = ((1 - x) * y).sum(self.dims)\n        \n        dc = (2 * tp + self.smooth) / (2 * tp + fp + fn + self.smooth)\n        dc = dc.mean()\n\n        return 1 - dc","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"bce_fn = nn.BCEWithLogitsLoss()\ndice_fn = SoftDiceLoss()\n\ndef loss_fn(y_pred, y_true):\n    bce = bce_fn(y_pred, y_true)\n    dice = dice_fn(y_pred.sigmoid(), y_true)\n    return 0.8*bce+ 0.2*dice\n\nprint(header)\n\nbest_loss = 600\nEPOCHES = 100\nfor epoch in range(1, EPOCHES+1):\n    losses = []\n    start_time = time.time()\n    model.train()\n    for image, target in loader:\n        \n        image, target = image.to(DEVICE), target.float().to(DEVICE)\n        optimizer.zero_grad()\n        output = model(image)#['out']\n        loss = loss_fn(output, target)\n        loss.backward()\n        optimizer.step()\n        losses.append(loss.item())\n    vloss = validation(model, vloader, loss_fn)\n    print(raw_line.format(epoch, np.array(losses).mean(), vloss,\n                              (time.time()-start_time)/60**1))\n    losses = []\n    \n    if vloss < best_loss:\n        best_loss = vloss\n        torch.save(model.state_dict(), './model_best.pth')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# del train set\ndel loader, vloader, train_ds, valid_ds, ds\ngc.collect();","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"WINDOW=1024\nMIN_OVERLAP=32\nNEW_SIZE=256\n\ntrfm = T.Compose([\n    T.ToPILImage(),\n    T.Resize(NEW_SIZE),\n    T.ToTensor(),\n    T.Normalize([0.625, 0.448, 0.688],\n                [0.131, 0.177, 0.101]),\n])\n\n\np = pathlib.Path(DATA_PATH)\n\nsubm = {}\n\nmodel.load_state_dict(torch.load(\"./model_best.pth\"))\nmodel.eval()\n\nfor i, filename in enumerate(p.glob('test/*.tiff')):\n    print(filename)\n    dataset = rasterio.open(filename.as_posix(), transform = identity)\n    slices = make_grid(dataset.shape, window=WINDOW, min_overlap=MIN_OVERLAP)\n    preds = np.zeros(dataset.shape, dtype=np.uint8)\n    for (x1,x2,y1,y2) in slices:\n        image = dataset.read([1,2,3],\n                    window=Window.from_slices((x1,x2),(y1,y2)))\n        image = np.moveaxis(image, 0, -1)\n        image = trfm(image)\n        with torch.no_grad():\n            image = image.to(DEVICE)[None]\n            score = model(image)[0][0]#['out'][0][0]\n\n            score2 = model(torch.flip(image, [0, 3]))#['out']\n            score2 = torch.flip(score2, [3, 0])[0][0]\n\n            score3 = model(torch.flip(image, [1, 2]))#['out']\n            score3 = torch.flip(score3, [2, 1])[0][0]\n\n            score_mean = (score + score2 + score3) / 3.0\n            score_sigmoid = score_mean.sigmoid().cpu().numpy()\n            score_sigmoid = cv2.resize(score_sigmoid, (WINDOW, WINDOW))\n            \n            preds[x1:x2,y1:y2] = (score_sigmoid > 0.5).astype(np.uint8)\n            \n    subm[i] = {'id':filename.stem, 'predicted': rle_numba_encode(preds)}\n    del preds\n    gc.collect();","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"submission = pd.DataFrame.from_dict(subm, orient='index')\nsubmission.to_csv('submission.csv', index=False)","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}