{"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":"# Exploratory Data Analysis🔎","metadata":{}},{"cell_type":"code","source":"!pip uninstall -y torchtext\n# !pip install -q --upgrade torch torchvision\n!pip install -q \"lightning-flash[image]\" \"torchmetrics<0.8\" --no-index --find-links ../input/demo-flash-semantic-segmentation/frozen_packages\n!pip install -q -U timm segmentation-models-pytorch --no-index --find-links ../input/demo-flash-semantic-segmentation/frozen_packages\n!pip install -q 'kaggle-image-segmentation' --no-index --find-links ../input/tract-segm-eda-3d-interactive-viewer/frozen_packages\n\n! pip list | grep torch\n! pip list | grep lightning\n! nvidia-smi -L","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-06-02T21:26:57.471056Z","iopub.execute_input":"2022-06-02T21:26:57.471359Z","iopub.status.idle":"2022-06-02T21:28:17.379214Z","shell.execute_reply.started":"2022-06-02T21:26:57.47128Z","shell.execute_reply":"2022-06-02T21:28:17.378344Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os, glob\nimport pandas as pd\nimport matplotlib.pyplot as plt\n\nDATASET_FOLDER = \"/kaggle/input/uw-madison-gi-tract-image-segmentation\"\ndf_train = pd.read_csv(os.path.join(DATASET_FOLDER, \"train.csv\"))\ndisplay(df_train.head())\n\ndf_pred = pd.read_csv(os.path.join(DATASET_FOLDER, \"sample_submission.csv\"))\nWITH_SUBMISSION = not df_pred.empty","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-06-02T21:28:17.381508Z","iopub.execute_input":"2022-06-02T21:28:17.381793Z","iopub.status.idle":"2022-06-02T21:28:17.899374Z","shell.execute_reply.started":"2022-06-02T21:28:17.381757Z","shell.execute_reply":"2022-06-02T21:28:17.898723Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"all_imgs = glob.glob(os.path.join(DATASET_FOLDER, \"train\", \"case*\", \"case*_day*\", \"scans\", \"*.png\"))\nall_imgs = [p.replace(DATASET_FOLDER, \"\") for p in all_imgs]\n\nprint(f\"images: {len(all_imgs)}\")\nprint(f\"annotated: {len(df_train['id'].unique())}\")","metadata":{"execution":{"iopub.status.busy":"2022-06-02T21:28:17.900877Z","iopub.execute_input":"2022-06-02T21:28:17.901144Z","iopub.status.idle":"2022-06-02T21:28:21.389925Z","shell.execute_reply.started":"2022-06-02T21:28:17.901108Z","shell.execute_reply":"2022-06-02T21:28:21.389135Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from pprint import pprint\nfrom kaggle_imsegm.data_io import extract_tract_details\n\npprint(extract_tract_details(df_train['id'].iloc[0], DATASET_FOLDER))\n\ndf_train[['Case','Day','Slice', 'image', 'image_path', 'height', 'width']] = df_train['id'].apply(\n    lambda x: pd.Series(extract_tract_details(x, DATASET_FOLDER))\n)\ndisplay(df_train.head())","metadata":{"execution":{"iopub.status.busy":"2022-06-02T21:28:21.392214Z","iopub.execute_input":"2022-06-02T21:28:21.392666Z","iopub.status.idle":"2022-06-02T21:30:32.644017Z","shell.execute_reply.started":"2022-06-02T21:28:21.392626Z","shell.execute_reply":"2022-06-02T21:30:32.643292Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Browse the 3D image\n\nsee the full version (without importing own package) in https://www.kaggle.com/code/jirkaborovec/tract-segm-eda-3d-data-browser","metadata":{}},{"cell_type":"code","source":"from ipywidgets import interact, IntSlider\nfrom kaggle_imsegm.data_io import load_volume_from_images, create_tract_segmentation\nfrom kaggle_imsegm.visual import show_tract_volume\n\nCASE = 108\nDAY = 10\nIMAGE_FOLDER = os.path.join(DATASET_FOLDER, \"train\", f\"case{CASE}\", f\"case{CASE}_day{DAY}\", \"scans\")\nvol = load_volume_from_images(img_dir=IMAGE_FOLDER)\nprint(vol.shape)\n\ndf_ = df_train[(df_train[\"Case\"] == CASE) & (df_train[\"Day\"] == DAY)]\nsegm = create_tract_segmentation(df_vol=df_, vol_shape=vol.shape)\n\ndef interactive_show(volume):\n    vol_shape = volume.shape\n    interact(\n        lambda x, y, z: plt.show(show_tract_volume(volume, segm, z, y, x)),\n        z=IntSlider(min=0, max=vol_shape[0], step=5, value=int(vol_shape[0] / 2)),\n        y=IntSlider(min=0, max=vol_shape[1], step=5, value=int(vol_shape[1] / 2)),\n        x=IntSlider(min=0, max=vol_shape[2], step=5, value=int(vol_shape[2] / 2)),\n    )","metadata":{"execution":{"iopub.status.busy":"2022-06-02T21:30:32.645207Z","iopub.execute_input":"2022-06-02T21:30:32.647626Z","iopub.status.idle":"2022-06-02T21:30:36.10418Z","shell.execute_reply.started":"2022-06-02T21:30:32.647583Z","shell.execute_reply":"2022-06-02T21:30:36.102635Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"interactive_show(vol)","metadata":{"execution":{"iopub.status.busy":"2022-06-02T21:30:36.105707Z","iopub.execute_input":"2022-06-02T21:30:36.105954Z","iopub.status.idle":"2022-06-02T21:30:36.749106Z","shell.execute_reply.started":"2022-06-02T21:30:36.105919Z","shell.execute_reply":"2022-06-02T21:30:36.748436Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Prepare flatten dataset","metadata":{}},{"cell_type":"code","source":"DATASET_IMAGES = \"/kaggle/temp/dataset-flash/images\"\nDATASET_SEGMS = \"/kaggle/temp/dataset-flash/segms\"\n\nfor rdir in (DATASET_IMAGES, DATASET_SEGMS):\n    for sdir in (\"train\", \"val\"):\n        os.makedirs(os.path.join(rdir, sdir), exist_ok=True)","metadata":{"execution":{"iopub.status.busy":"2022-06-02T21:30:36.75031Z","iopub.execute_input":"2022-06-02T21:30:36.752034Z","iopub.status.idle":"2022-06-02T21:30:36.765945Z","shell.execute_reply.started":"2022-06-02T21:30:36.751992Z","shell.execute_reply":"2022-06-02T21:30:36.765288Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_train['Case_Day'] = [f\"case{r['Case']}_day{r['Day']}\" for _, r in df_train.iterrows()]\n\nCASES_DAYS = list(df_train['Case_Day'].unique())\nVAL_SPLIT = 0.01 if WITH_SUBMISSION else 0.1\nVAL_CASES_DAYS = CASES_DAYS[-int(VAL_SPLIT * len(CASES_DAYS)):]\n\nprint(f\"all case-day: {len(CASES_DAYS)}\")\nprint(f\"val case-day: {len(VAL_CASES_DAYS)}\")","metadata":{"execution":{"iopub.status.busy":"2022-06-02T21:30:36.767511Z","iopub.execute_input":"2022-06-02T21:30:36.767971Z","iopub.status.idle":"2022-06-02T21:30:42.53716Z","shell.execute_reply.started":"2022-06-02T21:30:36.767933Z","shell.execute_reply":"2022-06-02T21:30:42.535663Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nfrom PIL import Image\nfrom joblib import Parallel, delayed\nfrom kaggle_imsegm.data_io import preprocess_tract_scan\n\nLABELS = sorted(df_train[\"class\"].unique())\nprint(LABELS)\n\ndef _chose_sfolder(df_, val_cases_days=VAL_CASES_DAYS) -> str:\n    case, day = df_.iloc[0][[\"Case\", \"Day\"]]\n    case_day = f\"case{case}_day{day}\"\n    return 'val' if case_day in val_cases_days else 'train'","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-06-02T21:30:42.538432Z","iopub.execute_input":"2022-06-02T21:30:42.539128Z","iopub.status.idle":"2022-06-02T21:30:42.602549Z","shell.execute_reply.started":"2022-06-02T21:30:42.539088Z","shell.execute_reply":"2022-06-02T21:30:42.601715Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from tqdm.auto import tqdm\n\n_args = dict(\n    dir_data=os.path.join(DATASET_FOLDER, \"train\"),\n    dir_imgs=DATASET_IMAGES,\n    dir_segm=DATASET_SEGMS,\n    labels=LABELS,\n)\n_= Parallel(n_jobs=6)(\n    delayed(preprocess_tract_scan)(dfg, sfolder=_chose_sfolder(dfg), **_args)\n    for _, dfg in tqdm(df_train.groupby(\"Case_Day\"))\n)","metadata":{"execution":{"iopub.status.busy":"2022-06-02T21:30:42.60565Z","iopub.execute_input":"2022-06-02T21:30:42.606182Z","iopub.status.idle":"2022-06-02T21:50:39.787065Z","shell.execute_reply.started":"2022-06-02T21:30:42.606151Z","shell.execute_reply":"2022-06-02T21:50:39.78626Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"spl_imgs = glob.glob(os.path.join(DATASET_IMAGES, \"*\", \"*.png\"))[:3]\nfig, axarr = plt.subplots(ncols=2, nrows=len(spl_imgs), figsize=(6, 3 * len(spl_imgs)))\n\nfor i, img in enumerate(spl_imgs):\n    segm = img.replace(DATASET_IMAGES, DATASET_SEGMS)\n    axarr[i, 0].imshow(plt.imread(img))\n    axarr[i, 1].imshow(plt.imread(segm))\nplt.tight_layout()","metadata":{"execution":{"iopub.status.busy":"2022-06-02T21:50:39.78885Z","iopub.execute_input":"2022-06-02T21:50:39.789138Z","iopub.status.idle":"2022-06-02T21:50:40.88895Z","shell.execute_reply.started":"2022-06-02T21:50:39.789101Z","shell.execute_reply":"2022-06-02T21:50:40.88827Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Lightning⚡Flash & DeepLab-v3 & albumentations\n\nlets follow the Semantinc segmentation example: https://lightning-flash.readthedocs.io/en/stable/reference/semantic_segmentation.html","metadata":{}},{"cell_type":"code","source":"import torch\n\nimport flash\nfrom flash.core.data.utils import download_data\nfrom flash.image import SemanticSegmentation, SemanticSegmentationData","metadata":{"execution":{"iopub.status.busy":"2022-06-02T21:50:40.890323Z","iopub.execute_input":"2022-06-02T21:50:40.890792Z","iopub.status.idle":"2022-06-02T21:50:51.037294Z","shell.execute_reply.started":"2022-06-02T21:50:40.890757Z","shell.execute_reply":"2022-06-02T21:50:51.036433Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 1. Create the DataModule","metadata":{}},{"cell_type":"code","source":"from dataclasses import dataclass\nfrom typing import Any, Callable, Dict, Mapping, Sequence, Tuple, Union\nimport albumentations as alb\n\nfrom flash.core.data.io.input_transform import InputTransform\nfrom flash.image.segmentation.input_transform import prepare_target, remove_extra_dimensions\nfrom kaggle_imsegm.transform import FlashAlbumentationsAdapter\n\nclass NormImage(alb.DualTransform):\n\n    def __init__(self, quantile: float = 0.01, norm: bool = True, always_apply=False, p=1):\n        super().__init__(always_apply, p)\n        self.quantile = quantile\n        self.norm = norm\n\n    def apply(self, img, **params):\n        if self.quantile > 0:\n            q_low, q_high = np.percentile(img, [self.quantile * 100, (1 - self.quantile) * 100])\n            img = np.clip(img, q_low, q_high)\n        if self.norm:\n            v_min, v_max = np.min(img), np.max(img)\n            img = (img - v_min) / float(v_max - v_min)\n        return img\n\n    def apply_to_mask(self, mask, **params):\n        # Bounding box coordinates are scale invariant\n        return mask\n\n@dataclass\nclass SemanticSegmentationInputTransform(InputTransform):\n    # https://albumentations.ai/docs/examples/pytorch_semantic_segmentation\n\n    image_size: Tuple[int, int] = (128, 128)\n\n    def train_per_sample_transform(self) -> Callable:\n        return FlashAlbumentationsAdapter([\n            NormImage(always_apply=True),\n            alb.Resize(*self.image_size),\n            alb.VerticalFlip(p=0.5),\n            alb.HorizontalFlip(p=0.5),\n            alb.RandomRotate90(p=0.5),\n            alb.ShiftScaleRotate(shift_limit=0.15, scale_limit=0.05, rotate_limit=5, p=1.),\n            alb.GaussNoise(var_limit=(0.001, 0.01), mean=0, per_channel=False, p=1.0),\n            # alb.OneOf([\n            #     alb.GridDistortion(num_steps=5, distort_limit=0.05, p=1.0),\n            #     alb.ElasticTransform(alpha=1, sigma=50, alpha_affine=50, p=1.0),\n            # ], p=0.25),\n            alb.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.8),\n        ])\n\n    def per_sample_transform(self) -> Callable:\n        return FlashAlbumentationsAdapter([NormImage(always_apply=True), alb.Resize(*self.image_size)])\n\n    def target_per_batch_transform(self) -> Callable:\n        return prepare_target\n\n    def predict_per_batch_transform(self) -> Callable:\n        return remove_extra_dimensions\n\n    def serve_per_batch_transform(self) -> Callable:\n        return remove_extra_dimensions","metadata":{"execution":{"iopub.status.busy":"2022-06-02T23:21:51.817629Z","iopub.execute_input":"2022-06-02T23:21:51.817972Z","iopub.status.idle":"2022-06-02T23:21:51.835407Z","shell.execute_reply.started":"2022-06-02T23:21:51.817937Z","shell.execute_reply":"2022-06-02T23:21:51.834531Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"IMAGE_SIZE = (320, 320)\n\ndatamodule = SemanticSegmentationData.from_folders(\n    train_folder=os.path.join(DATASET_IMAGES, 'train'),\n    train_target_folder=os.path.join(DATASET_SEGMS, 'train'),\n    val_folder=os.path.join(DATASET_IMAGES, 'val'),\n    val_target_folder=os.path.join(DATASET_SEGMS, 'val'),\n    predict_folder=os.path.join(DATASET_IMAGES, 'val'),\n    #val_split=0.1,\n    train_transform=SemanticSegmentationInputTransform,\n    val_transform=SemanticSegmentationInputTransform,\n    predict_transform=SemanticSegmentationInputTransform,\n    transform_kwargs=dict(image_size=IMAGE_SIZE),\n    num_classes=len(LABELS) + 1,\n    batch_size=9,\n    num_workers=3,\n)","metadata":{"execution":{"iopub.status.busy":"2022-06-02T23:21:23.11572Z","iopub.execute_input":"2022-06-02T23:21:23.115983Z","iopub.status.idle":"2022-06-02T23:21:23.660557Z","shell.execute_reply.started":"2022-06-02T23:21:23.115952Z","shell.execute_reply":"2022-06-02T23:21:23.65971Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# datamodule.show_train_batch()\n\nfig, axarr = plt.subplots(ncols=2, nrows=5, figsize=(8, 20))\nrunning_i = 0\n\nfor batch in datamodule.train_dataloader():\n    print(batch.keys())\n    for i in range(len(batch['input'])):\n        segm = batch['target'][i].numpy()\n        if np.sum(segm) == 0 or np.max(segm) <= 1:\n            continue\n        img = np.rollaxis(batch['input'][i].cpu().numpy(), 0, 3)\n        axarr[running_i, 0].imshow(img)\n        seg = axarr[running_i, 1].imshow(segm)\n        plt.colorbar(seg, ax=axarr[running_i, 1])\n        running_i += 1\n        if running_i >= 5:\n            break\n    if running_i >= 5:\n        break","metadata":{"execution":{"iopub.status.busy":"2022-06-02T23:21:25.004643Z","iopub.execute_input":"2022-06-02T23:21:25.005236Z","iopub.status.idle":"2022-06-02T23:21:29.563249Z","shell.execute_reply.started":"2022-06-02T23:21:25.005195Z","shell.execute_reply":"2022-06-02T23:21:29.562594Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fig, axarr = plt.subplots(ncols=4, nrows=1, figsize=(12, 3))\n\nfor batch in datamodule.predict_dataloader():\n    print(batch.keys())\n    for i in range(len(batch['input'])):\n        if i >= 4:\n            break\n        img = np.rollaxis(batch['input'][i].cpu().numpy(), 0, 3)\n        axarr[i].imshow(img)\n    break","metadata":{"execution":{"iopub.status.busy":"2022-06-02T23:15:30.224494Z","iopub.execute_input":"2022-06-02T23:15:30.225281Z","iopub.status.idle":"2022-06-02T23:15:31.457295Z","shell.execute_reply.started":"2022-06-02T23:15:30.225238Z","shell.execute_reply":"2022-06-02T23:15:31.456536Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 2. Build the task","metadata":{}},{"cell_type":"code","source":"import segmentation_models_pytorch as smp\n\nmodel = SemanticSegmentation(\n    backbone=\"efficientnet-b3\",\n    head=\"deeplabv3\",\n    pretrained=False,\n    optimizer=\"asgd\",\n    learning_rate=0.05,\n    lr_scheduler=(\"StepLR\", {\"step_size\": 500}),\n    # lr_scheduler=(\"cosineannealinglr\", {\"T_max\": 500, \"eta_min\": 1e-6}),\n    loss_fn=smp.losses.DiceLoss(mode=\"multiclass\"),\n    num_classes=datamodule.num_classes,\n)","metadata":{"execution":{"iopub.status.busy":"2022-06-02T22:32:52.163376Z","iopub.execute_input":"2022-06-02T22:32:52.163961Z","iopub.status.idle":"2022-06-02T22:32:52.350155Z","shell.execute_reply.started":"2022-06-02T22:32:52.163922Z","shell.execute_reply":"2022-06-02T22:32:52.349406Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 3. Create the trainer and finetune the model","metadata":{}},{"cell_type":"code","source":"import gc\nimport pytorch_lightning as pl\n\ngc.collect()\ntorch.cuda.empty_cache()\ntrainer = flash.Trainer(\n    max_epochs=9 if WITH_SUBMISSION else 5,\n    logger=pl.loggers.CSVLogger(save_dir='logs/'),\n    gpus=torch.cuda.device_count(),\n    # precision=16 if torch.cuda.is_available() else 32,\n    accumulate_grad_batches=24,\n    gradient_clip_val=0.01,\n    limit_train_batches=1.0 if WITH_SUBMISSION else 0.2,\n    limit_val_batches=1.0 if WITH_SUBMISSION else 0.3,\n)","metadata":{"execution":{"iopub.status.busy":"2022-06-02T22:32:52.35194Z","iopub.execute_input":"2022-06-02T22:32:52.352273Z","iopub.status.idle":"2022-06-02T22:32:52.822281Z","shell.execute_reply.started":"2022-06-02T22:32:52.352233Z","shell.execute_reply":"2022-06-02T22:32:52.821485Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import warnings\nwarnings.filterwarnings(\"ignore\", category=DeprecationWarning) \n\n# Train the model\ntrainer.finetune(model, datamodule=datamodule, strategy=\"no_freeze\")\n\n# Save the model!\ntrainer.save_checkpoint(\"semantic_segmentation_model.pt\")","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-06-02T22:32:52.823869Z","iopub.execute_input":"2022-06-02T22:32:52.824397Z","iopub.status.idle":"2022-06-02T22:53:52.502644Z","shell.execute_reply.started":"2022-06-02T22:32:52.82435Z","shell.execute_reply":"2022-06-02T22:53:52.501817Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import seaborn as sn\n\nmetrics = pd.read_csv(f'{trainer.logger.log_dir}/metrics.csv')\ndel metrics[\"step\"]\nmetrics.set_index(\"epoch\", inplace=True)\n# display(metrics.dropna(axis=1, how=\"all\").head())\ng = sn.relplot(data=metrics, kind=\"line\")\nplt.gcf().set_size_inches(12, 4)\nplt.grid()","metadata":{"execution":{"iopub.status.busy":"2022-06-02T22:53:52.505448Z","iopub.execute_input":"2022-06-02T22:53:52.505763Z","iopub.status.idle":"2022-06-02T22:53:53.22285Z","shell.execute_reply.started":"2022-06-02T22:53:52.505725Z","shell.execute_reply":"2022-06-02T22:53:53.221999Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 4. Segment a few images!","metadata":{}},{"cell_type":"code","source":"sample_imgs = glob.glob(os.path.join(DATASET_FOLDER, \"test\", \"**\", \"*.png\"), recursive=True)\nif not sample_imgs:\n    sample_imgs = glob.glob(os.path.join(DATASET_FOLDER, \"train\", \"**\", \"*.png\"), recursive=True)\nprint(f\"images: {len(sample_imgs)}\")\nsample_imgs = sample_imgs[:5]\n\ndatamodule = SemanticSegmentationData.from_files(\n    predict_files=sample_imgs,\n    predict_transform=SemanticSegmentationInputTransform,\n    transform_kwargs=dict(image_size=IMAGE_SIZE),\n    batch_size=3,\n)","metadata":{"execution":{"iopub.status.busy":"2022-06-02T23:24:02.60429Z","iopub.execute_input":"2022-06-02T23:24:02.605154Z","iopub.status.idle":"2022-06-02T23:24:04.393971Z","shell.execute_reply.started":"2022-06-02T23:24:02.605109Z","shell.execute_reply":"2022-06-02T23:24:04.393143Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fig, axarr = plt.subplots(ncols=5, nrows=len(sample_imgs), figsize=(15, 3 * len(sample_imgs)))\nrunning_i = 0\nfor preds in trainer.predict(model, datamodule=datamodule):\n    for pred in preds:\n        # print(pred.keys())\n        img = np.rollaxis(pred['input'].cpu().numpy(), 0, 3)\n        print(img.dtype, img.min(), img.max())\n        axarr[running_i, 0].imshow(img)\n        for j, seg in enumerate(pred['preds'].cpu().numpy()):\n            p = axarr[running_i, j + 1].imshow(seg, vmin=-10, vmax=10)\n            plt.colorbar(p, ax=axarr[running_i, j + 1])\n        running_i += 1","metadata":{"execution":{"iopub.status.busy":"2022-06-02T23:25:45.789352Z","iopub.execute_input":"2022-06-02T23:25:45.790003Z","iopub.status.idle":"2022-06-02T23:25:50.964911Z","shell.execute_reply.started":"2022-06-02T23:25:45.78996Z","shell.execute_reply":"2022-06-02T23:25:50.963341Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# fig, axarr = plt.subplots(ncols=2, nrows=len(sample_imgs), figsize=(8, 4 * len(sample_imgs)))\n# running_i = 0\n# for preds in trainer.predict(model, datamodule=datamodule, output=\"labels\"):\n#     for pred in preds:\n#         # print(pred)\n#         img = plt.imread(sample_imgs[running_i])\n#         axarr[running_i, 0].imshow(img, cmap=\"gray\")\n#         axarr[running_i, 1].imshow(pred)\n#         running_i += 1","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-06-02T22:53:59.300382Z","iopub.execute_input":"2022-06-02T22:53:59.300913Z","iopub.status.idle":"2022-06-02T22:53:59.304473Z","shell.execute_reply.started":"2022-06-02T22:53:59.300872Z","shell.execute_reply":"2022-06-02T22:53:59.303667Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Inference 🔥","metadata":{}},{"cell_type":"code","source":"model = SemanticSegmentation.load_from_checkpoint(\n    \"semantic_segmentation_model.pt\"\n)","metadata":{"execution":{"iopub.status.busy":"2022-06-02T22:53:59.305963Z","iopub.execute_input":"2022-06-02T22:53:59.306384Z","iopub.status.idle":"2022-06-02T22:53:59.873019Z","shell.execute_reply.started":"2022-06-02T22:53:59.30635Z","shell.execute_reply":"2022-06-02T22:53:59.872173Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_pred = pd.read_csv(os.path.join(DATASET_FOLDER, \"sample_submission.csv\"))\ndisplay(df_pred.head())\nsfolder = \"test\" if WITH_SUBMISSION else \"train\"\n\nif not WITH_SUBMISSION:\n    df_pred = pd.read_csv(os.path.join(DATASET_FOLDER, \"train.csv\"))\n    df_pred = df_pred[df_pred[\"id\"].str.startswith(\"case123_day\")]\n\nos.makedirs(os.path.join(DATASET_IMAGES, sfolder), exist_ok=True)","metadata":{"execution":{"iopub.status.busy":"2022-06-02T22:53:59.874432Z","iopub.execute_input":"2022-06-02T22:53:59.874675Z","iopub.status.idle":"2022-06-02T22:54:00.457863Z","shell.execute_reply.started":"2022-06-02T22:53:59.874641Z","shell.execute_reply":"2022-06-02T22:54:00.457063Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from pprint import pprint\nfrom kaggle_imsegm.data_io import extract_tract_details\n\npprint(extract_tract_details(df_pred['id'].iloc[0], DATASET_FOLDER, folder=sfolder))\n\ndf_pred[['Case','Day','Slice', 'image', 'image_path', 'height', 'width']] = df_pred['id'].apply(\n    lambda x: pd.Series(extract_tract_details(x, DATASET_FOLDER, folder=sfolder))\n)\ndf_pred[\"Case_Day\"] = [f\"case{r['Case']}_day{r['Day']}\" for _, r in df_pred.iterrows()]\ndisplay(df_pred.head())","metadata":{"execution":{"iopub.status.busy":"2022-06-02T22:54:00.460418Z","iopub.execute_input":"2022-06-02T22:54:00.460729Z","iopub.status.idle":"2022-06-02T22:54:02.37107Z","shell.execute_reply.started":"2022-06-02T22:54:00.460689Z","shell.execute_reply":"2022-06-02T22:54:02.370298Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Predictions for test scans","metadata":{}},{"cell_type":"code","source":"from joblib import Parallel, delayed\nfrom kaggle_imsegm.data_io import preprocess_tract_scan\n\n_args = dict(\n    dir_data=os.path.join(DATASET_FOLDER, sfolder),\n    dir_imgs=DATASET_IMAGES,\n    dir_segm=None,\n    labels=LABELS,\n    sfolder=sfolder,\n)\ntest_scans = Parallel(n_jobs=6)(\n    delayed(preprocess_tract_scan)(dfg, **_args)\n    for _, dfg in df_pred.groupby(\"Case_Day\")\n)","metadata":{"execution":{"iopub.status.busy":"2022-06-02T22:54:02.372441Z","iopub.execute_input":"2022-06-02T22:54:02.373011Z","iopub.status.idle":"2022-06-02T22:54:20.082506Z","shell.execute_reply.started":"2022-06-02T22:54:02.372967Z","shell.execute_reply":"2022-06-02T22:54:20.081404Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nfrom itertools import chain\nfrom kaggle_imsegm.mask import rle_encode\n\npreds = []\nfor test_imgs in test_scans:\n    dm = SemanticSegmentationData.from_files(\n        predict_files=test_imgs,\n        predict_transform=SemanticSegmentationInputTransform,\n        transform_kwargs=dict(image_size=IMAGE_SIZE),\n        num_classes=len(LABELS) + 1,\n        batch_size=5,\n        num_workers=3,\n    )\n    pred = trainer.predict(model, datamodule=dm, output=\"labels\")\n    pred = list(chain(*pred))\n    for img, seg in zip(test_imgs, pred):\n        rle = rle_encode(np.array(seg)) if np.sum(seg) > 1 else {}\n        name, _ = os.path.splitext(os.path.basename(img))\n        id_ = \"_\".join(name.split(\"_\")[:4])\n        preds += [{\"id\": id_, \"class\": lb, \"predicted\": rle.get(i + 1, \"\")} for i, lb in enumerate(LABELS)]\n\ndf_pred = pd.DataFrame(preds)\ndisplay(df_pred[df_pred[\"predicted\"] != \"\"].head())","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-06-02T22:54:20.08776Z","iopub.execute_input":"2022-06-02T22:54:20.090369Z","iopub.status.idle":"2022-06-02T22:56:46.684242Z","shell.execute_reply.started":"2022-06-02T22:54:20.09032Z","shell.execute_reply":"2022-06-02T22:56:46.683337Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Finalize submissions","metadata":{}},{"cell_type":"code","source":"df_ssub = pd.read_csv(os.path.join(DATASET_FOLDER, \"sample_submission.csv\"))\ndel df_ssub['predicted']\ndf_pred = df_ssub.merge(df_pred, on=['id','class'])\n\ndf_pred[['id', 'class', 'predicted']].to_csv(\"submission.csv\", index=False)\n\n!head submission.csv","metadata":{"execution":{"iopub.status.busy":"2022-06-02T22:56:46.68586Z","iopub.execute_input":"2022-06-02T22:56:46.686139Z","iopub.status.idle":"2022-06-02T22:56:47.688652Z","shell.execute_reply.started":"2022-06-02T22:56:46.686102Z","shell.execute_reply":"2022-06-02T22:56:47.687637Z"},"trusted":true},"execution_count":null,"outputs":[]}]}