{"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":"import sys\nsys.path.append(\"../input/segmentation-models-pytorch-install\")\nimport segmentation_models_pytorch as smp\n\nfrom torch.utils.data import Dataset, DataLoader\nimport cv2\nfrom pathlib import Path\nimport rasterio\nfrom rasterio.windows import Window\nimport pandas as pd\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport numpy as np\nimport torch\nfrom tqdm import tqdm\nimport math\nimport gc\nfrom typing import List, Dict\nfrom dataclasses import dataclass","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Unpack weights","metadata":{}},{"cell_type":"code","source":"!ls -l ../input/hubmap-folds-2\n!mkdir model_1\n!tar -xvf ../input/hubmap-folds-2/1024_avg_last.tar -C model_1\n!mkdir model_2\n!tar -xvf ../input/hubmap-folds-2/1024_512_pseudo_v1_avg.tar -C model_2","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Main pipline config","metadata":{}},{"cell_type":"code","source":"BATCH_SIZE = 1\nNUM_WORKERS = 0\nCROP_SIZE = 1024 * 4\nSTEP = 1024 * 2\nTHR = 0.5\nVOTE = 0\nSKIP_BACKGROUND = True\nSKIP_COMMIT = True\nRAMKA = 512\n\ndf_sub = pd.read_csv(\"../input/hubmap-kidney-segmentation/sample_submission.csv\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model configs","metadata":{}},{"cell_type":"code","source":"@dataclass\nclass ModelConfig:\n    encoder: str\n    weights_path: List[str]\n    img_size: int\n    weight_blend: float\n\nmy_models = [\n    ModelConfig(\n        \"resnet34\", \n        [\n            \"./model_1/fold0_avg_0.9238.pth\",\n            \"./model_1/fold1_avg_0.9162.pth\",\n            \"./model_1/fold2_avg_0.9303.pth\",\n            \"./model_1/fold3_avg_0.9336.pth\",\n            \"./model_1/fold4_avg_0.9407.pth\",\n        ], \n        4096, \n        0.35\n    ),\n    ModelConfig(\n        \"timm-efficientnet-b3\", \n        [\n            f'../input/hubmap5/exp_25_fold_0.pt',\n            f'../input/hubmap5/exp_24_fold_1.pt',\n            f'../input/hubmap5/exp_23_fold_2.pt',\n            f'../input/hubmap5/exp_22_fold_3.pt',\n            f'../input/hubmap5/exp_21_fold_4.pt',\n        ],\n        1024, \n        0.3\n    ),\n    ModelConfig(\n        \"se_resnext50_32x4d\", \n        [\n            \"./model_2/fold0_avg_0.9457.pth\",\n            \"./model_2/fold1_avg_0.9566.pth\",\n            \"./model_2/fold2_avg_0.9379.pth\",\n            \"./model_2/fold3_avg_0.9095.pth\",\n            \"./model_2/fold4_avg_0.9338.pth\",\n        ],\n        2048, \n        0.35\n    ),\n]\n\n# choosing reference image size for all models in enssamble\nTRAIN_IMG_SIZE = max([cfg.img_size for cfg in my_models])\nTRAIN_IMG_SIZE","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# import sys\n# sys.path.append(\"../input/segmentation-models-pytorch-install\")\n# import segmentation_models_pytorch as smp\n# import torch\n# model = smp.FPN(\"timm-efficientnet-b1\", upsampling=1, encoder_weights=None).cuda()\n# x = torch.rand(1, 3, 1344, 1344).cuda()\n# out = model(x)\n# print(out.shape)","metadata":{"scrolled":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Helping functions","metadata":{}},{"cell_type":"code","source":"def rle_encode_less_memory(img):\n    #the image should be transposed\n    pixels = img.T.flatten()\n    \n    # This simplified method requires first and last pixel to be zero\n    pixels[0] = 0\n    pixels[-1] = 0\n    runs = np.where(pixels[1:] != pixels[:-1])[0] + 2\n    runs[1::2] -= runs[::2]\n    \n    return ' '.join(str(x) for x in runs)\n\ndef valid_transform(img_size):\n    return A.Compose(\n        [\n            A.Resize(img_size, img_size),\n            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225],),\n            ToTensorV2(),\n        ]\n    )\n\ndef read_from_layers(layers, window):\n    if len(layers) == 1:\n        return np.stack([layers[0].read(x, window=window) for x in [1, 2, 3]], 2)\n    else:\n        return np.stack([layers[x].read(1, window=window) for x in range(3)], 2)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Pytorch dataset class","metadata":{}},{"cell_type":"code","source":"def _check_background(img, crop_size) -> bool:\n    s_th = 40  # saturation blancking threshold\n    p_th = 1000 * (crop_size // 256) ** 2 \n    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)\n    _, ss, _ = cv2.split(hsv)\n    \n    background = False if (ss > s_th).sum() <= p_th or img.sum() <= p_th else True\n    return background\n\n# Dataset class returns multiple instances of image with different sizes\nclass SingleTiffDataset(Dataset):\n    def __init__(self, tiff_path, all_img_sizes, crop_size=1024, step=512):\n        self.crop_size = crop_size\n        self.all_img_sizes = all_img_sizes\n\n        self.step = step\n        dataset = rasterio.open(tiff_path, num_threads=\"all_cpus\")\n        self.h = dataset.height\n        self.w = dataset.width\n        self.row_count = 1 + math.ceil((self.h - self.crop_size) / self.step)\n        self.col_count = 1 + math.ceil((self.w - self.crop_size) / self.step)\n\n        self.layers = []\n        if dataset.count != 3:\n            subdatasets = dataset.subdatasets\n            if len(subdatasets) > 0:\n                for i, subdataset in enumerate(subdatasets, 0):\n                    self.layers.append(\n                        rasterio.open(subdataset, num_threads=\"all_cpus\")\n                    )\n        else:\n            self.layers.append(rasterio.open(tiff_path, num_threads=\"all_cpus\"))\n\n    def __len__(self):\n        return self.row_count * self.col_count\n\n    def __getitem__(self, idx):\n        y = (idx // self.col_count) * self.step\n        x = (idx % self.col_count) * self.step\n        if x + self.crop_size > self.w:\n            x = self.w - self.crop_size\n        if y + self.crop_size > self.h:\n            y = self.h - self.crop_size\n        window = Window(x, y, self.crop_size, self.crop_size)\n        img = read_from_layers(self.layers, window=window)\n\n        # forming batch with different image sizes as keys\n        transormed_imgs = {img_size: valid_transform(img_size)(image=img)['image'] for img_size in self.all_img_sizes}\n        transormed_imgs[\"crop_names\"] = f\"{x}_{y}\"\n        transormed_imgs[\"not_background\"] = _check_background(img, self.crop_size)\n        return transormed_imgs","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## tiff inference function","metadata":{}},{"cell_type":"code","source":"    \ndef inference(data_loader, models_with_config, crop_size):\n    img_size = (data_loader.dataset.h, data_loader.dataset.w)\n    mask_pred = np.zeros(img_size, dtype=np.uint8)\n    #iterate over crops\n    for batch in tqdm(data_loader, ncols=70, leave=True):\n        #skip crop if there is no tissue\n        if SKIP_BACKGROUND is True and batch[\"not_background\"].item() is False:\n            continue\n        with torch.no_grad():\n            pred_total = None # accumulation for averaged by fold all models predictions\n            # iterate over folds of each model\n            for cfg, model_group in models_with_config:\n                image = batch[cfg.img_size].cuda()\n                pred_group = None # accumulation for all folds\n\n                for model in model_group:\n                    pred = model(image)\n                    pred = pred.sigmoid()\n                    if cfg.img_size != TRAIN_IMG_SIZE:\n                        pred = torch.nn.functional.interpolate(pred, size=TRAIN_IMG_SIZE, mode='bilinear')\n                    if pred_group == None:\n                        pred_group = pred\n                    else:\n                        pred_group += pred\n                        \n                pred_group = pred_group.squeeze()\n                if len(pred_group.shape) == 2:\n                    pred_group = pred_group.unsqueeze(0)\n                image.cpu()\n                del image\n                pred_group = cfg.weight_blend * (pred_group / len(model_group))\n                if pred_total == None:\n                    pred_total = pred_group\n                else:\n                    pred_total += pred_group\n\n            pred_total = (pred_total.cpu().data.numpy() > THR).astype(np.uint8)\n            for predict_single, crop_name in zip(pred_total, batch['crop_names']):\n                x = int(crop_name.split(\"_\")[-2])\n                y = int(crop_name.split(\"_\")[-1])\n\n                if crop_size != TRAIN_IMG_SIZE:\n                    predict_single = cv2.resize(predict_single, (crop_size, crop_size))\n      \n                predict_single[0:RAMKA] = 0\n                predict_single[-RAMKA:] = 0\n                predict_single[:, 0:RAMKA] = 0\n                predict_single[:, -RAMKA:] = 0\n                    \n                mask_pred[y : y + crop_size, x : x + crop_size] += predict_single\n\n    mask_pred = (mask_pred > VOTE)\n    mask_rle = rle_encode_less_memory(mask_pred)\n    del mask_pred, pred\n    gc.collect()\n    gc.collect()\n    return mask_rle","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Initialization of models and loading them to GPU","metadata":{}},{"cell_type":"code","source":"all_img_sizes = set([cfg.img_size for cfg in my_models])\nmodels_with_config = []\nfor cfg in my_models:\n    models_group = []\n    for w_path in cfg.weights_path:\n        model = smp.Unet(cfg.encoder, encoder_weights=None).cuda()\n        model.load_state_dict(torch.load(w_path))\n        model.eval()\n        models_group.append(model)\n    models_with_config.append((cfg, models_group))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Iterating over test tiffs","metadata":{}},{"cell_type":"code","source":"count_thr = 5 if SKIP_COMMIT is True else 4\nif len(df_sub) > count_thr:\n    for idx, row in df_sub.iterrows():\n        test_ds = SingleTiffDataset(\n                tiff_path=f\"../input/hubmap-kidney-segmentation/test/{row['id']}.tiff\",\n                all_img_sizes=all_img_sizes,\n                crop_size=CROP_SIZE, \n                step=STEP,\n            )\n\n        test_loader = DataLoader(\n                dataset=test_ds,\n                batch_size=BATCH_SIZE,\n                shuffle=False,\n                num_workers=NUM_WORKERS,\n                pin_memory=True,\n            )\n        rle = inference(test_loader, models_with_config, CROP_SIZE)\n        df_sub.loc[idx, \"predicted\"] = rle","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!rm -rf model_1\n!rm -rf model_2\n# !rm -rf model_3","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_sub.to_csv(\"submission.csv\", index=False)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_sub","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}