{"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":"","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom PIL import Image\nimport tifffile as tiff\nimport cv2\nimport os\nimport gc\nfrom tqdm.notebook import tqdm\nimport rasterio\nfrom rasterio.windows import Window\n\nfrom fastai.vision.all import *\nfrom torch.utils.data import Dataset, DataLoader\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-06-02T15:02:55.536543Z","iopub.execute_input":"2022-06-02T15:02:55.537153Z","iopub.status.idle":"2022-06-02T15:02:59.204502Z","shell.execute_reply.started":"2022-06-02T15:02:55.537074Z","shell.execute_reply":"2022-06-02T15:02:59.203557Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sz = 512   #the size of tiles\nsz_reduction = 4 #reduce the original images by 4 times\nexpansion = 128\nTH = 0.5  #threshold for positive predictions\nDATA = '../input/hubmap-kidney-segmentation/test/'\nMODELS_rsxt50 = [f'../input/hubmap-3rd-place-fast-weights/fold{i}.pth' for i in range(5)]\n#MODELS_rsxt101 = [f'../input/hubmaptest45/fold{i}.pth' for i in range(5)]\nMODELS_rsxt101 = []\n\ndf_sample = pd.read_csv('../input/hubmap-kidney-segmentation/sample_submission.csv')\nbs = 16\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')","metadata":{"execution":{"iopub.status.busy":"2022-06-02T15:02:59.206979Z","iopub.execute_input":"2022-06-02T15:02:59.207697Z","iopub.status.idle":"2022-06-02T15:02:59.228398Z","shell.execute_reply.started":"2022-06-02T15:02:59.207635Z","shell.execute_reply":"2022-06-02T15:02:59.227104Z"},"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    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\n#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)","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-06-02T15:02:59.230145Z","iopub.execute_input":"2022-06-02T15:02:59.230874Z","iopub.status.idle":"2022-06-02T15:02:59.249867Z","shell.execute_reply.started":"2022-06-02T15:02:59.230828Z","shell.execute_reply":"2022-06-02T15:02:59.248861Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# https://www.kaggle.com/iafoss/256x256-images\nmean = np.array([0.65459856,0.48386562,0.69428385])\nstd = np.array([0.15167958,0.23584107,0.13146145])\n\ns_th = 40  #saturation blancking threshold\np_th = 1000*(sz//256)**2 #threshold for the minimum number of pixels\nidentity = rasterio.Affine(1, 0, 0, 0, 1, 0)\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, sz=sz, sz_reduction=sz_reduction, expansion=expansion):\n        self.data = rasterio.open(os.path.join(DATA,idx+'.tiff'), transform = identity,\n                                 num_threads='all_cpus')\n        # some images have issues with their format \n        # and must be saved correctly before reading with rasterio\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        self.shape = self.data.shape\n        self.sz_reduction = sz_reduction\n        self.sz = sz_reduction*sz\n        self.expansion = sz_reduction*expansion\n        self.pad0 = (self.sz - self.shape[0]%self.sz)%self.sz\n        self.pad1 = (self.sz - self.shape[1]%self.sz)%self.sz\n        self.n0max = (self.shape[0] + self.pad0)//self.sz\n        self.n1max = (self.shape[1] + self.pad1)//self.sz\n        \n    def __len__(self):\n        return self.n0max*self.n1max\n    \n    def __getitem__(self, idx):\n        # the code below may be a little bit difficult to understand,\n        # but the thing it does is mapping the original image to\n        # tiles created with adding padding, as done in\n        # https://www.kaggle.com/iafoss/256x256-images ,\n        # and then the tiles are loaded with rasterio\n        # n0,n1 - are the x and y index of the tile (idx = n0*self.n1max + n1)\n        n0,n1 = idx//self.n1max, idx%self.n1max\n        # x0,y0 - are the coordinates of the lower left corner of the tile in the image\n        # negative numbers correspond to padding (which must not be loaded)\n        x0,y0 = -self.pad0//2 + n0*self.sz - self.expansion//2, -self.pad1//2 + n1*self.sz- self.expansion//2\n        # make sure that the region to read is within the image\n        p00,p01 = max(0,x0), min(x0+self.sz+self.expansion,self.shape[0])\n        p10,p11 = max(0,y0), min(y0+self.sz+self.expansion,self.shape[1])\n        img = np.zeros((self.sz+self.expansion,self.sz+self.expansion,3),np.uint8)\n        # mapping the loade region to the tile\n        if self.data.count == 3:\n            img[(p00-x0):(p01-x0),(p10-y0):(p11-y0)] = np.moveaxis(self.data.read([1,2,3],\n                window=Window.from_slices((p00,p01),(p10,p11))), 0, -1)\n        else:\n            for i,layer in enumerate(self.layers):\n                img[(p00-x0):(p01-x0),(p10-y0):(p11-y0),i] =\\\n                  layer.read(1,window=Window.from_slices((p00,p01),(p10,p11)))\n        \n        if self.sz_reduction != 1:\n            img = cv2.resize(img,((self.sz+self.expansion)//self.sz_reduction,(self.sz+self.expansion)//self.sz_reduction),\n                             interpolation = cv2.INTER_AREA)\n        #check for empty imges\n        hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)\n        h,s,v = cv2.split(hsv)\n        if (s>s_th).sum() <= p_th or img.sum() <= p_th:\n            #images with -1 will be skipped\n            return img2tensor((img/255.0 - mean)/std), -1\n        else: return img2tensor((img/255.0 - mean)/std), idx","metadata":{"execution":{"iopub.status.busy":"2022-06-02T15:02:59.252056Z","iopub.execute_input":"2022-06-02T15:02:59.252461Z","iopub.status.idle":"2022-06-02T15:02:59.287396Z","shell.execute_reply.started":"2022-06-02T15:02:59.252416Z","shell.execute_reply":"2022-06-02T15:02:59.286031Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#iterator like wrapper that returns predicted masks\nclass Model_pred:\n    def __init__(self, models, dl, tta:bool=True, half:bool=False):\n        self.models = models\n        self.dl = dl\n        self.tta = tta\n        self.half = half\n        \n    def __iter__(self):\n        count=0\n        with torch.no_grad():\n            for x,y in iter(self.dl):\n                if ((y>=0).sum() > 0): #exclude empty images\n                    x = x[y>=0].to(device)\n                    y = y[y>=0]\n                    if self.half: x = x.half()\n                    py = None\n                    for model in self.models:\n                        p = model(x)\n                        p = torch.sigmoid(p).detach()\n                        if py is None: py = p\n                        else: py += p\n                    if self.tta:\n                        #x,y,xy flips as TTA\n                        flips = [[-1],[-2],[-2,-1]]\n                        for f in flips:\n                            xf = torch.flip(x,f)\n                            for model in self.models:\n                                p = model(xf)\n                                p = torch.flip(p,f)\n                                py += torch.sigmoid(p).detach()\n                        py /= (1+len(flips))        \n                    py /= len(self.models)\n\n                    py = F.upsample(py, scale_factor=sz_reduction, mode=\"bilinear\")\n                    py = py.permute(0,2,3,1).float().cpu()\n                    \n                    batch_size = len(py)\n                    for i in range(batch_size):\n                        yield py[i,expansion*sz_reduction//2:-expansion*sz_reduction//2,expansion*sz_reduction//2:-expansion*sz_reduction//2],y[i]\n                        count += 1\n                    \n    def __len__(self):\n        return len(self.dl.dataset)","metadata":{"execution":{"iopub.status.busy":"2022-06-02T15:02:59.291644Z","iopub.execute_input":"2022-06-02T15:02:59.292584Z","iopub.status.idle":"2022-06-02T15:02:59.313417Z","shell.execute_reply.started":"2022-06-02T15:02:59.292466Z","shell.execute_reply":"2022-06-02T15:02:59.312119Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"import torch.nn as nn\nimport torch.tensor as Tensor\nimport torch\nfrom fastai.vision.all import *\nfrom torchvision.models.resnet import ResNet, Bottleneck\n\nclass FPN(nn.Module):\n    def __init__(self, input_channels:list, output_channels:list):\n        super().__init__()\n        self.convs = nn.ModuleList(\n            [nn.Sequential(nn.Conv2d(in_ch, out_ch*2, kernel_size=3, padding=1),\n             nn.ReLU(inplace=True), nn.BatchNorm2d(out_ch*2),\n             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 UnetBlock(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, **kwargs)\n        self.bn = 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 = ConvLayer(ni, nf, norm_type=None, **kwargs)\n        self.conv2 = ConvLayer(nf, nf, norm_type=None,\n            xtra=SelfAttention(nf) if self_attention else None, **kwargs)\n        self.relu = nn.ReLU(inplace=True)\n\n    def forward(self, up_in:Tensor, left_in:Tensor) -> 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(nn.Module):\n    def __init__(self, inplanes, planes, kernel_size, padding, dilation, groups=1):\n        super().__init__()\n        self.atrous_conv = nn.Conv2d(inplanes, planes, kernel_size=kernel_size,\n                stride=1, padding=padding, dilation=dilation, bias=False, groups=groups)\n        self.bn = nn.BatchNorm2d(planes)\n        self.relu = 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, nn.Conv2d):\n                torch.nn.init.kaiming_normal_(m.weight)\n            elif isinstance(m, nn.BatchNorm2d):\n                m.weight.data.fill_(1)\n                m.bias.data.zero_()\n\nclass ASPP(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 = nn.ModuleList(self.aspps)\n        self.global_pool = nn.Sequential(nn.AdaptiveMaxPool2d((1, 1)),\n                        nn.Conv2d(inplanes, mid_c, 1, stride=1, bias=False),\n                        nn.BatchNorm2d(mid_c), nn.ReLU())\n        out_c = out_c if out_c is not None else mid_c\n        self.out_conv = nn.Sequential(nn.Conv2d(mid_c*(2+len(dilations)), out_c, 1, bias=False),\n                                    nn.BatchNorm2d(out_c), nn.ReLU(inplace=True))\n        self.conv1 = nn.Conv2d(mid_c*(2+len(dilations)), out_c, 1, bias=False)\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, nn.Conv2d):\n                torch.nn.init.kaiming_normal_(m.weight)\n            elif isinstance(m, nn.BatchNorm2d):\n                m.weight.data.fill_(1)\n                m.bias.data.zero_()\n\n\nimport torch.hub\n# hub_model = torch.hub.load(\n#     'moskomule/senet.pytorch',\n#     'se_resnet50',\n#     pretrained=True,)\nimport torchvision\nclass UneXt(nn.Module):\n    def __init__(self, m, stride=1, **kwargs):\n        super().__init__()\n        #encoder\n        # m = torch.hub.load('facebookresearch/semi-supervised-ImageNet1K-models',\n        #                    'resnext101_32x4d_swsl')\n#         m = torch.hub.load('facebookresearch/semi-supervised-ImageNet1K-models',\n#                            'resnext50_32x4d_swsl', pretrained=False)\n        #m = ResNet(Bottleneck, [3, 4, 23, 3], groups=32, width_per_group=4)\n        #m = torchvision.models.resnext50_32x4d(pretrained=False)\n        # m = torch.hub.load(\n        #     'moskomule/senet.pytorch',\n        #     'se_resnet101',\n        #     pretrained=True,)\n\n        #m=torch.hub.load('zhanghang1989/ResNeSt', 'resnest50', pretrained=True)\n        self.enc0 = nn.Sequential(m.conv1, m.bn1, nn.ReLU(inplace=True))\n        self.enc1 = nn.Sequential(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 = nn.Dropout2d(0.5)\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 = nn.Dropout2d(0.1)\n        self.final_conv = ConvLayer(32+16*4, 1, ks=1, norm_type=None, act_cls=None)\n\n    def forward(self, x):\n        enc0 = self.enc0(x)\n        enc1 = self.enc1(enc0)\n        enc2 = self.enc2(enc1)\n        enc3 = self.enc3(enc2)\n        enc4 = self.enc4(enc3)\n        enc5 = self.aspp(enc4)\n        dec3 = self.dec4(self.drop_aspp(enc5),enc3)\n        dec2 = self.dec3(dec3,enc2)\n        dec1 = self.dec2(dec2,enc1)\n        dec0 = self.dec1(dec1,enc0)\n        x = self.fpn([enc5, dec3, dec2, dec1], dec0)\n        x = self.final_conv(self.drop(x))\n        x = F.interpolate(x,scale_factor=2,mode='bilinear')\n        return x\n\n","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-06-02T15:02:59.315407Z","iopub.execute_input":"2022-06-02T15:02:59.315903Z","iopub.status.idle":"2022-06-02T15:02:59.376667Z","shell.execute_reply.started":"2022-06-02T15:02:59.315873Z","shell.execute_reply":"2022-06-02T15:02:59.375411Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"models = []\nfor path in MODELS_rsxt50:\n    state_dict = torch.load(path)\n    model = UneXt(m=torchvision.models.resnext50_32x4d(pretrained=False)).cuda()\n    model = nn.DataParallel(model)\n    model.load_state_dict(state_dict)\n    model.float()\n    model.eval()\n    #model.to(device)\n    models.append(model)\n\nfor path in MODELS_rsxt101:\n    state_dict = torch.load(path)\n    model = UneXt(m=ResNet(Bottleneck, [3, 4, 23, 3], groups=32, width_per_group=4)).cuda()\n    model = nn.DataParallel(model)\n    model.load_state_dict(state_dict)\n    model.float()\n    model.eval()\n    #model.to(device)\n    models.append(model)    \n    \n    \ndel state_dict","metadata":{"execution":{"iopub.status.busy":"2022-06-02T15:02:59.37845Z","iopub.execute_input":"2022-06-02T15:02:59.379224Z","iopub.status.idle":"2022-06-02T15:03:15.646935Z","shell.execute_reply.started":"2022-06-02T15:02:59.379174Z","shell.execute_reply":"2022-06-02T15:03:15.645786Z"},"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    idx = row['id']\n    ds = HuBMAPDataset(idx)\n    #rasterio cannot be used with multiple workers\n    dl = DataLoader(ds,bs,num_workers=0,shuffle=False,pin_memory=True)\n    mp = Model_pred(models,dl)\n    #generate masks\n    mask = torch.zeros(len(ds),ds.sz,ds.sz,dtype=torch.int8)\n    for p,i in iter(mp): \n        print(p.shape)\n        mask[i.item()] = p.squeeze(-1) > TH\n    \n    #reshape tiled masks into a single mask and crop padding\n    mask = mask.view(ds.n0max,ds.n1max,ds.sz,ds.sz).\\\n        permute(0,2,1,3).reshape(ds.n0max*ds.sz,ds.n1max*ds.sz)\n    mask = mask[ds.pad0//2:-(ds.pad0-ds.pad0//2) if ds.pad0 > 0 else ds.n0max*ds.sz,\n        ds.pad1//2:-(ds.pad1-ds.pad1//2) if ds.pad1 > 0 else ds.n1max*ds.sz]\n    \n    #convert to rle\n    #https://www.kaggle.com/bguberfain/memory-aware-rle-encoding\n    rle = rle_encode_less_memory(mask.numpy())\n    names.append(idx)\n    preds.append(rle)\n    del mask, ds, dl\n    gc.collect()","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-06-02T15:03:15.649971Z","iopub.execute_input":"2022-06-02T15:03:15.650636Z","iopub.status.idle":"2022-06-02T15:19:18.129653Z","shell.execute_reply.started":"2022-06-02T15:03:15.65059Z","shell.execute_reply":"2022-06-02T15:19:18.128177Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df = pd.DataFrame({'id':names,'predicted':preds})\ndf.to_csv('submission.csv',index=False)","metadata":{"execution":{"iopub.status.busy":"2022-06-02T15:19:18.131467Z","iopub.execute_input":"2022-06-02T15:19:18.132204Z","iopub.status.idle":"2022-06-02T15:19:18.714804Z","shell.execute_reply.started":"2022-06-02T15:19:18.132152Z","shell.execute_reply":"2022-06-02T15:19:18.713782Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df.head()","metadata":{"execution":{"iopub.status.busy":"2022-06-02T15:19:18.717819Z","iopub.execute_input":"2022-06-02T15:19:18.718533Z","iopub.status.idle":"2022-06-02T15:19:18.75684Z","shell.execute_reply.started":"2022-06-02T15:19:18.718469Z","shell.execute_reply":"2022-06-02T15:19:18.755939Z"},"trusted":true},"execution_count":null,"outputs":[]}]}