{"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 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\nimport segmentation_models_pytorch as smp\nfrom segmentation_models_pytorch import Unet\nfrom torch.nn import functional as F\nfrom rasterio.windows import Window\nimport torch\nfrom torch import nn\nfrom torch.utils.data import Dataset, DataLoader\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sz = 256 # 256   #the size of tiles\nreduce = 4 # 4 #reduce the original images by 4 times\nTH = 0.5  #threshold for positive predictions\n\n\nTEST = True\n\n\nif not TEST:\n    DATA = '../input/hubmap-kidney-segmentation/train/'\n    df_sample = pd.read_csv('../input/hubmap-kidney-segmentation/train.csv')\nelse:\n    DATA = '../input/hubmap-kidney-segmentation/test/'\n    df_sample = pd.read_csv('../input/hubmap-kidney-segmentation/sample_submission.csv')\n\nMODELS = [f'../input/b4newdatanewtest/efficientnet-b4-unet-BCELoss-256-FOLD-{i}-model.pth' for i in range(5)]\n# MODELS2 = [f'../input/b4newdatanewtest/efficientnet-b4-unet-BCELoss-256-FOLD-{i}-model.pth' for i in range(5)]\n# MODELS += MODELS2\nbs = 2\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n# model_name = 'efficientnet-b4' # efficientnet-b4, se_resnext50_32x4d\nEXPAND = 8\nminoverlap = 1 / 32\nTTA = True\ntile_size = int(sz * EXPAND) # 1024 # 2048\ntile_resized = int(tile_size * reduce) # 4096 # 8192","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"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): \n            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":{"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\n\ndef img2tensor(img, dtype:np.dtype=np.float32):\n    if img.ndim == 2 : \n        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\n\ndef 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    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    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)\n\n\n# ---------------------------------------------------------------------------------------------------------------\n\nclass HuBMAPDataset(Dataset):\n    def __init__(self, idx, sz=sz, reduce=reduce):\n        self.data = rasterio.open(os.path.join(DATA, idx + '.tiff'), transform=identity,\n                                  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        self.shape = self.data.shape  # 25784*34937\n        self.reduce = reduce  # 4\n        self.sz = reduce * sz  # 4*256 = 1024\n        self.mask_grid = make_grid(self.shape, window=tile_resized, min_overlap=int(tile_resized * minoverlap))\n\n    def __len__(self):\n        return len(self.mask_grid)\n\n    def __getitem__(self, idx):\n        x1, x2, y1, y2 = self.mask_grid[idx]\n        if self.data.count == 3:\n            img = np.moveaxis(self.data.read([1, 2, 3], window=Window.from_slices((x1, x2), (y1, y2))), 0, -1)\n        else:\n            img = np.zeros((tile_resized, tile_resized, 3), np.uint8)\n            for i, layer in enumerate(self.layers):\n                img[:, :, i] = layer.read(1, window=Window.from_slices((x1, x2), (y1, y2)))\n\n        if self.reduce != 1:\n            img = cv2.resize(img, (tile_size, tile_size),\n                             interpolation=cv2.INTER_AREA)\n        # check for empty images\n        hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)\n        h, s, v = cv2.split(hsv)\n        vertices = torch.tensor([x1, x2, y1, y2])\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), vertices, -1\n        else:\n            return img2tensor((img / 255.0 - mean) / std), vertices, idx\n\n\nclass Model_pred:\n    def __init__(self, models, dl, tta: bool = TTA, 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        # x: images, z: vertices, y: idx\n        with torch.no_grad():\n            for x, z, y in iter(self.dl):\n                if (y >= 0).sum() > 0:  # exclude empty images\n                    x = x[y >= 0].to(device)\n                    z = z[y >= 0]\n                    y = y[y >= 0]\n                    if self.half:\n                        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:\n                            py = p\n                        else:\n                            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=reduce, mode=\"bilinear\")\n                    # 1 class\n                    py = py.permute(0, 2, 3, 1).float().cpu()\n\n                    py = py.squeeze(-1).numpy()\n                    z = z.numpy()\n\n                    batch_size = len(py)\n                    for i in range(batch_size):\n                        yield py[i], z[i], y[i]\n                        count += 1\n\n    def __len__(self):\n        return len(self.dl.dataset)\n\n    \n\n    \nmodels = []\nfor path in MODELS:\n    state_dict = torch.load(path, map_location=torch.device('cpu'))\n    if \"efficientnet-b4\" in path:\n        model = smp.Unet(\"efficientnet-b4\", encoder_weights=None, classes=1)\n    elif \"efficientnet-b2\" in path:\n        model = smp.Unet(\"efficientnet-b2\", encoder_weights=None, classes=1)\n    model.load_state_dict(state_dict)\n    model.float()\n    model.eval()\n    model.to(device)\n    models.append(model)\n\ndel state_dict\n\nnames, preds = [], []\n\nfor idx, row in tqdm(df_sample.iterrows(), total=len(df_sample)):\n    idx = row['id']\n    ds = HuBMAPDataset(idx)\n    dl = DataLoader(ds, batch_size=bs, pin_memory=True, shuffle=False, num_workers=0)\n    mp = Model_pred(models, dl)\n    \n    \n    # generate masks\n    mask = np.zeros(ds.shape, dtype=np.uint8)\n    for pred, vert, i in iter(mp):\n        x1, x2, y1, y2 = vert\n        mask[x1:x2, y1:y2] += (pred > TH).astype(np.uint8)\n    \n    mask = (mask > 0.5).astype(np.uint8)\n    \n    rle = rle_encode_less_memory(mask)\n    names.append(idx)\n    preds.append(rle)\n    del mask, ds, dl\n    gc.collect()\n","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)\ndisplay(df)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}