{"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":7527810,"sourceType":"datasetVersion","datasetId":4384599},{"sourceId":150248402,"sourceType":"kernelVersion"},{"sourceId":155873259,"sourceType":"kernelVersion"},{"sourceId":161265529,"sourceType":"kernelVersion"}],"dockerImageVersionId":30627,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!python -m pip install --no-index --find-links=/kaggle/input/pip-download-for-segmentation-models-pytorch segmentation-models-pytorch\n!python -m pip install --no-index --find-links=/kaggle/input/connected-components-3d connected-components-3d\n!python -m pip install --no-index --find-links=/kaggle/input/empatches empatches","metadata":{"execution":{"iopub.status.busy":"2024-02-01T07:07:47.085203Z","iopub.execute_input":"2024-02-01T07:07:47.085515Z","iopub.status.idle":"2024-02-01T07:08:06.239174Z","shell.execute_reply.started":"2024-02-01T07:07:47.085488Z","shell.execute_reply":"2024-02-01T07:08:06.238238Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport cv2\nimport matplotlib.pyplot as plt\nimport torch\nfrom torch.utils.data import DataLoader\nimport torch.nn as nn\nimport segmentation_models_pytorch as smp\nimport albumentations as A\nimport glob\nimport os\nfrom tqdm import tqdm\nfrom empatches import EMPatches\nimport joblib\nimport gc\nimport sys","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-02-01T07:28:52.543166Z","iopub.execute_input":"2024-02-01T07:28:52.543883Z","iopub.status.idle":"2024-02-01T07:28:52.55286Z","shell.execute_reply.started":"2024-02-01T07:28:52.543848Z","shell.execute_reply":"2024-02-01T07:28:52.551907Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CFG:\n    data_root_path = '/kaggle/input/blood-vessel-segmentation'\n    workers = 4\n    seed = 42\n    threshold = 0.2\n    test_overlap = 0.1\n    batch_size = 16\n    patch_size = 256\n    debug = False\n    \n    seg_model = \"Unet\" \n    encoder_name = \"tu-maxvit_tiny_tf_512\"\n\n    device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")","metadata":{"execution":{"iopub.status.busy":"2024-02-01T07:28:53.458092Z","iopub.execute_input":"2024-02-01T07:28:53.45851Z","iopub.status.idle":"2024-02-01T07:28:53.464791Z","shell.execute_reply.started":"2024-02-01T07:28:53.458474Z","shell.execute_reply":"2024-02-01T07:28:53.463802Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ls_images = glob.glob(os.path.join(CFG.data_root_path, \"test\", \"*\", \"*\", \"*.tif\"))\nprint(f\"found images: {len(ls_images)}\")\n\nif CFG.debug:\n    print('Debugging..')\n    ls_images = glob.glob(os.path.join(CFG.data_root_path, \"train\", \"kidney_2\", \"images\", \"*.tif\"))\n    print(f\"found images: {len(ls_images)}\")\n    kidney_ids = []\n    image_ids = []\n    for p_img in tqdm(ls_images):\n        path_ = p_img.split(os.path.sep)\n        # parse the submission ID\n        kidney_id = path_[-3]\n        image_id, _ = os.path.splitext(path_[-1])\n        kidney_ids.append(f\"{kidney_id}\")\n        image_ids.append(f\"{image_id}\")\nelse: \n    kidney_ids = []\n    image_ids = []\n    for p_img in tqdm(ls_images):\n        path_ = p_img.split(os.path.sep)\n        # parse the submission ID\n        kidney_id = path_[-3]\n        image_id, _ = os.path.splitext(path_[-1])\n        kidney_ids.append(f\"{kidney_id}\")\n        image_ids.append(f\"{image_id}\")","metadata":{"execution":{"iopub.status.busy":"2024-02-01T07:28:53.964453Z","iopub.execute_input":"2024-02-01T07:28:53.964804Z","iopub.status.idle":"2024-02-01T07:28:53.996818Z","shell.execute_reply.started":"2024-02-01T07:28:53.964778Z","shell.execute_reply":"2024-02-01T07:28:53.995974Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df = pd.DataFrame({'kidney_ids':kidney_ids, 'image_ids':image_ids}).sort_values('image_ids', ascending=True).reset_index(drop=True)","metadata":{"execution":{"iopub.status.busy":"2024-02-01T07:28:54.698859Z","iopub.execute_input":"2024-02-01T07:28:54.699223Z","iopub.status.idle":"2024-02-01T07:28:54.709557Z","shell.execute_reply.started":"2024-02-01T07:28:54.699192Z","shell.execute_reply":"2024-02-01T07:28:54.708615Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if CFG.debug:\n    print('shortening df for debugging')\n    df = df.iloc[900:900+256].copy()\n    print(df.shape)\n    df['folder'] = 'train'\nelse:\n    df['folder'] = 'test'","metadata":{"execution":{"iopub.status.busy":"2024-02-01T07:28:55.578902Z","iopub.execute_input":"2024-02-01T07:28:55.579646Z","iopub.status.idle":"2024-02-01T07:28:55.586323Z","shell.execute_reply.started":"2024-02-01T07:28:55.579611Z","shell.execute_reply":"2024-02-01T07:28:55.585323Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"os.mkdir('/kaggle/working/images/')\nos.mkdir('/kaggle/working/indices/')","metadata":{"execution":{"iopub.status.busy":"2024-02-01T07:28:56.946853Z","iopub.execute_input":"2024-02-01T07:28:56.947721Z","iopub.status.idle":"2024-02-01T07:28:56.988186Z","shell.execute_reply.started":"2024-02-01T07:28:56.947685Z","shell.execute_reply":"2024-02-01T07:28:56.986814Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def create_kidney_volume(kidney, folder):\n    images = []\n    kidney_ids = df[df['kidney_ids'] == kidney].sort_values('image_ids', ascending=True).image_ids\n    for image_id in tqdm(kidney_ids):\n        img = cv2.imread(f'/kaggle/input/blood-vessel-segmentation/{folder}/{kidney}/images/{str(image_id).zfill(4)}.tif',cv2.IMREAD_GRAYSCALE)\n        img = torch.from_numpy(img.copy())\n        img = img.to(torch.uint8)\n        images.append(img)\n    images = torch.stack(images)\n    return images, kidney_ids\n\ndef preprocess_image(image, lo, hi):\n    image = image.to(torch.float32)\n    image = (image - lo) / (hi - lo)\n    image = torch.clamp(image, min=0.5)\n    return image\n\ndef get_patch_id_list(data, truncate=0, return_indices=False):\n    emp = EMPatches()\n    img = data[0]\n    img_patches, image_indices = emp.extract_patches(img, patchsize=CFG.patch_size, overlap=CFG.test_overlap)\n    patch_ids = []\n    for image_id in range(data.shape[0]-truncate):\n        for patch in range(len(img_patches)):\n            patch_ids.append(f'{image_id}_{patch}')\n    if return_indices:\n        return patch_ids, len(img_patches), image_indices\n    return patch_ids\n\ndef create_grid(images):\n    row1 = np.concatenate([images[0, :, :],images[1, :, :]], axis=1)\n    row2 = np.concatenate([images[2, :, :],images[3, :, :]], axis=1)\n    image = row2 = np.concatenate([row1, row2], axis=0)\n    return image\n\nclass Dataset(torch.utils.data.Dataset):\n    def __init__(self, data, patch_ids):\n        self.data = data\n        self.patch_ids = patch_ids\n        self.emp = EMPatches()\n        self.lo, self.hi = np.percentile(data.numpy(), (2, 98))\n\n    def __getitem__(self, index):\n        orig_image_id, patch_id = self.patch_ids[index].split('_')\n        images = []\n        for i in range(4):\n            image_id = int(orig_image_id)+i\n            img = preprocess_image(self.data[image_id], self.lo, self.hi)\n            img_patches, image_indices = self.emp.extract_patches(img, patchsize=CFG.patch_size, overlap=CFG.test_overlap)\n            img = img_patches[int(patch_id)]\n            images.append(img)\n        images = np.stack(images)\n        image = create_grid(images)\n        \n        image = torch.tensor(image) \n        orig_image_id = torch.tensor(int(orig_image_id), dtype=torch.int16)\n        patch_id = torch.tensor(int(patch_id), dtype=torch.int8)\n        return image.unsqueeze(0), orig_image_id.unsqueeze(0),  patch_id.unsqueeze(0)\n\n    def __len__(self):\n        return len(self.patch_ids)","metadata":{"execution":{"iopub.status.busy":"2024-02-01T07:28:58.827193Z","iopub.execute_input":"2024-02-01T07:28:58.828052Z","iopub.status.idle":"2024-02-01T07:28:58.839596Z","shell.execute_reply.started":"2024-02-01T07:28:58.828015Z","shell.execute_reply":"2024-02-01T07:28:58.838654Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"torch.manual_seed(CFG.seed)\n\nseg_models = {\n    \"Unet\": smp.Unet,\n    \"Unet++\": smp.UnetPlusPlus,\n    \"MAnet\": smp.MAnet,\n    \"Linknet\": smp.Linknet,\n    \"FPN\": smp.FPN,\n    \"PSPNet\": smp.PSPNet,\n    \"PAN\": smp.PAN,\n    \"DeepLabV3\": smp.DeepLabV3,\n    \"DeepLabV3+\": smp.DeepLabV3Plus,\n}\n\nclass Model(torch.nn.Module):\n    def __init__(self):\n        super(Model, self).__init__()\n        self.model = seg_models[CFG.seg_model](\n            encoder_name=CFG.encoder_name,\n            encoder_weights=None, \n            in_channels=1,\n            classes=1,\n            activation=None,\n        )\n\n    def forward(self, images):\n        preds = self.model(images)\n        return preds","metadata":{"execution":{"iopub.status.busy":"2024-02-01T07:28:59.972516Z","iopub.execute_input":"2024-02-01T07:28:59.973321Z","iopub.status.idle":"2024-02-01T07:28:59.98205Z","shell.execute_reply.started":"2024-02-01T07:28:59.973285Z","shell.execute_reply":"2024-02-01T07:28:59.980994Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def rle_encode(mask):\n    pixel = mask.flatten()\n    pixel = np.concatenate([[0], pixel, [0]])\n    run = np.where(pixel[1:] != pixel[:-1])[0] + 1\n    run[1::2] -= run[::2]\n    rle = ' '.join(str(r) for r in run)\n    if rle == '':\n        rle = '1 0'\n    return rle","metadata":{"execution":{"iopub.status.busy":"2024-02-01T07:29:00.51437Z","iopub.execute_input":"2024-02-01T07:29:00.515205Z","iopub.status.idle":"2024-02-01T07:29:00.521162Z","shell.execute_reply.started":"2024-02-01T07:29:00.515171Z","shell.execute_reply":"2024-02-01T07:29:00.520105Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def rotate_grid(images, k=0):\n    row1 = torch.concat([torch.rot90(images[:,:,:CFG.patch_size, :CFG.patch_size], k=k, dims=[2, 3]), torch.rot90(images[:,:,:CFG.patch_size, CFG.patch_size:], k=k, dims=[2, 3])], axis=3)\n    row2 = torch.concat([torch.rot90(images[:,:,CFG.patch_size:, :CFG.patch_size], k=k, dims=[2, 3]), torch.rot90(images[:,:,CFG.patch_size:, CFG.patch_size:], k=k, dims=[2, 3])], axis=3)\n    image = torch.concat([row1, row2], axis=2)\n    return image\n\ndef rotate_grid_tta(model, images, rot90):\n    with torch.cuda.amp.autocast(enabled=True):\n        with torch.no_grad():\n            images = rotate_grid(images, k=rot90)\n            preds_tta = model(images)\n            preds_tta = nn.Sigmoid()(preds_tta)\n            preds_tta = rotate_grid(preds_tta, k=-rot90)\n            return preds_tta\n        \ndef predict_axis(kidney_volume, all_preds, axis=0):\n    print('Predicting axis:', axis)\n    if axis == 1:\n        kidney_volume = kidney_volume.permute(1,2,0)\n        all_preds = all_preds.transpose(1,2,0)\n    elif axis == 2:\n        kidney_volume = kidney_volume.permute(2,0,1)\n        all_preds = all_preds.transpose(2,0,1)\n    \n    kidney_patched_ids, num_patches_kidney, indices_kidney = get_patch_id_list(kidney_volume, truncate=3, return_indices=True)\n\n    test_dataset = Dataset(kidney_volume, kidney_patched_ids)\n    test_dataloader = DataLoader(test_dataset, batch_size=CFG.batch_size, shuffle=False, num_workers=CFG.workers,pin_memory=True)\n    \n    kidney_patched_ids = get_patch_id_list(kidney_volume, truncate=0)\n    test_dict = {}\n    test_num = {}\n    for id in kidney_patched_ids:\n        test_dict[id] = torch.zeros((CFG.patch_size,CFG.patch_size), device='cpu', dtype=torch.float16)\n        test_num[id] = 0\n    pbar = tqdm(enumerate(test_dataloader), total=len(test_dataloader), desc='Test ')\n    for step, (images, orig_image_id, patch_id) in pbar:        \n        images  = images.to(CFG.device, dtype=torch.float)\n        preds = torch.zeros(images.shape, device='cpu')\n\n        for i in range(4): \n            preds_tta = rotate_grid_tta(model, images, i)\n            preds += preds_tta.detach().cpu()\n        preds /= 4 \n        preds = preds.to(torch.float16)\n        orig_image_id = orig_image_id.cpu().numpy()\n        patch_id = patch_id.cpu().numpy()\n        for i, pred in enumerate(preds):\n            pred = pred.squeeze(0)\n            patches = torch.stack([pred[:CFG.patch_size, :CFG.patch_size], pred[:CFG.patch_size, CFG.patch_size:], pred[CFG.patch_size:, :CFG.patch_size], pred[CFG.patch_size:, CFG.patch_size:]])\n            for x, patch in enumerate(patches):\n                image_id = orig_image_id[i].item()+x\n                image_patch_id = f'{str(image_id)}_{patch_id[i].item()}'\n                test_dict[image_patch_id] += patch\n                test_num[image_patch_id] += 1\n    torch.cuda.empty_cache()\n    gc.collect()     \n    \n    for id in kidney_patched_ids:\n        test_dict[id] /= test_num[id]\n    \n    print('Adding to predictions:')\n    for x, image_id in enumerate(tqdm(range(kidney_volume.shape[0]))):\n        preds_image = []\n        for i in range(num_patches_kidney):\n            patch_id = f'{image_id}_{i}'\n            preds_image.append(test_dict[patch_id].cpu().numpy())\n            test_dict[patch_id] = 0\n        merged_preds = emp.merge_patches(preds_image, indices_kidney, mode='max')\n        merged_preds = merged_preds.astype(np.float16)\n        all_preds[x] += merged_preds        \n        torch.cuda.empty_cache()\n        gc.collect()\n    print(all_preds.shape)   \n    if axis == 1:\n        kidney_volume = kidney_volume.permute(2,0,1)\n        all_preds = all_preds.transpose(2,0,1)\n    elif axis == 2:\n        kidney_volume = kidney_volume.permute(1,2,0)\n        all_preds = all_preds.transpose(1,2,0)\n    print(all_preds.shape)\n        \n    return all_preds\n\ndef predict_kidney(kidney='kidney_2', folder='train'):\n    \n    kidney_volume, kidney_ids = create_kidney_volume(kidney, folder)\n    all_preds = np.zeros(kidney_volume.shape, dtype=np.float16)\n    all_preds = predict_axis(kidney_volume, all_preds, axis=0)\n    all_preds = predict_axis(kidney_volume, all_preds, axis=1)\n    all_preds = predict_axis(kidney_volume, all_preds, axis=2)\n    counter = 3\n    all_preds /= counter\n    \n    all_preds = (all_preds>CFG.threshold).astype(np.int8)\n    all_rle = []\n    for pred in all_preds:\n        rle = rle_encode(pred)\n        all_rle.append(rle)\n    submission = pd.DataFrame.from_dict({\n    \"id\": kidney_ids,\n    \"rle\": all_rle\n    })\n    submission.id = submission.id.apply(lambda x: f'{kidney}_{x}')\n    return submission","metadata":{"execution":{"iopub.status.busy":"2024-02-01T07:29:01.251444Z","iopub.execute_input":"2024-02-01T07:29:01.252116Z","iopub.status.idle":"2024-02-01T07:29:01.280257Z","shell.execute_reply.started":"2024-02-01T07:29:01.252082Z","shell.execute_reply":"2024-02-01T07:29:01.279332Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"weights = torch.load(f'/kaggle/input/model2-5d-v17/tu-maxvit_tiny_tf_512-Unet_loss_V17_last_3epochs.pth', map_location=CFG.device) #['state_dict']\nmodel = Model()\nmodel.to(CFG.device)\nmodel.load_state_dict(weights)\nemp = EMPatches()\n\nmodel.eval()\n\nif CFG.debug:\n    submission = predict_kidney(kidney='kidney_2', folder='train')\n    \nelse:\n    df_kidney_5 = predict_kidney(kidney='kidney_5', folder='test')\n    torch.cuda.empty_cache()\n    gc.collect()\n    df_kidney_6 = predict_kidney(kidney='kidney_6', folder='test')\n    torch.cuda.empty_cache()\n    gc.collect()\n    submission = pd.concat([df_kidney_5,df_kidney_6]).reset_index(drop=True)\n","metadata":{"execution":{"iopub.status.busy":"2024-02-01T07:29:02.130792Z","iopub.execute_input":"2024-02-01T07:29:02.131529Z","iopub.status.idle":"2024-02-01T07:32:29.729962Z","shell.execute_reply.started":"2024-02-01T07:29:02.131494Z","shell.execute_reply":"2024-02-01T07:32:29.728874Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission.to_csv(\"submission.csv\", index=False)","metadata":{"execution":{"iopub.status.busy":"2024-02-01T07:33:00.198006Z","iopub.execute_input":"2024-02-01T07:33:00.198985Z","iopub.status.idle":"2024-02-01T07:33:00.218448Z","shell.execute_reply.started":"2024-02-01T07:33:00.198946Z","shell.execute_reply":"2024-02-01T07:33:00.217593Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission","metadata":{"execution":{"iopub.status.busy":"2024-02-01T07:33:00.571166Z","iopub.execute_input":"2024-02-01T07:33:00.571577Z","iopub.status.idle":"2024-02-01T07:33:00.588768Z","shell.execute_reply.started":"2024-02-01T07:33:00.571545Z","shell.execute_reply":"2024-02-01T07:33:00.587887Z"},"trusted":true},"execution_count":null,"outputs":[]}]}