{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":61446,"databundleVersionId":6962461,"sourceType":"competition"},{"sourceId":7607219,"sourceType":"datasetVersion","datasetId":4429227},{"sourceId":158724263,"sourceType":"kernelVersion"}],"dockerImageVersionId":30636,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install /kaggle/input/save-smp/segmentation_models_pytorch/{segmentation_models_pytorch-0.3.3-py3-none-any.whl,pretrainedmodels-0.7.4-py3-none-any.whl,efficientnet_pytorch-0.7.1-py3-none-any.whl,timm-0.9.2-py3-none-any.whl,munch-4.0.0-py2.py3-none-any.whl}","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-02-11T17:43:46.019848Z","iopub.execute_input":"2024-02-11T17:43:46.020447Z","iopub.status.idle":"2024-02-11T17:44:22.492464Z","shell.execute_reply.started":"2024-02-11T17:43:46.020413Z","shell.execute_reply":"2024-02-11T17:44:22.491394Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import segmentation_models_pytorch as smp\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\n\nfrom glob import glob\nimport os\nimport logging\nfrom tqdm import tqdm\n\nimport tifffile\nimport numpy as np\nimport pandas as pd\n\nimport albumentations as A\nimport albumentations.pytorch as AT\nfrom torchvision import transforms\nimport gc","metadata":{"execution":{"iopub.status.busy":"2024-02-11T17:44:22.494902Z","iopub.execute_input":"2024-02-11T17:44:22.495332Z","iopub.status.idle":"2024-02-11T17:44:30.948234Z","shell.execute_reply.started":"2024-02-11T17:44:22.495280Z","shell.execute_reply":"2024-02-11T17:44:30.947210Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class UnetUpscale(nn.Module):\n    def __init__(\n        self,\n        encoder_name,\n        decoder_use_batchnorm,\n        in_channels,\n        classes,\n        encoder_weights,\n        upscale_factor,\n    ):\n        super().__init__()\n        self.upscale_factor = upscale_factor\n\n        self.model = smp.Unet(\n            encoder_name=encoder_name,\n            decoder_use_batchnorm=decoder_use_batchnorm,\n            in_channels=in_channels,\n            classes=classes,\n            encoder_weights=encoder_weights,\n        )\n\n    def forward(self, x):\n        x = torch.nn.functional.interpolate(\n            x, (x.shape[-2] * self.upscale_factor, x.shape[-1] * self.upscale_factor), mode=\"bilinear\"\n        )\n        x = self.model(x)\n        x = torch.nn.functional.interpolate(\n            x, (x.shape[-2] // self.upscale_factor, x.shape[-1] // self.upscale_factor), mode=\"bilinear\"\n        )\n        return x\n    \n    ","metadata":{"execution":{"iopub.status.busy":"2024-02-11T17:44:30.950043Z","iopub.execute_input":"2024-02-11T17:44:30.950715Z","iopub.status.idle":"2024-02-11T17:44:30.959576Z","shell.execute_reply.started":"2024-02-11T17:44:30.950673Z","shell.execute_reply":"2024-02-11T17:44:30.958491Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Dataset2DMultiPlanesTest(Dataset):\n    def __init__(\n        self,\n        full_image,\n        crop_size,\n        overlap_size,\n        planes,\n        transform=None,\n    ):\n        super().__init__()\n        step_size = crop_size - overlap_size\n        self.crop_size = crop_size\n        self.image = full_image\n         \n        self.depth, self.height, self.width = self.image.shape\n        \n        #calculate XY coordinates\n        xy_coordinates = []\n        if 'xy' in planes:\n            for z in range(self.depth):\n                for y in range(0, self.height - step_size, step_size):\n                    for x in range(0, self.width - step_size, step_size):\n                        crop_end_y = min(y + crop_size, self.height)\n                        crop_end_x = min(x + crop_size, self.width)\n\n                        xy_coordinates.append((z, z+1, y, crop_end_y, x, crop_end_x))\n\n        # calculate XZ coordinates\n        xz_coordinates = []\n        if 'xz' in planes:\n            for z in range(0, self.depth - step_size, step_size):\n                for y in range(self.height):\n                    for x in range(0, self.width - step_size, step_size):\n                        crop_end_z = min(z + crop_size, self.depth)\n                        crop_end_x = min(x + crop_size, self.width)\n\n                        xz_coordinates.append((z, crop_end_z, y, y+1, x, crop_end_x))\n\n        # calculate YZ coordinates\n        yz_coordinates = []\n        if 'yz' in planes:\n            for z in range(0, self.depth - step_size, step_size):\n                for y in range(0, self.height - step_size, step_size):\n                    for x in range(self.width):\n                        crop_end_z = min(z + crop_size, self.depth)\n                        crop_end_y = min(y + crop_size, self.height)\n\n                        yz_coordinates.append((z, crop_end_z, y, crop_end_y, x, x+1))\n\n        print(f'num xy slices: {len(xy_coordinates)} num xz slices: {len(xz_coordinates)} num yz slices: {len(yz_coordinates)}')\n        self.coordinates = xy_coordinates + xz_coordinates + yz_coordinates\n        print(f'total num of coordinates across 3 planes: {len(self.coordinates)}')\n\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.coordinates)\n\n    def __getitem__(self, idx):\n        coordinates = self.coordinates[idx]\n        z1, z2, y1, y2, x1, x2 = coordinates\n        \n        image_crop = self.image[z1:z2, y1:y2, x1:x2].copy().squeeze()\n        \n        height_pad_before = height_pad_after = width_pad_before = width_pad_after = 0\n        if image_crop.shape[0] != self.crop_size:\n            height_pad_size = self.crop_size - image_crop.shape[0]\n            height_pad_before = height_pad_size // 2\n            height_pad_after = height_pad_size - height_pad_before\n\n        if image_crop.shape[1] != self.crop_size:\n            width_pad_size = self.crop_size - image_crop.shape[1]\n            width_pad_before = width_pad_size // 2\n            width_pad_after = width_pad_size - width_pad_before\n            \n        image_crop = np.pad(image_crop, ((height_pad_before, height_pad_after), (width_pad_before, width_pad_after)), mode=\"constant\", constant_values=0)        \n               \n        if self.transform:\n            sample = self.transform(image=image_crop)\n            image_crop = sample['image']\n\n        image_mean = torch.mean(image_crop.float())\n        image_std = torch.std(image_crop.float())\n\n        image_crop = (image_crop - image_mean) / (image_std + 1e-4)\n        \n        return image_crop, torch.tensor([z1, z2, y1, y2, x1, x2]), torch.tensor([height_pad_before, height_pad_after, width_pad_before, width_pad_after])","metadata":{"execution":{"iopub.status.busy":"2024-02-11T17:44:30.962620Z","iopub.execute_input":"2024-02-11T17:44:30.962939Z","iopub.status.idle":"2024-02-11T17:44:30.984631Z","shell.execute_reply.started":"2024-02-11T17:44:30.962911Z","shell.execute_reply":"2024-02-11T17:44:30.983369Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def create_dataset(dataset_root):\n    paths = sorted(glob(f'{dataset_root}/*.tif'))\n    height, width = tifffile.memmap(paths[0], mode='r').shape\n    \n    full_image = np.zeros((len(paths), height, width), dtype=np.uint8)\n    \n    for path_index, path in enumerate(paths):\n        full_image[path_index] = (tifffile.imread(path) / 256).astype(np.uint8)\n    \n    return full_image","metadata":{"execution":{"iopub.status.busy":"2024-02-11T17:44:30.985748Z","iopub.execute_input":"2024-02-11T17:44:30.986120Z","iopub.status.idle":"2024-02-11T17:44:31.000207Z","shell.execute_reply.started":"2024-02-11T17:44:30.986082Z","shell.execute_reply":"2024-02-11T17:44:30.999147Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def predict(\n    batch_size,\n    num_workers,\n    dataset_params,\n    model_name,\n    model_params,\n    test_kidney=None,\n):  \n\n    if model_name == 'unet':\n        model = smp.Unet(\n            encoder_name=model_params['encoder_name'],\n            decoder_use_batchnorm=model_params['decoder_use_batchnorm'],\n            in_channels=1,\n            classes=1,\n            encoder_weights=None,\n        )\n    elif model_name == 'unet_upscale':\n        model = UnetUpscale(\n            encoder_name=model_params['encoder_name'],\n            decoder_use_batchnorm=model_params['decoder_use_batchnorm'],\n            upscale_factor=model_params['upscale_factor'],\n            in_channels=1,\n            classes=1,\n            encoder_weights=None,\n        )\n    else:\n        raise ValueError('Wrong model_name')\n\n\n    checkpoint = torch.load(model_params['checkpoint_path'], map_location='cpu')\n\n    model.load_state_dict(checkpoint['model'], strict=True)\n    model.cuda().eval();\n\n    transform = A.Compose(\n        [\n            AT.ToTensorV2(),\n        ]\n    )\n\n    dataset = Dataset2DMultiPlanesTest(\n        full_image=test_kidney,\n        crop_size=dataset_params['crop_size'],\n        overlap_size=dataset_params['overlap_size'],\n        planes=dataset_params['planes'],\n        transform=transform,\n    )\n\n    loader = DataLoader(\n        dataset=dataset,\n        shuffle=False,\n        drop_last=False,\n        batch_size=batch_size,\n        pin_memory=True,\n        num_workers=num_workers,\n    )\n\n\n    y_pred_shape = (loader.dataset.depth, loader.dataset.height, loader.dataset.width)\n    y_pred = torch.zeros(y_pred_shape, dtype=torch.float16)\n    y_stats = torch.zeros(y_pred_shape, dtype=torch.uint8)\n\n    for (input, coordinates, paddings) in (tqdm(loader)):\n        input = input.cuda()\n\n        with torch.no_grad():\n            with torch.cuda.amp.autocast(enabled=True):\n                preds = model(input)\n\n                for coordinates_sample, paddings_sample, preds_sample in zip(coordinates, paddings, preds):\n                    z1, z2, y1, y2, x1, x2 = coordinates_sample\n                    \n                    height_pad_before, height_pad_after, width_pad_before, width_pad_after = paddings_sample\n                    if height_pad_before:\n                        preds_sample = preds_sample[:, height_pad_before:, :]\n                    if height_pad_after:\n                        preds_sample = preds_sample[:, :-height_pad_after, :]\n                    if width_pad_before:\n                        preds_sample = preds_sample[:, :, width_pad_before:]\n                    if width_pad_after:\n                        preds_sample = preds_sample[:, :, :-width_pad_after]\n\n                    slice_shape = y_pred[z1:z2, y1:y2, x1:x2].shape\n\n                    y_pred[z1:z2, y1:y2, x1:x2] += preds_sample.view(slice_shape).cpu()\n                    y_stats[z1:z2, y1:y2, x1:x2] += 1\n\n\n    y_pred /= y_stats\n\n    del model, y_stats\n    gc.collect()\n\n    return y_pred\n\n\ndef rle_encode(img):\n    '''\n    img: numpy array, 1 - mask, 0 - background\n    Returns run length as string formated\n    '''\n    pixels = img.flatten()\n    pixels = np.concatenate([[0], pixels, [0]])\n    runs = np.where(pixels[1:] != pixels[:-1])[0] + 1\n    runs[1::2] -= runs[::2]\n    return ' '.join(str(x) for x in runs)","metadata":{"execution":{"iopub.status.busy":"2024-02-11T17:48:10.567947Z","iopub.execute_input":"2024-02-11T17:48:10.568341Z","iopub.status.idle":"2024-02-11T17:48:10.588450Z","shell.execute_reply.started":"2024-02-11T17:48:10.568308Z","shell.execute_reply":"2024-02-11T17:48:10.587153Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"th = 0.025","metadata":{"execution":{"iopub.status.busy":"2024-02-11T17:44:31.024068Z","iopub.execute_input":"2024-02-11T17:44:31.024398Z","iopub.status.idle":"2024-02-11T17:44:31.036572Z","shell.execute_reply.started":"2024-02-11T17:44:31.024360Z","shell.execute_reply":"2024-02-11T17:44:31.035568Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model1 = {\n    'batch_size': 2,\n    'num_workers': 2,\n    'dataset_params' : {\n            'crop_size': 512,\n            'overlap_size': 256,\n            'planes': ['xy', 'xz', 'yz'],\n        },\n    'model_name': 'unet',\n    'model_params':\n        {\n            'encoder_name': 'tu-maxvit_base_tf_512.in21k_ft_in1k',\n            'decoder_use_batchnorm': False,\n            'checkpoint_path': '/kaggle/input/sennet-final-weights/maxvit_base.pt/epoch_33_surface_dice_at_mean_0.8023.pt',\n        },\n}","metadata":{"execution":{"iopub.status.busy":"2024-02-11T17:44:31.037755Z","iopub.execute_input":"2024-02-11T17:44:31.038086Z","iopub.status.idle":"2024-02-11T17:44:31.047540Z","shell.execute_reply.started":"2024-02-11T17:44:31.038059Z","shell.execute_reply":"2024-02-11T17:44:31.046441Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model2 = {\n    'batch_size': 4,\n    'num_workers': 2,\n    'dataset_params' : {\n            'crop_size': 512,\n            'overlap_size': 256,\n            'planes': ['xy', 'xz', 'yz'],\n        },\n    'model_name': 'unet_upscale',\n    'model_params':\n        {\n            'encoder_name': 'tu-tf_efficientnetv2_s.in21k_ft_in1k',\n            'decoder_use_batchnorm': False,\n            'checkpoint_path': '/kaggle/input/sennet-final-weights/effnet_v2_m.pt/epoch_23_surface_dice_at_mean_0.8133.pt',\n            'upscale_factor': 2,\n        },\n}","metadata":{"execution":{"iopub.status.busy":"2024-02-11T17:44:31.048870Z","iopub.execute_input":"2024-02-11T17:44:31.049197Z","iopub.status.idle":"2024-02-11T17:44:31.058101Z","shell.execute_reply.started":"2024-02-11T17:44:31.049169Z","shell.execute_reply":"2024-02-11T17:44:31.057081Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model3 = {\n    'batch_size': 2,\n    'num_workers': 2,\n    'dataset_params' : {\n            'crop_size': 512,\n            'overlap_size': 256,\n            'planes': ['xy', 'xz', 'yz'],\n        },\n    'model_name': 'unet_upscale',\n    'model_params':\n        {\n            'encoder_name': 'tu-dpn68b',\n            'decoder_use_batchnorm': False,\n            'checkpoint_path': '/kaggle/input/sennet-final-weights/dpn_68.pt/epoch_38_surface_dice_at_mean_0.80779.pt',\n            'upscale_factor': 2,\n        },\n}","metadata":{"execution":{"iopub.status.busy":"2024-02-11T17:44:31.060895Z","iopub.execute_input":"2024-02-11T17:44:31.061268Z","iopub.status.idle":"2024-02-11T17:44:31.071888Z","shell.execute_reply.started":"2024-02-11T17:44:31.061234Z","shell.execute_reply":"2024-02-11T17:44:31.071000Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"models = [model1, model2, model3]","metadata":{"execution":{"iopub.status.busy":"2024-02-11T17:44:31.072954Z","iopub.execute_input":"2024-02-11T17:44:31.073240Z","iopub.status.idle":"2024-02-11T17:44:31.082576Z","shell.execute_reply.started":"2024-02-11T17:44:31.073214Z","shell.execute_reply":"2024-02-11T17:44:31.081615Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ids, rles = [], []\n\nfor test_kidney in [6,5]:\n    images_paths = sorted(glob(f'/kaggle/input/blood-vessel-segmentation/test/kidney_{test_kidney}/images/*.tif')) \n    test_kidney_image = create_dataset(\n        dataset_root=f'/kaggle/input/blood-vessel-segmentation/test/kidney_{test_kidney}/images/'\n    )\n\n    if test_kidney == 6:\n        private_res = 63.08\n        public_res = 50.0\n\n        scale = private_res / public_res\n\n        d_original, h_original, w_original = test_kidney_image.shape\n        test_kidney_image = torch.tensor(test_kidney_image).view(1, 1, d_original, h_original, w_original)\n        test_kidney_image = test_kidney_image.to(dtype=torch.float32)\n        test_kidney_image = torch.nn.functional.interpolate(test_kidney_image, (\n            int(d_original*scale),\n            int(h_original*scale),\n            int(w_original*scale),\n        ), mode='trilinear').squeeze().numpy()\n\n    for model_index, model in enumerate(models):\n        preds = predict(\n            **model,\n            test_kidney=test_kidney_image,\n        )\n\n        if model_index == 0:\n            preds_ensemble = preds\n        else:\n            preds_ensemble += preds\n\n        del preds\n        gc.collect()\n\n    del test_kidney_image\n    gc.collect()\n\n    preds_ensemble /= len(models)\n    if test_kidney == 6:\n        d_preds, h_preds, w_preds = preds_ensemble.shape \n        preds_ensemble = preds_ensemble.view(1, 1, d_preds, h_preds, w_preds)\n        preds_ensemble = preds_ensemble.to(dtype=torch.float32)\n\n        preds_ensemble = torch.nn.functional.interpolate(preds_ensemble, (\n            d_original,\n            h_original,\n            w_original,\n        ), mode='trilinear').squeeze()\n\n    preds_ensemble_th = torch.sigmoid(preds_ensemble.cuda()).cpu() > th\n    for pred_index, pred in enumerate(preds_ensemble_th):\n        ids.append(f'kidney_{test_kidney}_{images_paths[pred_index].split(\"/\")[-1].split(\".\")[0]}')\n        rle = rle_encode(pred)\n        if rle == '':\n            rle = '1 0'\n        rles.append(rle)\n\n    del preds_ensemble, preds_ensemble_th\n    gc.collect()\n\n\nsubmission = pd.DataFrame({\n    'id': ids,\n    'rle': rles,\n})","metadata":{"execution":{"iopub.status.busy":"2024-02-11T17:48:12.975322Z","iopub.execute_input":"2024-02-11T17:48:12.976232Z","iopub.status.idle":"2024-02-11T17:48:51.978595Z","shell.execute_reply.started":"2024-02-11T17:48:12.976192Z","shell.execute_reply":"2024-02-11T17:48:51.977748Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission.to_csv('submission.csv', index=False)","metadata":{"execution":{"iopub.status.busy":"2024-02-11T17:49:33.256509Z","iopub.execute_input":"2024-02-11T17:49:33.256946Z","iopub.status.idle":"2024-02-11T17:49:33.267467Z","shell.execute_reply.started":"2024-02-11T17:49:33.256910Z","shell.execute_reply":"2024-02-11T17:49:33.266565Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}