{"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 numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport cv2\nimport os\nfrom pathlib import Path\nimport gc\nfrom tqdm.notebook import tqdm\nimport rasterio\nfrom rasterio.windows import Window\nimport torch, torchvision\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nimport warnings\nwarnings.filterwarnings(\"ignore\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import fastai\nfastai.__version__","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import fastai.layers as fastai_layers","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sz = 480  #the size of tiles\nTH = 0.438  #threshold for positive predictions\nWINDOW = 1024\nOVERLAP = 308\nEDGE_IGNORE = 150\ntta = True\nDATA = '/kaggle/input/hubmap-kidney-segmentation/test/'\nMODELS = [\n         '/kaggle/input/hubmap-win1024-sz480-subds-fold0-8epochs/fullTraining_8epochs',\n         '/kaggle/input/hubmap-win1024-sz480-subds-fold3-8epochs/fullTraining_8epochs',\n         '/kaggle/input/hubmap-win1024-sz480-subds-fold1-8epochs/fullTraining_8epochs',\n         '/kaggle/input/hubmap-sz480win1024-train-fold2-subds/model_checkpoints_fold_2/fullTraining',\n         '/kaggle/input/hubmap-win1024-sz480-subds-fold4-8epochs/fullTraining_8epochs',\n         ]\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Data","metadata":{}},{"cell_type":"code","source":"#functions to convert encoding to mask and mask to encoding\ndef enc2mask(encs, shape):\n    # shape: (width, height)\n    img = np.zeros(shape[0]*shape[1], dtype=np.uint8)\n    for m,enc in enumerate(encs):\n        if isinstance(enc,np.float) and np.isnan(enc): continue\n        s = enc.split()\n        for i in range(len(s)//2):\n            start = int(s[2*i]) - 1\n            length = int(s[2*i+1])\n            img[start:start+length] = 1 + m\n    return img.reshape(shape).T\n\ndef mask2enc(mask, n=1):\n    pixels = mask.T.flatten()\n    encs = []\n    for i in range(1,n+1):\n        p = (pixels == i).astype(np.int8)\n        if p.sum() == 0: encs.append(np.nan)\n        else:\n            p = np.concatenate([[0], p, [0]])\n            runs = np.where(p[1:] != p[:-1])[0] + 1\n            runs[1::2] -= runs[::2]\n            encs.append(' '.join(str(x) for x in runs))\n    return encs\n\ndef rle_encode(img:np.ndarray):\n    img = img.T.flatten()\n    idxs = np.nonzero(img[:-1]!=img[1:])[0]\n    if img[0] == 1:\n        idxs = np.concatenate([[-1],idxs])\n    if img[-1] == 1:\n        idxs = np.concatenate([idxs,[len(img)-1]])\n    # index for 1start\n    idxs[::2] = idxs[::2] + 1\n    idxs[1::2] = (idxs[1::2] - idxs[::2]) + 1\n    idxs[::2] += 1\n    return ' '.join(str(x) for x in idxs)","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def 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    assert np.all(x2-x1 == window), \"Row or height not equal to window. All tiles must be window_x_window\"\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    assert np.all(y2-y1 == window), \"Column or width not equal to window. All tiles must be window_x_window\"\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)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"mean = np.array([0.63468326, 0.48969275, 0.67348264])\nstd = np.array([0.20481194, 0.25495073, 0.1935244])\n\nstd_th = 7\n\ndef img2tensor(img,dtype:np.dtype=np.float32):\n    if img.ndim==2 : img = np.expand_dims(img,2)\n    img = np.transpose(img,(2,0,1))\n    return torch.from_numpy(img.astype(dtype, copy=False))\n\nclass HuBMAPDataset(Dataset):\n    def __init__(self, idx):\n        self.data = rasterio.open(os.path.join(DATA,idx+'.tiff'))\n        if self.data.count != 3:\n            self.layers = []\n            if len(self.data.subdatasets) > 0:\n                for i, subdataset in enumerate(self.data.subdatasets, 0):\n                    self.layers.append(rasterio.open(subdataset))\n        self.slices = make_grid(self.data.shape, window=WINDOW, min_overlap=OVERLAP) #[num_slices,4]\n        \n    def __len__(self):\n        return self.slices.shape[0]\n    \n    def __getitem__(self, idx):\n        x1,x2,y1,y2 = self.slices[idx]\n        if self.data.count == 3:\n            img = self.data.read(window=Window.from_slices((x1,x2),(y1,y2))) # shape: [C,H,W]\n        else:\n            img = np.zeros((3,WINDOW,WINDOW),np.uint8)\n            for j, layer in enumerate(self.layers):\n                img[j,:,:] = layer.read(window=Window.from_slices((x1,x2),(y1,y2)))[0]\n        img = np.moveaxis(img, 0, -1)\n        #img = cv2.resize(img, (sz, sz), interpolation=cv2.INTER_AREA)\n        \n        #check for images with no real content\n        if np.all(np.array([np.std(img[:,:,i]) for i in range(3)]) <= std_th):\n            #images with -1 will be skipped\n            img = cv2.resize(img, (sz, sz), interpolation=cv2.INTER_AREA)\n            return img2tensor((img/255.0 - mean)/std), torch.tensor(-1), torch.tensor((x1,x2,y1,y2))\n        else:\n            img = cv2.resize(img, (sz, sz), interpolation=cv2.INTER_AREA)\n            return img2tensor((img/255.0 - mean)/std), torch.tensor(1), torch.tensor((x1,x2,y1,y2))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"class FPN(torch.nn.Module):\n    def __init__(self, input_channels:list, output_channels:list):\n        super().__init__()\n        self.convs = torch.nn.ModuleList(\n            [torch.nn.Sequential(torch.nn.Conv2d(in_ch, out_ch*2, kernel_size=3, padding=1),\n             torch.nn.ReLU(inplace=True), torch.nn.BatchNorm2d(out_ch*2),\n             torch.nn.Conv2d(out_ch*2, out_ch, kernel_size=3, padding=1))\n            for in_ch, out_ch in zip(input_channels, output_channels)])\n        \n    def forward(self, xs:list, last_layer):\n        hcs = [F.interpolate(c(x),scale_factor=2**(len(self.convs)-i),mode='bilinear')\n               for i,(c,x) in enumerate(zip(self.convs, xs))]\n        hcs.append(last_layer)\n        return torch.cat(hcs, dim=1)\n\nclass PixelShuffle_ICNR(torch.nn.Sequential):\n    \"Upsample by `scale` from `ni` filters to `nf` (default `ni`), using `nn.PixelShuffle`.\"\n    def __init__(self, ni, nf=None, scale=2, blur=False, norm_type=fastai_layers.NormType.Weight, \n                 act_cls=fastai_layers.defaults.activation):\n        super().__init__()\n        nf = fastai_layers.ifnone(nf, ni)\n        layers = [fastai_layers.ConvLayer(ni, nf*(scale**2), ks=1, norm_type=norm_type, act_cls=act_cls, bias_std=0),\n                  torch.nn.PixelShuffle(scale)]\n        if norm_type == fastai_layers.NormType.Weight:\n            layers[0][0].weight_v.data.copy_(fastai_layers.icnr_init(layers[0][0].weight_v.data))\n            layers[0][0].weight_g.data.copy_(((layers[0][0].weight_v.data**2).sum(dim=[1,2,3])**0.5)[:,None,None,None])\n        else:\n            layers[0][0].weight.data.copy_(fastai_layers.icnr_init(layers[0][0].weight.data))\n        \n        if blur: layers += [torch.nn.ReplicationPad2d((1,0,1,0)), torch.nn.AvgPool2d(2, stride=1)]\n        super().__init__(*layers)\n        \nclass UnetBlock(torch.nn.Module):\n    def __init__(self, up_in_c:int, x_in_c:int, nf:int=None, blur:bool=False,\n                 self_attention:bool=False, **kwargs):\n        super().__init__()\n        self.shuf = PixelShuffle_ICNR(up_in_c, up_in_c//2, blur=blur, norm_type=None, **kwargs)\n        self.bn = torch.nn.BatchNorm2d(x_in_c)\n        ni = up_in_c//2 + x_in_c\n        nf = nf if nf is not None else max(up_in_c//2,32)\n        self.conv1 = fastai_layers.ConvLayer(ni, nf, norm_type=None, **kwargs)\n        self.conv2 = fastai_layers.ConvLayer(nf, nf, norm_type=None, **kwargs)\n        self.relu = torch.nn.ReLU(inplace=True)\n\n    def forward(self, up_in:torch.Tensor, left_in:torch.Tensor) -> torch.Tensor:\n        s = left_in\n        up_out = self.shuf(up_in)\n        cat_x = self.relu(torch.cat([up_out, self.bn(s)], dim=1))\n        return self.conv2(self.conv1(cat_x))\n\nclass _ASPPModule(torch.nn.Module):\n    def __init__(self, inplanes, planes, kernel_size, padding, dilation, groups=1):\n        super().__init__()\n        self.atrous_conv = torch.nn.Conv2d(inplanes, planes, kernel_size=kernel_size,\n                stride=1, padding=padding, dilation=dilation, bias=False, groups=groups)\n        self.bn = torch.nn.BatchNorm2d(planes)\n        self.relu = torch.nn.ReLU()\n\n        self._init_weight()\n\n    def forward(self, x):\n        x = self.atrous_conv(x)\n        x = self.bn(x)\n\n        return self.relu(x)\n\n    def _init_weight(self):\n        for m in self.modules():\n            if isinstance(m, torch.nn.Conv2d):\n                torch.nn.init.kaiming_normal_(m.weight)\n            elif isinstance(m, torch.nn.BatchNorm2d):\n                m.weight.data.fill_(1)\n                m.bias.data.zero_()\n\nclass ASPP(torch.nn.Module):\n    def __init__(self, inplanes=512, mid_c=256, dilations=[6, 12, 18, 24], out_c=None):\n        super().__init__()\n        self.aspps = [_ASPPModule(inplanes, mid_c, 1, padding=0, dilation=1)] + \\\n            [_ASPPModule(inplanes, mid_c, 3, padding=d, dilation=d, groups=4) for d in dilations]\n        self.aspps = torch.nn.ModuleList(self.aspps)\n        self.global_pool = torch.nn.Sequential(torch.nn.AdaptiveMaxPool2d((1, 1)),\n                        torch.nn.Conv2d(inplanes, mid_c, 1, stride=1, bias=False),\n                        torch.nn.BatchNorm2d(mid_c), torch.nn.ReLU())\n        out_c = out_c if out_c is not None else mid_c\n        self.out_conv = torch.nn.Sequential(torch.nn.Conv2d(mid_c*(2+len(dilations)), out_c, 1, bias=False),\n                                    torch.nn.BatchNorm2d(out_c), torch.nn.ReLU(inplace=True))\n        self._init_weight()\n\n    def forward(self, x):\n        x0 = self.global_pool(x)\n        xs = [aspp(x) for aspp in self.aspps]\n        x0 = F.interpolate(x0, size=xs[0].size()[2:], mode='bilinear', align_corners=True)\n        x = torch.cat([x0] + xs, dim=1)\n        return self.out_conv(x)\n    \n    def _init_weight(self):\n        for m in self.modules():\n            if isinstance(m, torch.nn.Conv2d):\n                torch.nn.init.kaiming_normal_(m.weight)\n            elif isinstance(m, torch.nn.BatchNorm2d):\n                m.weight.data.fill_(1)\n                m.bias.data.zero_()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class UneXt50(torch.nn.Module):\n    def __init__(self, stride=1):\n        super().__init__()\n        #encoder\n        #m = torch.hub.load('facebookresearch/semi-supervised-ImageNet1K-models',\n        #                   'resnext50_32x4d_swsl')\n        m = torchvision.models.resnext50_32x4d(pretrained=False)\n        self.enc0 = torch.nn.Sequential(m.conv1, m.bn1, torch.nn.ReLU(inplace=True))\n        self.enc1 = torch.nn.Sequential(torch.nn.MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1),\n                            m.layer1) #256\n        self.enc2 = m.layer2 #512\n        self.enc3 = m.layer3 #1024\n        self.enc4 = m.layer4 #2048\n        #aspp with customized dilatations\n        self.aspp = ASPP(2048,256,out_c=512,dilations=[stride*1,stride*2,stride*3,stride*4])\n        self.drop_aspp = torch.nn.Dropout2d(0.25)\n        #decoder\n        self.dec4 = UnetBlock(512,1024,256)\n        self.dec3 = UnetBlock(256,512,128)\n        self.dec2 = UnetBlock(128,256,64)\n        self.dec1 = UnetBlock(64,64,32)\n        self.fpn = FPN([512,256,128,64],[16]*4)\n        self.drop = torch.nn.Dropout2d(0.1)\n        self.final_conv = fastai_layers.ConvLayer(32+16*4, 1, ks=1, norm_type=None, act_cls=None)\n        \n    def forward(self, x): #N,3,H,W\n        enc0 = self.enc0(x) #N,64,H/2,W/2\n        enc1 = self.enc1(enc0) #N,256,H/(2**2),W/(2**2)\n        enc2 = self.enc2(enc1) #N,512,H/(2**3),W/(2**3)\n        enc3 = self.enc3(enc2) #N,1024,H/(2**4),W/(2**4)\n        enc4 = self.enc4(enc3) #N,2048,H/(2**5),W/(2**5)\n        enc5 = self.aspp(enc4) #N,512,H/(2**5),W/(2**5)\n        dec3 = self.dec4(self.drop_aspp(enc5),enc3) #N,256,H/(2**4),W/(2**4)\n        dec2 = self.dec3(dec3,enc2) #N,128,H/(2**3),W/(2**3)\n        dec1 = self.dec2(dec2,enc1) #N,64,H/(2**2),W/(2**2)\n        dec0 = self.dec1(dec1,enc0) #N,32,H/(2**1),W/(2**1)\n        x = self.fpn([enc5, dec3, dec2, dec1], dec0) #N,96,H/(2**1),W/(2**1)\n        x = self.final_conv(self.drop(x)) #N,1,H/(2**1),W/(2**1)\n        x = F.interpolate(x,scale_factor=2,mode='bilinear') #N,1,H,W\n        return x","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"models = []\nfor path in MODELS:\n    model = UneXt50(stride=2)\n    model.load_state_dict(torch.load(path,map_location=torch.device('cpu')))\n    model.to(device)\n    model.eval()\n    models.append(model)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(models)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Prediction","metadata":{}},{"cell_type":"code","source":"\"\"\"\nidx = 'c68fe75ea'\nnum_ex = 50\nplt.figure(figsize=(10,20))\nds = HuBMAPDataset(idx)\n\nfor i in range(num_ex):\n    img, b, _ = ds[100+i]\n    img = img.numpy()\n    img = np.transpose(img, (1,2,0))\n    img = ((img * std + mean) * 255.0).astype(np.uint8)\n    plt.subplot(10,5,i+1)\n    plt.imshow(img)\n    plt.title(f'{b}')\n    plt.axis('off')\n\"\"\"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_paths = list(Path('/kaggle/input/hubmap-kidney-segmentation/test').glob('*tiff')); test_paths","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"names,preds_rle = [],[]\nbs = 12\nwith torch.no_grad():\n    for test_path in tqdm(test_paths):\n        idx = test_path.stem\n        ds = HuBMAPDataset(idx)\n        dl = DataLoader(ds, batch_size=bs, shuffle=False, drop_last=False, num_workers=0)\n        mask = torch.zeros(ds.data.shape, dtype=torch.uint8)\n        for imgs, igns, coords in iter(dl):\n            imgs = imgs[igns==1]\n            if len(imgs) == 0:\n                continue\n            imgs = imgs.to(device)\n            preds = torch.zeros((imgs.shape[0],1,*imgs.shape[-2:]), dtype=torch.float32, device=device)\n            for model in models:\n                preds += torch.sigmoid(model(imgs))\n                if tta:\n                    flips = [[-1],[-2],[-2,-1]]\n                    for f in flips:\n                        _preds_f = torch.sigmoid(model(torch.flip(imgs,f)))\n                        preds += torch.flip(_preds_f, f)\n            if tta:\n                preds = preds / (len(models)*(len(flips)+1))\n            else:\n                preds = preds / len(models)\n            preds = F.interpolate(preds, size=(WINDOW,WINDOW), mode='bilinear')\n            \n            coords = coords[igns==1]\n            for i in range(len(coords)):\n                x1,x2,y1,y2 = coords[i].tolist()    \n                x1_ign, x2_ign, y1_ign, y2_ign = 0,0,0,0\n                if x1 != 0:\n                    x1_ign = EDGE_IGNORE\n                if x2 != mask.shape[0]:\n                    x2_ign = EDGE_IGNORE\n                if y1 != 0:\n                    y1_ign = EDGE_IGNORE\n                if y2 != mask.shape[1]:\n                    y2_ign = EDGE_IGNORE\n                mask[(x1+x1_ign):(x2-x2_ign), (y1+y1_ign):(y2-y2_ign)] += (preds[i, 0, x1_ign:WINDOW-x2_ign, y1_ign:WINDOW-y2_ign] > TH).to(dtype=torch.uint8, device='cpu')\n\n        mask.clamp_(0,1)\n        mask_np = mask.numpy()\n\n        rle = rle_encode(mask_np)\n        names.append(idx)\n        preds_rle.append(rle)\n        \n        del mask, mask_np, ds, dl\n        gc.collect()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df = pd.DataFrame({'id':names,'predicted':preds_rle})\ndf.to_csv('submission.csv',index=False)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}