{"cells":[{"metadata":{},"cell_type":"markdown","source":"\n# Hubmap Kaggle Competition"},{"metadata":{},"cell_type":"markdown","source":"## Import needed libraries"},{"metadata":{"trusted":true},"cell_type":"code","source":"%cd ../input/thoplib\n! pip install thop-0.0.31.post2005241907-py3-none-any.whl\n%cd /kaggle/working","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 rasterio\nimport numba, cv2, gc\nimport pathlib, sys, os, random, time\nfrom datetime import datetime\nfrom rasterio.windows import Window\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\n\n%cd ../input/pranet-for-hubmap\nfrom PraNet.lib.PraNet_Res2Net import PraNet\nfrom PraNet.utils.utils import clip_gradient, adjust_lr, AvgMeter\n%cd /kaggle/working","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch, torchvision\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.utils.data as D\nimport torchvision.transforms as T\nfrom torch.autograd import Variable\n\nDEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'  # Check GPU\nprint(DEVICE)","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":{},"cell_type":"markdown","source":"## Hyperparameters Config"},{"metadata":{"trusted":true},"cell_type":"code","source":"class HubmapConfig:\n\n    EPOCHS = 20  # epoch number\n\n    LR = 4e-4  # learning rate\n\n    WEIGHT_DECAY = 1e-4  # weight decay\n\n    BATCH_SIZE = 2 # training batch size\n\n    CLIP = 0.5  # gradient clipping margin\n\n    DECAY_RATE = 0.1  # decay rate of learning rate\n\n    DECAY_EPOCH = 5  # every n epochs decay learning rate\n\n    DATA_PATH = '../input/hubmap-kidney-segmentation'  # path to our dataset\n\n    OUTPUT_PATH = './'  # path to output model\n\nhubmapConfig = HubmapConfig()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Load Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"def rle2mask(rle, mask_shape):\n    \"\"\"\n    decode the rle to image mask\n    :param rle: a column ordered rle\n    :param mask_shape: the shape of mask\n    :return: a mask ndarray, 1-mask, 0-background\n    \"\"\"\n    mask = np.zeros(np.prod(mask_shape), dtype=np.uint8)  # 1d mask array\n    rle = np.array(rle.split()).astype(int)  # rle values to int\n    starts = rle[::2]\n    lengths = rle[1::2]\n    for s, l in zip(starts, lengths):\n        mask[s:s + l] = 1\n    return mask.reshape(np.flip(mask_shape)).T  # flip because of the column-first ordered\n\ndef mask2rle(mask):\n    \"\"\"\n    encode mask to rle\n    :param mask: a mask adarray\n    :return: column ordered rle\n    \"\"\"\n    mask = mask.T.reshape(-1)  # make the mask 1d, column-first\n    mask = np.pad(mask, 1)  # make sure that the 1d mask starts and ends with a 0\n    starts = np.nonzero((~mask[:-1] & mask[1:]))[0]  # start points\n    ends = np.nonzero(mask[:-1] & (~mask[1:]))[0]  # end points\n    rle = np.empty(2 * starts.size, dtype=int)\n    rle[0::2] = starts\n    rle[1::2] = ends - starts\n    rle = \" \".join([str(elem) for elem in rle])\n    return rle\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    slice the image into N tiles, where N is the number of tiles\n    :param shape: a tuple (x, y) represents the input shape\n    :param window: the sliding window size\n    :param min_overlap: overlap between tiles\n    :return: Array of size (N, 4), 2nd axis represent 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)  # generate a array representing x1\n    x1[-1] = x - window  # make the last tile can be a window size tile\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)  # generate a array representing y1\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 HubmapDataset(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(rle2mask(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) > 110:\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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"WINDOW = 1024  # The tile size\nMIN_OVERLAP = 64  # Overlapping between tiles\nNEW_SIZE = 512  # Apply transformation to images and masks, this is the new size","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Define\ntransform = A.Compose([\n    A.Resize(NEW_SIZE,NEW_SIZE),\n    A.OneOf([\n        A.HorizontalFlip(p=0.5),\n        A.VerticalFlip(p=0.5),\n        A.RandomRotate90(p=0.5),\n    ]),\n    A.OneOf([\n        A.RandomContrast(p=0.5),\n        A.RandomGamma(p=0.5),\n        A.RandomBrightness(p=0.5),\n        A.ColorJitter(brightness=0.07, contrast=0.07,\n                    saturation=0.1, hue=0.1, always_apply=False, p=0.5),\n        A.CLAHE(p=0.5),\n    ]),\n\n    A.ShiftScaleRotate(),\n])\n   \nhubmapDataset = HubmapDataset(hubmapConfig.DATA_PATH, window=WINDOW, overlap=MIN_OVERLAP, transform=transform)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Randomly show 3 images\nrandom_list = random.sample(range(0, len(hubmapDataset)), 3)\nfor i in random_list:\n    image, mask = hubmapDataset[i]\n    plt.figure(figsize=(16, 8))\n    plt.subplot(121)\n    plt.imshow(mask[0], cmap='gray')\n    plt.subplot(122)\n    plt.imshow(image[0])\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Split dataset into train and val data\nval_idx, train_idx = [], []\nfor i in range(len(hubmapDataset)):\n    if i % 4 == 0:\n        val_idx.append(i)\n    else:\n        train_idx.append(i)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_data = D.Subset(hubmapDataset, train_idx)\nval_data = D.Subset(hubmapDataset, val_idx)\n\n# Define training and validation data loader\ntrain_loader = D.DataLoader(train_data, batch_size=hubmapConfig.BATCH_SIZE, shuffle=False, num_workers=2)\nprint(len(train_loader))\nval_loader = D.DataLoader(val_data, batch_size=hubmapConfig.BATCH_SIZE, shuffle=False, num_workers=2)\nprint(len(val_loader))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Define train and loss function"},{"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\n\nbce_fn = nn.BCEWithLogitsLoss()\ndice_fn = SoftDiceLoss()\n\ndef loss_fn(y_pred, y_true):\n    y_true = y_true.type_as(y_pred)\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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train(train_loader, val_loader, model, optimizer, epoch, best_pred):\n    # --- model save path ---\n    save_path = hubmapConfig.OUTPUT_PATH\n    os.makedirs(save_path, exist_ok=True)\n    # --- start training ---\n    model.train()\n    # --- multi-scale training ---\n    size_rates = [0.75, 1, 1.25]\n    loss_record2, loss_record3, loss_record4, loss_record5 = AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter()\n    val_loss_record2, val_loss_record3, val_loss_record4, val_loss_record5 = AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter()\n    for i, pack in enumerate(train_loader, start=1):\n        for rate in size_rates:\n            optimizer.zero_grad()\n            # --- data prepare ---\n            images, gts = pack\n            images = images.cuda()\n            gts = gts.float().cuda()\n            # --- rescale ---\n            trainsize = int(round(NEW_SIZE * rate / 32) * 32)\n            if rate != 1:\n                images = F.upsample(images, size=(trainsize, trainsize), mode='bilinear', align_corners=True)\n                gts = F.upsample(gts, size=(trainsize, trainsize), mode='bilinear', align_corners=True)\n            # --- forward ---\n            lateral_map_5, lateral_map_4, lateral_map_3, lateral_map_2 = model(images)\n            # --- loss function ---\n            loss5 = loss_fn(lateral_map_5, gts)\n            loss4 = loss_fn(lateral_map_4, gts)\n            loss3 = loss_fn(lateral_map_3, gts)\n            loss2 = loss_fn(lateral_map_2, gts)\n            loss = loss2 + loss3 + loss4 + loss5\n            # --- backward ---\n            loss.backward()\n            clip_gradient(optimizer, hubmapConfig.CLIP)\n            optimizer.step()\n            # --- recording loss and val loss---\n            if rate == 1:\n                loss_record2.update(loss2.data, hubmapConfig.BATCH_SIZE)\n                loss_record3.update(loss3.data, hubmapConfig.BATCH_SIZE)\n                loss_record4.update(loss4.data, hubmapConfig.BATCH_SIZE)\n                loss_record5.update(loss5.data, hubmapConfig.BATCH_SIZE)\n        # --- train visualization and validation---\n        if i == 1 or i % 100 == 0 or i == len(train_loader):\n            with torch.no_grad():\n                for val_pack in val_loader:\n                    val_images, val_gts = val_pack\n                    val_images = val_images.cuda()\n                    val_gts = val_gts.float().cuda()\n\n                    val_lateral_map_5, val_lateral_map_4, val_lateral_map_3, val_lateral_map_2 = model(val_images)\n                    val_loss5 = loss_fn(val_lateral_map_5, val_gts)\n                    val_loss4 = loss_fn(val_lateral_map_4, val_gts)\n                    val_loss3 = loss_fn(val_lateral_map_3, val_gts)\n                    val_loss2 = loss_fn(val_lateral_map_2, val_gts)\n\n                    val_loss_record2.update(val_loss2.data, 1)\n                    val_loss_record3.update(val_loss3.data, 1)\n                    val_loss_record4.update(val_loss4.data, 1)\n                    val_loss_record5.update(val_loss5.data, 1)\n            print('{} Epoch [{:03d}/{:03d}], Step [{:04d}/{:04d}], '\n                  '[lateral-2: {:.4f}, lateral-3: {:0.4f}, lateral-4: {:0.4f}, lateral-5: {:0.4f}], '\n                  '[val_lateral-2: {:.4f}, val_lateral-3: {:0.4f}, val_lateral-4: {:0.4f}, val_lateral-5: {:0.4f}]'.\n                  format(datetime.now(), epoch, hubmapConfig.EPOCHS, i, len(train_loader),\n                         loss_record2.show(), loss_record3.show(), loss_record4.show(), loss_record5.show(),\n                         val_loss_record2.show(), val_loss_record3.show(), val_loss_record4.show(), val_loss_record5.show()))\n            if val_loss_record2.show() < best_pred[-1]:\n                best_pred.append(val_loss_record2.show())\n                print(\"Best Val Loss: {}\".format(val_loss_record2.show()))\n                torch.save(model.state_dict(), os.path.join(save_path, 'model-best.pth'))\n                print('[Saving Snapshot:]', os.path.join(save_path, 'model-best.pth'))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Build model"},{"metadata":{"trusted":true},"cell_type":"code","source":"# torch.cuda.set_device(0)  # Set your gpu device\nmodel = PraNet().cuda()  # Load model\nmodel.load_state_dict(torch.load('../input/pranet-for-hubmap/PraNet/pretrained/PraNet-19.pth'))\nparams = model.parameters()  # Show model parameters\noptimizer = torch.optim.AdamW(params, lr=hubmapConfig.LR, weight_decay=hubmapConfig.WEIGHT_DECAY)\n# optimizer = torch.optim.Adam(params, lr=hubmapConfig.LR)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(\"#\"*20, \"Start Training\", \"#\"*20)\n\nbest_pred = [10]\nfor epoch in range(1, hubmapConfig.EPOCHS+1):\n    adjust_lr(optimizer, hubmapConfig.LR, epoch, hubmapConfig.DECAY_RATE, hubmapConfig.DECAY_EPOCH)\n    train(train_loader, val_loader, model, optimizer, epoch, best_pred)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"del train_loader, val_loader, train_data, val_data, hubmapDataset\ngc.collect()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"trfm1 = 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\np = pathlib.Path(hubmapConfig.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    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 tqdm(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 = trfm1(image)\n        with torch.no_grad():\n            image = image.to(DEVICE)[None]\n            _, _, _, score = model(image)\n\n            _, _, _, score2 = model(torch.flip(image, [0, 3]))\n            score2 = torch.flip(score2, [3, 0])\n\n            _, _, _, score3 = model(torch.flip(image, [1, 2]))\n            score3 = torch.flip(score3, [2, 1])\n\n            score_mean = (score + score2 + score3) / 3.0\n            score_sigmoid = score_mean.sigmoid()\n            score_sigmoid= F.upsample(score_sigmoid, size=(WINDOW, WINDOW), mode='bilinear', align_corners=False)\n            score_sigmoid = score_sigmoid.cpu().numpy()\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}