{"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":"markdown","source":"# HuBMAP - deepflash2 submission\n\n> Submission kernel for model trained with efficient region based sampling. \n\n***\n\n\n## Highlights\n\n- *Super fast submission*: ~30 min for a 5-model ensemble on the public test set\n- Using overlapping tiles and gaussian weighting similar to [nnuet](https://www.nature.com/articles/s41592-020-01008-z), which removes almost all prediction artifacts\n\n## Related Kernels\n- Training: https://www.kaggle.com/matjes/hubmap-deepflash2-train\n- Sampling: https://www.kaggle.com/matjes/hubmap-efficient-sampling-ii-deepflash2\n\nTo make this kernel fast and reliable, special thanks go to @leighplt ([kernel](https://www.kaggle.com/leighplt/pytorch-fcn-resnet50)) and @iafoss ([kernel](https://www.kaggle.com/iafoss/hubmap-pytorch-fast-ai-starter-sub))!\n\n## Overview\n\n1. Installation and package loading\n2. Functions and classes for prediction\n3. Configuration\n4. Prediction\n5. Submission\n\n### Installation and package loading","metadata":{"id":"OcsetTMwKXqC"}},{"cell_type":"code","source":"# Install deepflash2 and dependencies\nimport sys\nsys.path.append(\"../input/segmentation-models-pytorch-install\")\n!pip install -q --no-deps ../input/deepflash2-lfs\nimport cv2, torch, gc, rasterio\nimport torch.nn.functional as F\nimport deepflash2.tta as tta\nimport matplotlib.pyplot as plt\nimport pandas as pd, numpy as np\nimport segmentation_models_pytorch as smp\nfrom pathlib import Path\nfrom rasterio.windows import Window\nfrom torch.utils.data import Dataset, DataLoader\nfrom scipy .ndimage.filters import gaussian_filter\nfrom tqdm.notebook import tqdm\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Functions and classes for prediction","metadata":{}},{"cell_type":"code","source":"#https://www.kaggle.com/bguberfain/memory-aware-rle-encoding\n#with transposed mask\ndef 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 load_model_weights(model, file, strict=True):\n    state = torch.load(file, map_location='cpu')\n    stats = state['stats']\n    model_state = state['model']\n    model.load_state_dict(model_state, strict=strict)\n    return model, stats\n\n# from https://github.com/MIC-DKFZ/nnUNet/blob/2fade8f32607220f8598544f0d5b5e5fa73768e5/nnunet/network_architecture/neural_network.py#L250\ndef _get_gaussian(patch_size, sigma_scale=1. / 8) -> np.ndarray:\n    tmp = np.zeros(patch_size)\n    center_coords = [i // 2 for i in patch_size]\n    sigmas = [i * sigma_scale for i in patch_size]\n    tmp[tuple(center_coords)] = 1\n    gaussian_importance_map = gaussian_filter(tmp, sigmas, 0, mode='constant', cval=0)\n    gaussian_importance_map = gaussian_importance_map / np.max(gaussian_importance_map) * 1\n    gaussian_importance_map = gaussian_importance_map.astype(np.float32)\n\n    # gaussian_importance_map cannot be 0, otherwise we may end up with nans!\n    gaussian_importance_map[gaussian_importance_map == 0] = np.min(\n        gaussian_importance_map[gaussian_importance_map != 0])\n\n    return gaussian_importance_map","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Some code adapted from https://www.kaggle.com/iafoss/hubmap-pytorch-fast-ai-starter-sub\nclass HubmapDataset(Dataset):\n    'HubmapDataset class that does not load the full tiff files.'\n    def __init__(self, file, stats, scale=3, shift=.8, output_shape=(512,512), s_th = 40):\n        \n        self.mean, self.std = stats\n        self.scale = scale\n        self.shift = shift\n        self.output_shape = output_shape\n        self.input_shape = tuple(int(t*scale) for t in self.output_shape)      \n        self.s_th = s_th #saturation blancking threshold\n        self.p_th = 1000*(self.output_shape[0]//256)**2 #threshold for the minimum number of pixels\n\n        identity = rasterio.Affine(1, 0, 0, 0, 1, 0)\n        self.data = rasterio.open(file, transform = identity, num_threads='all_cpus')\n        if self.data.count != 3:\n            subdatasets = self.data.subdatasets\n            self.layers = []\n            if len(subdatasets) > 0:\n                for i, subdataset in enumerate(subdatasets, 0):\n                    self.layers.append(rasterio.open(subdataset))\n            \n        # Tiling\n        self.slices = []\n        self.out_slices = []\n        self.out_data_shape = tuple(int(x//self.scale) for x in self.data.shape)\n        start_points = [o//2 for o in self.output_shape]\n        end_points = [(s - st) for s, st in zip(self.out_data_shape, start_points)]\n        n_points = [int(s//(o*self.shift))+1 for s, o in zip(self.out_data_shape, self.output_shape)]\n        center_points = [np.linspace(st, e, num=n, endpoint=True, dtype=np.int64) for st, e, n in zip(start_points, end_points, n_points)]\n        for cx in center_points[1]:\n            for cy in center_points[0]:\n                # Calculate output slices for whole image\n                slices = tuple(slice(int((c*self.scale - o/2).clip(0, s)), int((c*self.scale + o/2).clip(max=s)))\n                                 for (c, o, s) in zip((cy, cx), self.input_shape, self.data.shape))\n                self.slices.append(slices)\n                \n                out_slices = tuple(slice(int((c - o/2).clip(0, s)), int((c + o/2).clip(max=s)))\n                                 for (c, o, s) in zip((cy, cx), self.output_shape, self.out_data_shape))\n                self.out_slices.append(out_slices)\n                \n\n    def __len__(self):\n        return len(self.slices)\n\n    def __getitem__(self, idx):\n        if torch.is_tensor(idx):\n            idx = idx.tolist()\n            \n        slices = self.slices[idx]\n        if self.data.count == 3: # normal\n            img = self.data.read([1, 2, 3], \n                window=Window.from_slices(*slices)\n            )\n            img = np.moveaxis(img, 0, -1)\n        else: # with subdatasets/layers\n            img = np.zeros((*self.input_shape, 3), dtype=np.uint8)\n            for fl in range(3):\n                img[:, :, fl] = self.layers[fl].read(\n                    window=Window.from_slices(*slices)\n                )\n        \n        if self.scale!=1:\n            img = cv2.resize(img, self.output_shape, interpolation = cv2.INTER_AREA)\n        \n        #check for empty imges\n        hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)\n        h,s,v = cv2.split(hsv)\n        if (s>self.s_th).sum() <= self.p_th or img.sum() <= self.p_th:\n            # Remove if idx=-1\n            idx = -1\n        \n        img = (img/255.0 - self.mean)/self.std\n        img = img.transpose(2, 0, 1).astype('float32')\n        \n        return torch.from_numpy(img), idx\n    \nclass Model_pred:\n    'Class for prediction with multiple models'\n    def __init__(self, models, use_tta=True, batch_size=32):\n        self.models = models\n        self.bs = batch_size\n        self.tfms = [tta.HorizontalFlip(), tta.VerticalFlip()] if use_tta else []\n        \n    def predict(self, ds):\n        #rasterio cannot be used with multiple workers\n        dl = DataLoader(ds, self.bs, num_workers=0, shuffle=False, pin_memory=True)\n        \n        # Create zero arrays\n        pred = np.zeros(ds.out_data_shape, dtype='float32')\n        merge_map = np.zeros(ds.out_data_shape, dtype='float32')\n        \n        # Gaussian weights\n        gw_numpy = _get_gaussian(ds.output_shape)\n        gw = torch.from_numpy(gw_numpy).to(device)\n        \n        with torch.no_grad():\n            for images, idxs in tqdm(iter(dl), total=len(dl)):\n                if ((idxs>=0).sum() > 0): #exclude empty images\n                    images = images[idxs>=0].to(device)\n                    idxs = idxs[idxs>=0]\n                    merger = tta.Merger()\n                    for t in tta.Compose(self.tfms):\n                        aug_images = t.augment_image(images)\n                        model_merger = tta.Merger()\n                        for model in self.models:\n                            out = model(aug_images)\n                            out = F.softmax(out, dim=1)\n                            model_merger.append(out)\n                        out = t.deaugment_mask(model_merger.result())\n                        merger.append(out)\n            \n                    # Apply gaussian weigthing\n                    batch_smx = merger.result()*gw.view(1,1,*gw.shape)\n                    batch_smx = [x for x in batch_smx.permute(0,2,3,1).cpu().numpy()]\n\n                    for smx, idx in zip(batch_smx, idxs):\n                        slcs = ds.out_slices[idx]\n                        # Only using positive class here\n                        pred[slcs] += smx[...,1]\n                        merge_map[slcs] += gw_numpy\n\n        pred /= merge_map\n        return pred","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Configuration","metadata":{}},{"cell_type":"code","source":"class CONFIG():\n    \n    # data paths\n    data_path = Path('../input/hubmap-kidney-segmentation')\n    # Local train on all data\n    model_path = Path('../input/hubmap-models-final/Unet_efficientnet-b2_1.5_2500_Unet_b2_aughard_DiceCELoss_no-d48/Unet_efficientnet-b2_1.5_2500_Unet_b2_aughard_DiceCELoss_no-d48')\n    #model_path = Path('../input/hubmap-deepflash2-train')\n   \n    # zoom factor (e.g., 3 means downscaling from 1536 to 512)\n    scale = 3 \n    # tile shift for prediction\n    shift = 0.8 \n    tile_shape = (512, 512)\n\n    # pytorch model (https://github.com/qubvel/segmentation_models.pytorch)\n    encoder_name = \"efficientnet-b2\"\n    encoder_weights = None\n    in_channels = 3\n    classes = 2\n    \n    # dataloader \n    batch_size = 32\n    \n    # test time augmentation\n    tta = True\n    # prediction threshold\n    threshold = 0.5\n    \ncfg = CONFIG()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Sample submissions for ids\ndf_sample = pd.read_csv(cfg.data_path/'sample_submission.csv',  index_col='id')\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\n# Models (see https://github.com/qubvel/segmentation_models.pytorch)\nMODELS = [f for f in cfg.model_path.iterdir() if f.suffix=='.pth']\nprint(f'Found {len(MODELS)} models', *MODELS)\n\nmodels = []\nfor i, m_path in enumerate(MODELS):\n    #state_dict = torch.load(path,map_location=torch.device('cpu'))\n    model = smp.Unet(encoder_name=cfg.encoder_name, \n                     encoder_weights=cfg.encoder_weights, \n                     in_channels=cfg.in_channels, \n                     classes=cfg.classes)\n    model, stats = load_model_weights(model, m_path)\n    model.float()\n    model.eval()\n    model.to(device)\n    models.append(model)\n\nmp = Model_pred(models, use_tta=cfg.tta, batch_size=cfg.batch_size)\nprint(len(models))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Prediction","metadata":{}},{"cell_type":"code","source":"names,preds = [],[]\nfor idx,row in tqdm(df_sample.iterrows(),total=len(df_sample)):\n    \n    print(f'###### File {idx} ######')\n    f = cfg.data_path/'test'/f'{idx}.tiff'\n    ds = HubmapDataset(f, stats, scale=cfg.scale, shift=cfg.shift, output_shape=cfg.tile_shape)\n    \n    print('Predicting...')   \n    pred = mp.predict(ds)\n       \n    print('Rezising...')\n    shape = ds.data.shape\n    pred = cv2.resize((pred*255).astype('uint8'), (shape[1], shape[0]))\n    \n    pred = (pred>cfg.threshold*255).astype(np.uint8)\n    \n    #convert to rle\n    #https://www.kaggle.com/bguberfain/memory-aware-rle-encoding\n    rle = rle_encode_less_memory(pred)\n    names.append(idx)\n    preds.append(rle)\n    \n    print('Plotting')\n    fig, ax = plt.subplots(figsize=(15,15))\n    ax.imshow(cv2.resize(pred, (1024, 1024*shape[0]//shape[1])))\n    plt.show()\n    \n    del pred\n    gc.collect()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df = pd.DataFrame({'id':names,'predicted':preds})\ndf.to_csv('submission.csv',index=False)\ndf.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}