{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":61446,"databundleVersionId":6962461,"sourceType":"competition"},{"sourceId":1144103,"sourceType":"datasetVersion","datasetId":645458},{"sourceId":4053053,"sourceType":"datasetVersion","datasetId":2398033},{"sourceId":5220538,"sourceType":"datasetVersion","datasetId":3037137},{"sourceId":7284438,"sourceType":"datasetVersion","datasetId":4223999},{"sourceId":7445815,"sourceType":"datasetVersion","datasetId":4333929},{"sourceId":7526908,"sourceType":"datasetVersion","datasetId":4079395},{"sourceId":7529557,"sourceType":"datasetVersion","datasetId":4327897},{"sourceId":7545569,"sourceType":"datasetVersion","datasetId":4308668},{"sourceId":7551993,"sourceType":"datasetVersion","datasetId":4348159},{"sourceId":126502507,"sourceType":"kernelVersion"},{"sourceId":150248402,"sourceType":"kernelVersion"}],"dockerImageVersionId":30580,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install monai lovely-numpy -q --no-index --find-links=../input/vesuvis-downloads\n!python -m pip install -q --no-index --find-links=/kaggle/input/pip-download-for-segmentation-models-pytorch segmentation-models-pytorch\n!python -m pip install -q /kaggle/input/omegaconf222py3/omegaconf-2.2.2-py3-none-any.whl --no-index --find-links=/kaggle/input/omegaconf222py3/\n!pip install -q /kaggle/input/tensordict/tensordict-0.2.1-cp310-cp310-manylinux1_x86_64.whl","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_kg_hide-output":true,"execution":{"iopub.status.busy":"2024-01-31T13:48:50.119563Z","iopub.execute_input":"2024-01-31T13:48:50.120178Z","iopub.status.idle":"2024-01-31T13:49:25.465038Z","shell.execute_reply.started":"2024-01-31T13:48:50.120152Z","shell.execute_reply":"2024-01-31T13:49:25.463995Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import sys\nsys.path.append(\"../input/ttach-kaggle/\")","metadata":{"execution":{"iopub.status.busy":"2024-01-31T13:51:15.88551Z","iopub.execute_input":"2024-01-31T13:51:15.885811Z","iopub.status.idle":"2024-01-31T13:51:15.890562Z","shell.execute_reply.started":"2024-01-31T13:51:15.885786Z","shell.execute_reply":"2024-01-31T13:51:15.889635Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\n\nimport cv2\nfrom glob import glob\nfrom tqdm import tqdm\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nimport tensordict\n\nimport albumentations as A\nimport segmentation_models_pytorch as smp\nimport gc\nimport monai\nimport ttach as tta\nfrom typing import Union, Dict, Tuple\n\n\nimport re","metadata":{"execution":{"iopub.status.busy":"2024-01-31T13:51:15.891823Z","iopub.execute_input":"2024-01-31T13:51:15.892088Z","iopub.status.idle":"2024-01-31T13:52:14.850372Z","shell.execute_reply.started":"2024-01-31T13:51:15.892065Z","shell.execute_reply":"2024-01-31T13:52:14.849485Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DATASET_FOLDER = \"/kaggle/input/blood-vessel-segmentation\"\n\nis_test = not len(glob(os.path.join(DATASET_FOLDER, \"test\", \"*\", \"*\", \"*.tif\"))) == 6\nif is_test:\n    datasets = sorted(glob(f\"{DATASET_FOLDER}/test/*\"))[::-1]\nelse:\n    datasets = sorted(glob(f\"{DATASET_FOLDER}/train/kidney_2\"))\n\nprint(len(datasets))","metadata":{"execution":{"iopub.status.busy":"2024-01-31T13:52:14.851571Z","iopub.execute_input":"2024-01-31T13:52:14.852297Z","iopub.status.idle":"2024-01-31T13:52:14.889002Z","shell.execute_reply.started":"2024-01-31T13:52:14.852258Z","shell.execute_reply":"2024-01-31T13:52:14.887957Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def rename_keys(original_dict, pattern):\n    new_dict = {}\n    \n    for old_key, value in original_dict.items():\n        new_key = re.sub(pattern, '', old_key)\n        \n        new_dict[new_key] = value\n    \n    return new_dict\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    rle = ' '.join(str(x) for x in runs)\n    if rle=='':\n        rle = '1 0'\n    return rle\n\ndef find_highest_score_filename(file_list):\n    highest_score = float('-inf')\n    highest_score_filename = None\n\n    for filename in file_list:\n        # Extract the score from the filename using regular expression\n        match = re.search(r'dice_(\\d+\\.\\d+)', filename)\n        if match:\n            current_score = float(match.group(1))\n            if current_score > highest_score:\n                highest_score = current_score\n                highest_score_filename = filename\n\n    return highest_score_filename\n\ndef to_device(x: torch.Tensor, cuda_id: int = 0) -> torch.Tensor:\n    return x.cuda(cuda_id) if torch.cuda.is_available() else x\n\n\ndef load_jit_model(model_path: str, cuda_id: int = 0) -> torch.nn.Module:\n    model = torch.jit.load(\n        model_path,\n        map_location=f\"cuda:{cuda_id}\" if torch.cuda.is_available() else \"cpu\",\n    )\n    return model","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-01-31T13:52:14.890264Z","iopub.execute_input":"2024-01-31T13:52:14.890556Z","iopub.status.idle":"2024-01-31T13:52:14.901987Z","shell.execute_reply.started":"2024-01-31T13:52:14.89053Z","shell.execute_reply":"2024-01-31T13:52:14.901021Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"predict_on = [\n    [\n        \"Unet3d\",\n        192,\n        \"baseline_3d_unet_192_bs4_d4_scaled_pseudo_0.1_random\",\n        1.0,\n    ],\n]\n\n\nclass BuildDataset:\n    def __init__(self, dataset: str, is_test: bool = True):\n        self.ids = []\n        self.is_test = is_test\n\n        self.xmin, self.xmax = 0, 0\n\n        self.data_tensor = self.load_volume(dataset)\n        self.shape_orig = self.data_tensor.shape\n        \n    def normilize(self, image: np.ndarray) -> np.ndarray:\n        if image.dtype != np.half:\n            image = image.astype(np.half, copy=False)\n            \n        image -= self.xmin\n        image /= (self.xmax - self.xmin)\n        \n        np.clip(image, 0, 1, out=image)\n        return image\n    \n    @staticmethod\n    def norm_by_percentile(\n        volume: np.ndarray, low: float = 10, high: float = 99.8\n    ) -> Tuple:\n        xmin = np.percentile(volume, low)\n        print(xmin)\n        xmax = np.max([np.percentile(volume, high), 1])\n        print(xmax)\n        return int(xmin), int(xmax)\n\n    def load_volume(self, dataset: str) -> np.ndarray:\n        path = os.path.join(dataset, \"images\", \"*.tif\")\n        \n        dataset = sorted(glob(path)) if self.is_test else sorted(glob(path))[:192]\n\n        for p_img in tqdm(dataset):\n            path_ = p_img.split(os.path.sep)\n            slice_id, _ = os.path.splitext(path_[-1])\n            self.ids.append(f\"{path_[-3]}_{slice_id}\")\n\n        volume = None\n\n        for z, path in enumerate(tqdm(dataset)):\n            image = cv2.imread(path, cv2.IMREAD_ANYDEPTH).astype(np.half, copy=False)\n            \n            if volume is None:\n                volume = np.zeros((len(dataset), *image.shape[-2:]), dtype=np.float16)\n            volume[z, :, :] = image\n            \n        self.xmin, self.xmax = self.norm_by_percentile(volume)\n        return volume\n    \n    \nclass ModelWrapper(torch.nn.Module):\n    def __init__(self, base_model):\n        super(ModelWrapper, self).__init__()\n        self.base_model = base_model\n\n    def forward(self, x):\n        return torch.sigmoid(self.base_model(x)).half()","metadata":{"execution":{"iopub.status.busy":"2024-01-31T13:52:14.903438Z","iopub.execute_input":"2024-01-31T13:52:14.903751Z","iopub.status.idle":"2024-01-31T13:52:14.92001Z","shell.execute_reply.started":"2024-01-31T13:52:14.903719Z","shell.execute_reply":"2024-01-31T13:52:14.919193Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"tta_models = []\nweights = []\nfolds2predict = [0, 1]\n\nfor model_config in tqdm(predict_on):\n    for fold in folds2predict:\n        model_path = sorted(\n            glob(\n                f\"/kaggle/input/senet-3d-models/{model_config[2]}/{fold}/checkpoints/epoch*.ckpt\"\n            )\n        )[-1]\n        print(model_path)\n        state_dict = rename_keys(\n            torch.load(model_path, map_location=\"cpu\")[\"state_dict\"], \"net.\"\n        )\n        model_base = to_device(\n            monai.networks.nets.DynUNet(spatial_dims=3, in_channels=1, out_channels=1, kernel_size=[ [ 3, 3, 3 ], [ 3, 3, 3 ], [ 3, 3, 3 ], [ 3, 3, 3 ], [ 3, 3, 3 ], [ 3, 3, 3 ] ], strides=[ [ 1, 1, 1 ], [ 2, 2, 2 ], [ 2, 2, 2 ], [ 2, 2, 2 ], [ 2, 2, 2 ], [ 2, 2, 2 ] ], upsample_kernel_size=[[ 2, 2, 2 ], [ 2, 2, 2 ], [ 2, 2, 2 ], [ 2, 2, 2 ], [ 2, 2, 2 ]], dropout=0.2)\n        )\n        model_base.load_state_dict(state_dict)\n        model = ModelWrapper(model_base)\n\n        model.eval()\n\n        model = torch.nn.DataParallel(model)\n\n        tta_models.append(\n            tta.SegmentationTTAWrapper(\n                model.half(), tta.aliases.d4_transform(), merge_mode=\"mean\"\n            )\n        )\n\n        weights.append(model_config[-1])","metadata":{"execution":{"iopub.status.busy":"2024-01-31T13:52:14.923656Z","iopub.execute_input":"2024-01-31T13:52:14.924011Z","iopub.status.idle":"2024-01-31T13:52:20.545633Z","shell.execute_reply.started":"2024-01-31T13:52:14.923984Z","shell.execute_reply":"2024-01-31T13:52:20.544599Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"os.makedirs(f\"preds_3d\")\n\nrles, ids = [], []\nwith torch.no_grad():\n    for dataset in datasets:\n        folder = dataset.split(\"/\")[-1]\n\n        test_dataset = BuildDataset(dataset, is_test=is_test)\n\n        ids += test_dataset.ids\n        \n        preds = 0\n        input_tensor = tensordict.MemmapTensor.from_tensor(torch.from_numpy(test_dataset.normilize(test_dataset.data_tensor.astype(np.half))).unsqueeze(0).unsqueeze(0))\n        for tta_model, weight in zip(tta_models, weights):\n            preds += weight * monai.inferers.sliding_window_inference(\n                inputs=input_tensor, # if is_test else torch.rand(1, 1, 512, 512, 512),\n                predictor=tta_model,\n                sw_batch_size=2,\n                roi_size=(256, 256, 256),\n                overlap=0.25,\n                padding_mode=\"reflect\",\n                mode=\"gaussian\",\n                sw_device=\"cuda\",\n                device=\"cpu\",\n                progress=True,\n            ).squeeze().cpu().numpy().astype(np.half) / sum(weights)\n\n        for idx, pred in enumerate(preds):\n            cv2.imwrite(f\"preds_3d/{test_dataset.ids[idx]}.png\", (255*pred).astype(np.uint8))\n\n            \n        if is_test:\n            del input_tensor, test_dataset, preds\n            gc.collect()\n            torch.cuda.empty_cache()\n","metadata":{"execution":{"iopub.status.busy":"2024-01-31T13:52:20.547022Z","iopub.execute_input":"2024-01-31T13:52:20.547656Z","iopub.status.idle":"2024-01-31T13:59:59.201975Z","shell.execute_reply.started":"2024-01-31T13:52:20.54762Z","shell.execute_reply":"2024-01-31T13:59:59.201063Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if is_test:\n    del tta_models\n    gc.collect()\n    torch.cuda.empty_cache()","metadata":{"execution":{"iopub.status.busy":"2024-01-31T13:59:59.203105Z","iopub.execute_input":"2024-01-31T13:59:59.203419Z","iopub.status.idle":"2024-01-31T13:59:59.20775Z","shell.execute_reply.started":"2024-01-31T13:59:59.203392Z","shell.execute_reply":"2024-01-31T13:59:59.206894Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"predict_on = [\n     [\"UnetPlusPlus\", \"tu-tf_efficientnet_b5\", 1, \"UnetPlusPlus_tu-tf_efficientnet_b5_BoundaryDoULoss_size_1_512_bs32_hard_pseudo_v2\", 3., \"scse\", \"senet-models\", 800],  #0878 new sampling + cutmix\n]\n\n\ntta_models = []\nweights = []\n\nuse_top_only = True #True\nuse_best = False # False\nfolds2predict = [0, 1]\n\nuse_tta = True\n\nTH3d = 0.5 #()\nTH2d = 0.05\n\n","metadata":{"execution":{"iopub.status.busy":"2024-01-31T13:59:59.208979Z","iopub.execute_input":"2024-01-31T13:59:59.209269Z","iopub.status.idle":"2024-01-31T13:59:59.223731Z","shell.execute_reply.started":"2024-01-31T13:59:59.209246Z","shell.execute_reply":"2024-01-31T13:59:59.222862Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset","metadata":{}},{"cell_type":"code","source":"class BuildDataset(torch.utils.data.Dataset):\n    def __init__(self, dataset, in_channels=3, is_test=False):\n        self.window = in_channels // 2\n        self.is_test = is_test\n        self.ids = []\n\n        self.data_tensor = self.load_volume(dataset)\n        self.shape_orig = self.data_tensor.shape\n\n        padding = (\n            (self.window, self.window),\n        ) * self.data_tensor.ndim\n\n        self.padding = tuple(\n            (max(0, before), max(0, after)) for (before, after) in padding\n        )\n        self.data_tensor = np.pad(\n            self.data_tensor, padding, mode=\"constant\", constant_values=0\n        )\n\n    def __len__(self):\n        return sum(self.shape_orig) if self.is_test else self.shape_orig[0]\n\n    def normilize(self, image):\n        image = (image - self.xmin) / (\n                self.xmax - self.xmin)\n        image = np.clip(image, 0, 1)\n        return image.astype(np.float32)\n    \n    @staticmethod\n    def norm_by_percentile(volume, low=10, high=99.8):\n        xmin = np.percentile(volume, low)\n        print(xmin)\n        xmax = np.max([np.percentile(volume, high), 1])\n        print(xmax)\n        return xmin, xmax\n\n    def load_volume(self, dataset):\n        path = os.path.join(dataset, \"images\", \"*.tif\")\n        dataset = sorted(glob(path)) if self.is_test else sorted(glob(path))[:192]\n        for p_img in tqdm(dataset):\n            path_ = p_img.split(os.path.sep)\n            slice_id, _ = os.path.splitext(path_[-1])\n            self.ids.append(f\"{path_[-3]}_{slice_id}\")\n            \n        volume = None\n\n        for z, path in enumerate(tqdm(dataset)):\n            image = cv2.imread(path, cv2.IMREAD_ANYDEPTH)\n            image = np.array(image, dtype=np.uint16)\n            if volume is None:\n                volume = np.zeros((len(dataset), *image.shape[-2:]), dtype=np.uint16)\n            volume[z] = image\n        self.xmin, self.xmax = self.norm_by_percentile(volume)\n        return volume\n\n    def __getitem__(self, idx):\n        # Determine which axis to sample from based on the index\n        if idx < self.shape_orig[0]:\n            idx = idx + self.window\n            slice_data = self.normilize(\n                self.data_tensor[\n                    idx - self.window : 1 + idx + self.window, :, :\n                ].transpose(1, 2, 0)[self.window:-self.window, self.window:-self.window, :]\n            )\n            axis = \"X\"\n            idx -= self.window\n        elif idx < self.shape_orig[0] + self.shape_orig[1]:\n            idx -= (self.shape_orig[0] - self.window)\n            slice_data = self.normilize(\n                self.data_tensor[\n                    :, idx - self.window : 1 + idx + self.window, :\n                ].transpose(0, 2, 1)[self.window:-self.window, self.window:-self.window, :]\n            )\n            axis = \"Y\"\n            idx -= self.window\n\n            \n        else:\n            idx -= (\n                self.shape_orig[0]\n                + self.shape_orig[1]\n                - self.window\n            ) \n            \n            slice_data = self.normilize(\n                self.data_tensor[\n                    :, :, idx - self.window : 1 + idx + self.window\n                ][self.window:-self.window, self.window:-self.window, :]\n            )\n            axis = \"Z\"\n            idx -= self.window\n\n        slice_data = torch.tensor(slice_data.transpose(2, 0, 1))\n\n        return {\n            \"slice\": slice_data.half(),\n            \"slice_index\": idx,\n            \"axis\": axis\n        }","metadata":{"execution":{"iopub.status.busy":"2024-01-31T13:59:59.225002Z","iopub.execute_input":"2024-01-31T13:59:59.225264Z","iopub.status.idle":"2024-01-31T13:59:59.246204Z","shell.execute_reply.started":"2024-01-31T13:59:59.225241Z","shell.execute_reply":"2024-01-31T13:59:59.245509Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"execution":{"iopub.status.busy":"2024-01-31T13:59:59.288476Z","iopub.execute_input":"2024-01-31T13:59:59.288793Z","iopub.status.idle":"2024-01-31T13:59:59.298683Z","shell.execute_reply.started":"2024-01-31T13:59:59.288769Z","shell.execute_reply":"2024-01-31T13:59:59.297848Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"in_chans = []\nresolutions = []\n\nfor model_config in tqdm(predict_on):\n    for fold in folds2predict:\n        model_path = sorted(glob(f\"/kaggle/input/{model_config[6]}/{model_config[3]}/{fold}/checkpoints/last*.ckpt\"))[-1]\n        print(f\"use_top_only, loading: {model_path}\")\n        state_dict = rename_keys(torch.load(model_path, map_location=\"cpu\")[\"state_dict\"], \"net.\")\n        model = to_device(smp.create_model(arch=model_config[0], encoder_name=model_config[1], in_channels=model_config[2], encoder_weights=None, decoder_attention_type=model_config[5]))\n        model.load_state_dict(state_dict)\n        model.eval()\n\n        model = torch.nn.DataParallel(model)\n\n        if use_tta:\n            tta_models.append(tta.SegmentationTTAWrapper(model.half(), tta.aliases.flip_transform(), merge_mode='mean')) #flip_transform d4_transform\n        else:\n            tta_models.append(model)\n\n        weights.append(model_config[4])\n        in_chans.append(model_config[2])\n        resolutions.append(model_config[7])\n","metadata":{"execution":{"iopub.status.busy":"2024-01-31T13:59:59.300158Z","iopub.execute_input":"2024-01-31T13:59:59.300533Z","iopub.status.idle":"2024-01-31T14:00:07.82326Z","shell.execute_reply.started":"2024-01-31T13:59:59.300499Z","shell.execute_reply":"2024-01-31T14:00:07.822339Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"del state_dict","metadata":{"execution":{"iopub.status.busy":"2024-01-31T14:00:07.824674Z","iopub.execute_input":"2024-01-31T14:00:07.825073Z","iopub.status.idle":"2024-01-31T14:00:07.832001Z","shell.execute_reply.started":"2024-01-31T14:00:07.825035Z","shell.execute_reply":"2024-01-31T14:00:07.831009Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Inference","metadata":{}},{"cell_type":"code","source":"def merge_preds(mask1, mask2):\n    binary_mask = (255 * (mask1 > TH2d)).astype(np.uint8)\n    edged = cv2.Canny(binary_mask, 12, 200, L2gradient=True)\n    contours, hierarchy = cv2.findContours(edged,  \n        cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) \n    interest_mask = np.zeros_like(binary_mask)\n\n    if len(contours) > 0:\n        all_contours = np.vstack(contours[i] for i in range(len(contours)))\n        hull = cv2.convexHull(all_contours)\n        cv2.drawContours(interest_mask, [hull], -1, (1), thickness=cv2.FILLED)\n\n        interest_mask = cv2.dilate(interest_mask, np.ones((5, 5), np.uint8), iterations=5) \n        return ((interest_mask * mask2 + mask1) > TH2d + TH3d).astype(np.uint8)   \n    else:\n        return (mask1 > TH2d).astype(np.uint8)","metadata":{"execution":{"iopub.status.busy":"2024-01-31T14:14:34.283875Z","iopub.execute_input":"2024-01-31T14:14:34.284695Z","iopub.status.idle":"2024-01-31T14:14:34.293942Z","shell.execute_reply.started":"2024-01-31T14:14:34.28466Z","shell.execute_reply":"2024-01-31T14:14:34.292814Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"rles, ids = [], []\n\n\nwith torch.no_grad():\n    for dataset in datasets:\n        test_dataset = BuildDataset(dataset, is_test=is_test, in_channels=3) # TODO: refactor this\n        test_loader = DataLoader(test_dataset, batch_size=1, num_workers=4, shuffle=False, pin_memory=False)\n\n        y_preds = np.zeros(test_dataset.shape_orig, dtype=np.half)\n        ids += test_dataset.ids\n\n        pbar = tqdm(enumerate(test_loader), total=len(test_loader), desc=f'Inference {dataset}')\n        for step, batch in pbar:\n            images = to_device(batch[\"slice\"])\n            \n            axis = batch[\"axis\"][0]\n            idx = batch[\"slice_index\"].numpy()[0]\n\n            preds = 0\n            for tta_model, weight, in_chan, resolution in zip(tta_models, weights, in_chans, resolutions):\n                preds += weight * monai.inferers.sliding_window_inference(\n                    inputs=images.half() if in_chan != 1 else images[:, 1,...].unsqueeze(0).half(), # TODO: Refactor this\n                    predictor=tta_model.half(),\n                    sw_batch_size=8,\n                    roi_size=(resolution, resolution),\n                    overlap=0.25,\n                    padding_mode=\"reflect\",\n                    mode=\"gaussian\",\n                    sw_device=\"cuda\",\n                    device=\"cuda\",\n                    progress=False,\n                )\n            if axis == \"X\":\n                y_preds[idx, :, :] += ((preds / sum(weights)).squeeze().sigmoid().cpu().numpy() / 3.).astype(np.half)\n            elif axis == \"Y\":\n                y_preds[:, idx, :] += ((preds / sum(weights)).squeeze().sigmoid().cpu().numpy() / 3.).astype(np.half)\n            elif axis == \"Z\":\n                y_preds[:, :, idx] += ((preds / sum(weights)).squeeze().sigmoid().cpu().numpy() / 3.).astype(np.half)\n        \n        for idx, pred_2d in enumerate(y_preds):\n            \n            pred_3d = cv2.imread(f\"preds_3d/{test_dataset.ids[idx]}.png\", 0) / 255.\n            \n            rles.append(rle_encode(merge_preds(pred_2d, pred_3d)))\n\n        del test_dataset, test_loader, y_preds\n        gc.collect()\n        torch.cuda.empty_cache()","metadata":{"execution":{"iopub.status.busy":"2024-01-31T14:14:45.452094Z","iopub.execute_input":"2024-01-31T14:14:45.45251Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"del tta_models, tta_model, batch, preds, images, model\ngc.collect()\ntorch.cuda.empty_cache()\n","metadata":{"execution":{"iopub.status.busy":"2024-01-22T10:33:12.202646Z","iopub.execute_input":"2024-01-22T10:33:12.203002Z","iopub.status.idle":"2024-01-22T10:33:12.604875Z","shell.execute_reply.started":"2024-01-22T10:33:12.202968Z","shell.execute_reply":"2024-01-22T10:33:12.604117Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!rm -rf preds_3d","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission = pd.DataFrame.from_dict({\n    \"id\": ids,\n    \"rle\": rles\n})\nsubmission.to_csv(\"submission.csv\", index=False)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}