{"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":"# GI-Seg PyTorch ⚡ MONAI 3D Train & Infer\n- Downloaded weight and requirements come from [GI-Seg Downloads](https://www.kaggle.com/clemchris/gi-seg-download)\n- Dataset: [UW-Madison GI Tract Image Segmentation Masks](https://www.kaggle.com/datasets/clemchris/uw-madison-gi-tract-image-segmentation-masks)\n\n\n## Sources --> please upvote them if you find this notebook useful\n- Yiheng's [[LB 0.877] A 3D solution with MONAI](https://www.kaggle.com/competitions/uw-madison-gi-tract-image-segmentation/discussion/325646)\n- Awsaf's [UWMGI: Unet [Train] [PyTorch]](https://www.kaggle.com/code/awsaf49/uwmgi-unet-train-pytorch)\n- Awsaf's [UWMGI: 2.5D stride=2 Data](https://www.kaggle.com/code/awsaf49/uwmgi-2-5d-stride-2-data)\n- Awsaf's [UWMGI: Unet [Infer] [PyTorch]](https://www.kaggle.com/code/awsaf49/uwmgi-unet-infer-pytorch)\n\n## Scores\n- V01 - 10: Mixing IDs bug\n- V11: XXX","metadata":{}},{"cell_type":"markdown","source":"# Installs","metadata":{}},{"cell_type":"code","source":"!cd ../input/gi-seg-downloads && \\\npip install -q monai-0.8.1-202202162213-py3-none-any.whl torchmetrics-0.8.2-py3-none-any.whl","metadata":{"execution":{"iopub.status.busy":"2022-07-05T08:18:33.977307Z","iopub.execute_input":"2022-07-05T08:18:33.977926Z","iopub.status.idle":"2022-07-05T08:19:04.549858Z","shell.execute_reply.started":"2022-07-05T08:18:33.977837Z","shell.execute_reply":"2022-07-05T08:19:04.549016Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Imports","metadata":{}},{"cell_type":"code","source":"from pathlib import Path\nfrom typing import Callable\nfrom typing import List\nfrom typing import Optional\nfrom typing import Tuple\n\nimport cupy as cp\nimport cv2\nimport matplotlib.pyplot as plt\nimport monai\nimport numpy as np\nimport pandas as pd\nimport pytorch_lightning as pl\nimport seaborn as sns\nimport torch\nfrom joblib import delayed\nfrom joblib import Parallel\nfrom monai.data import CSVDataset\nfrom monai.data import DataLoader\nfrom pytorch_lightning.callbacks import ModelCheckpoint\nfrom torchmetrics import Metric\nfrom torchmetrics import MetricCollection\nfrom torch.utils.data import DataLoader\nfrom tqdm import tqdm\nfrom tqdm.auto import tqdm","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-07-05T08:19:04.552903Z","iopub.execute_input":"2022-07-05T08:19:04.553321Z","iopub.status.idle":"2022-07-05T08:19:11.506169Z","shell.execute_reply.started":"2022-07-05T08:19:04.553289Z","shell.execute_reply":"2022-07-05T08:19:11.505378Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Paths & Settings","metadata":{}},{"cell_type":"code","source":"KAGGLE_DIR = Path(\"/\") / \"kaggle\"\nINPUT_DIR = KAGGLE_DIR / \"input\"\nOUTPUT_DIR = KAGGLE_DIR / \"working\"\n\nINPUT_DATA_DIR = INPUT_DIR / \"uw-madison-gi-tract-image-segmentation\"\nINPUT_DATA_NPY_DIR = INPUT_DIR / \"uw-madison-gi-tract-image-segmentation-masks\"\n\nSPATIAL_SIZE = (224, 224, 80)\nN_SPLITS = 5\nRANDOM_SEED = 2022\nVAL_FOLD = 0\nBATCH_SIZE = 4\nNUM_WORKERS = 2\nLOSS = \"Dice\"\nLEARNING_RATE = 2e-4\nWEIGHT_DECAY = 1e-6\nSCHEDULER = \"CosineAnnealingWarmRestarts\"\nMIN_LR = 5e-4\n\nFAST_DEV_RUN = False # Debug training\nGPUS = 1\nMAX_EPOCHS = 100\nPRECISION = 16\n\nDEVICE = \"cuda\"\nTHR = 0.45\n\nDEBUG = False # Debug complete pipeline\n","metadata":{"execution":{"iopub.status.busy":"2022-07-05T08:19:11.509652Z","iopub.execute_input":"2022-07-05T08:19:11.509856Z","iopub.status.idle":"2022-07-05T08:19:11.519287Z","shell.execute_reply.started":"2022-07-05T08:19:11.509830Z","shell.execute_reply":"2022-07-05T08:19:11.518605Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Prepare 3D Data","metadata":{}},{"cell_type":"code","source":"def add_3d_paths(df, stage):\n    df[\"image_3d\"] = df[\"image_path\"].str.split(\"/scans\").str[0] + \"_image_3d.npy\"\n    df[\"image_3d\"] = df[\"image_3d\"].str.replace(\"input\", \"working\")\n    \n    if stage == \"train\":\n        df[\"mask_3d\"] = df[\"image_3d\"].str.replace(\"_image_\", \"_mask_\")\n        \n    return df","metadata":{"execution":{"iopub.status.busy":"2022-07-05T08:19:11.522744Z","iopub.execute_input":"2022-07-05T08:19:11.522992Z","iopub.status.idle":"2022-07-05T08:19:11.528901Z","shell.execute_reply.started":"2022-07-05T08:19:11.522941Z","shell.execute_reply":"2022-07-05T08:19:11.528260Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_df = pd.read_csv(INPUT_DATA_NPY_DIR / \"train_preprocessed.csv\")\n\nif DEBUG:\n    train_df = train_df.head(1_000)\n\ntrain_df = add_3d_paths(train_df, stage=\"train\")","metadata":{"execution":{"iopub.status.busy":"2022-07-05T08:19:11.530634Z","iopub.execute_input":"2022-07-05T08:19:11.531476Z","iopub.status.idle":"2022-07-05T08:19:12.485078Z","shell.execute_reply.started":"2022-07-05T08:19:11.531436Z","shell.execute_reply":"2022-07-05T08:19:12.484354Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def load_image(path):\n    image = cv2.imread(path, cv2.IMREAD_UNCHANGED) # uint16\n    return image\n\n\ndef load_mask(row):\n    shape = (row.height, row.width, 3)\n    mask = np.zeros(shape, dtype=np.uint8)\n\n    rles = eval(row.segmentation.replace(\"nan\", \"''\"))\n    for i, rle in enumerate(rles):\n        if rle:\n            mask[..., i] = rle_decode(rle, shape[:2])\n\n    return mask * 255\n\n\n# ref: https://www.kaggle.com/paulorzp/run-length-encode-and-decode\ndef rle_decode(mask_rle, shape):\n    \"\"\"\n    mask_rle: run-length as string formated (start length)\n    shape: (height,width) of array to return\n    Returns numpy array, 1 - mask, 0 - background\n\n    \"\"\"\n    s = np.asarray(mask_rle.split(), dtype=int)\n    starts = s[0::2] - 1\n    lengths = s[1::2]\n    ends = starts + lengths\n\n    mask = np.zeros(shape[0] * shape[1], dtype=np.uint8)\n    for lo, hi in zip(starts, ends):\n        mask[lo:hi] = 1\n\n    return mask.reshape(shape)  # Needed to align to RLE direction\n\n\ndef create_3d_image_mask(group_df, stage):\n    image_3d, mask_3d = [], []\n    for row in group_df.itertuples():\n        image_3d.append(load_image(row.image_path))  # uint16\n        \n        if stage == \"train\":\n            mask_3d.append(load_mask(row))  # uint8\n\n    image_3d = np.stack(image_3d, axis=-1)\n\n    dir_3d = Path(row.image_3d).parent\n    dir_3d.mkdir(parents=True, exist_ok=True)\n    np.save(row.image_3d, image_3d)\n\n    if stage == \"train\":\n        mask_3d = np.stack(mask_3d, axis=-1)\n        np.save(row.mask_3d, mask_3d)\n\n    return group_df.id.to_list()\n\n\ndef create_3d_npy_data(df, stage):\n    grouped = df.groupby([\"case\", \"day\"])\n    Parallel(n_jobs=NUM_WORKERS)(\n        delayed(create_3d_image_mask)(group_df, stage)\n        for _, group_df in tqdm(grouped, total=len(grouped), desc=\"Iterating over case-day groups\")\n    )\n    \n    # Aggregate ids\n    ids = df.groupby([\"case\", \"day\"])[\"id\"].agg(list).to_frame().reset_index()\n\n    columns_to_drop = [\"id\", \"slice\", \"image_path\"]\n    if stage == \"train\":\n        columns_to_drop += [\"classes\", \"segmentation\", \"rle_len\", \"empty\", \"mask_path\", \"image_paths\"]\n\n    df = df.drop(columns=columns_to_drop)\n    df = df.drop_duplicates().reset_index(drop=True)\n\n    df = df.merge(ids, on=[\"case\", \"day\"])\n\n    return df","metadata":{"execution":{"iopub.status.busy":"2022-07-05T08:19:12.487207Z","iopub.execute_input":"2022-07-05T08:19:12.487615Z","iopub.status.idle":"2022-07-05T08:19:12.504104Z","shell.execute_reply.started":"2022-07-05T08:19:12.487578Z","shell.execute_reply":"2022-07-05T08:19:12.503191Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_df = create_3d_npy_data(train_df, stage=\"train\")\n\nif DEBUG:\n    print(len(train_df))\n    display(train_df.head())","metadata":{"execution":{"iopub.status.busy":"2022-07-05T08:19:12.505244Z","iopub.execute_input":"2022-07-05T08:19:12.505601Z","iopub.status.idle":"2022-07-05T08:19:21.177531Z","shell.execute_reply.started":"2022-07-05T08:19:12.505552Z","shell.execute_reply":"2022-07-05T08:19:21.176669Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_df.to_csv(f\"train_preprocessed_3d.csv\")","metadata":{"execution":{"iopub.status.busy":"2022-07-05T08:19:21.179236Z","iopub.execute_input":"2022-07-05T08:19:21.179550Z","iopub.status.idle":"2022-07-05T08:19:21.187830Z","shell.execute_reply.started":"2022-07-05T08:19:21.179508Z","shell.execute_reply":"2022-07-05T08:19:21.187089Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# LitDataModule","metadata":{}},{"cell_type":"code","source":"class LitDataModule(pl.LightningDataModule):\n    def __init__(\n        self,\n        train_csv_path: str,\n        test_csv_path: Optional[str],\n        spatial_size: Tuple[int, int, int],\n        val_fold: int,\n        batch_size: int,\n        num_workers: int,\n    ):\n        super().__init__()\n\n        self.save_hyperparameters()\n\n        self.train_df = pd.read_csv(train_csv_path)\n\n        if test_csv_path is not None:\n            self.test_df = pd.read_csv(test_csv_path)\n        else:\n            self.test_df = None\n\n        self.train_transforms, self.val_transforms, self.test_transforms = self._init_transforms()\n\n    def _init_transforms(self):\n        spatial_size = self.hparams.spatial_size\n\n        train_transforms = [\n            monai.transforms.LoadImaged(keys=[\"image_3d\", \"mask_3d\"]),\n            monai.transforms.AddChanneld(keys=\"image_3d\"),\n            monai.transforms.AsChannelFirstd(keys=\"mask_3d\", channel_dim=2),\n            monai.transforms.RandSpatialCropd(keys=[\"image_3d\", \"mask_3d\"], roi_size=spatial_size, random_size=False),\n            monai.transforms.Lambdad(keys=[\"image_3d\", \"mask_3d\"], func=lambda x: x / x.max()),\n            monai.transforms.RandFlipd(keys=[\"image_3d\", \"mask_3d\"], prob=0.5, spatial_axis=[0]),\n            monai.transforms.RandFlipd(keys=[\"image_3d\", \"mask_3d\"], prob=0.5, spatial_axis=[1]),\n            monai.transforms.RandAffined(\n                keys=[\"image_3d\", \"mask_3d\"],\n                prob=0.5,\n                rotate_range=np.pi / 12,\n                translate_range=(spatial_size[0] * 0.0625, spatial_size[1] * 0.0625),\n                scale_range=(0.1, 0.1),\n                mode=\"nearest\",\n                padding_mode=\"reflection\",\n            ),\n            monai.transforms.OneOf(\n                [\n                    monai.transforms.RandGridDistortiond(\n                        keys=[\"image_3d\", \"mask_3d\"],\n                        prob=0.5,\n                        distort_limit=(-0.05, 0.05),\n                        mode=\"nearest\",\n                        padding_mode=\"reflection\",\n                    ),\n                    monai.transforms.RandCoarseDropoutd(\n                        keys=[\"image_3d\", \"mask_3d\"],\n                        holes=5,\n                        max_holes=8,\n                        spatial_size=(1, 1, 1),\n                        max_spatial_size=(12, 12, 12),\n                        fill_value=0.0,\n                        prob=0.5,\n                    ),\n                ]\n            ),\n            monai.transforms.RandScaleIntensityd(keys=\"image_3d\", factors=(-0.2, 0.2), prob=0.5),\n            monai.transforms.RandShiftIntensityd(keys=\"image_3d\", offsets=(-0.1, 0.1), prob=0.5),\n            monai.transforms.EnsureTyped(keys=[\"image_3d\", \"mask_3d\"], dtype=torch.float32),\n        ]\n\n        val_transforms = [\n            monai.transforms.LoadImaged(keys=[\"image_3d\", \"mask_3d\"]),\n            monai.transforms.AddChanneld(keys=\"image_3d\"),\n            monai.transforms.AsChannelFirstd(keys=\"mask_3d\", channel_dim=2),\n            monai.transforms.Lambdad(keys=[\"image_3d\", \"mask_3d\"], func=lambda x: x / x.max()),\n            monai.transforms.EnsureTyped(keys=[\"image_3d\", \"mask_3d\"], dtype=torch.float32),\n        ]\n\n        test_transforms = [\n            monai.transforms.LoadImaged(keys=\"image_3d\"),\n            monai.transforms.AddChanneld(keys=\"image_3d\"),\n            monai.transforms.Lambdad(keys=\"image_3d\", func=lambda x: x / x.max()),\n            monai.transforms.EnsureTyped(keys=\"image_3d\", dtype=torch.float32),\n        ]\n\n        train_transforms = monai.transforms.Compose(train_transforms)\n        val_transforms = monai.transforms.Compose(val_transforms)\n        test_transforms = monai.transforms.Compose(test_transforms)\n\n        return train_transforms, val_transforms, test_transforms\n\n    def setup(self, stage: Optional[str] = None):\n        train_df = self.train_df[self.train_df.fold != self.hparams.val_fold].reset_index(drop=True)\n        val_df = self.train_df[self.train_df.fold == self.hparams.val_fold].reset_index(drop=True)\n\n        if stage == \"fit\" or stage is None:\n            self.train_dataset = self._dataset(train_df, transforms=self.train_transforms)\n            self.val_dataset = self._dataset(val_df, transforms=self.val_transforms)\n\n        if stage == \"test\" or stage is None:\n            if self.test_df is not None:\n                self.test_dataset = self._dataset(self.test_df, transforms=self.test_transforms)\n            else:\n                self.test_dataset = self._dataset(val_df, transforms=self.val_transforms)\n\n    def _dataset(self, df: pd.DataFrame, transforms: Callable) -> CSVDataset:\n        return CSVDataset(src=df, transform=transforms)\n\n    def train_dataloader(self) -> DataLoader:\n        return self._dataloader(self.train_dataset, train=True)\n\n    def val_dataloader(self) -> DataLoader:\n        return self._dataloader(self.val_dataset)\n\n    def test_dataloader(self) -> DataLoader:\n        return self._dataloader(self.test_dataset)\n\n    def _dataloader(self, dataset: CSVDataset, train: bool = False) -> DataLoader:\n        return DataLoader(\n            dataset,\n            batch_size=self.hparams.batch_size if train else 1,\n            shuffle=train,\n            num_workers=self.hparams.num_workers,\n            drop_last=train,\n        )","metadata":{"execution":{"iopub.status.busy":"2022-07-05T08:19:21.189236Z","iopub.execute_input":"2022-07-05T08:19:21.189542Z","iopub.status.idle":"2022-07-05T08:19:21.218139Z","shell.execute_reply.started":"2022-07-05T08:19:21.189505Z","shell.execute_reply":"2022-07-05T08:19:21.217414Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Visualize Cases","metadata":{}},{"cell_type":"code","source":"data_module = LitDataModule(\n    train_csv_path=\"train_preprocessed_3d.csv\",\n    test_csv_path=None,\n    spatial_size=SPATIAL_SIZE,\n    val_fold=VAL_FOLD,\n    batch_size=4,\n    num_workers=NUM_WORKERS,\n)\ndata_module.setup()\n\ntrain_dataloader = data_module.train_dataloader()\nbatch = next(iter(train_dataloader))","metadata":{"execution":{"iopub.status.busy":"2022-07-05T08:19:21.221238Z","iopub.execute_input":"2022-07-05T08:19:21.222534Z","iopub.status.idle":"2022-07-05T08:19:26.447908Z","shell.execute_reply.started":"2022-07-05T08:19:21.221679Z","shell.execute_reply":"2022-07-05T08:19:26.447028Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for batch_idx, _ in enumerate(batch[\"image_3d\"]):\n    image_3d = batch[\"image_3d\"][batch_idx]\n    mask_3d = batch[\"mask_3d\"][batch_idx]\n\n    fig, ax = plt.subplots()\n    _, images_grid = monai.visualize.utils.matshow3d(volume=image_3d, every_n=10, frame_dim=-1, fig=fig)\n    _, masks_grid = monai.visualize.utils.matshow3d(volume=mask_3d, every_n=10, frame_dim=-1, channel_dim=0, fig=fig)\n    plt.title(f\"Case {batch['case'][batch_idx]}, Day {batch['day'][batch_idx]}\")\n    plt.imshow(images_grid, cmap=\"bone\")\n    plt.imshow(masks_grid, alpha=0.5)\n    plt.axis(\"off\")\n    plt.tight_layout()","metadata":{"execution":{"iopub.status.busy":"2022-07-05T08:19:26.449581Z","iopub.execute_input":"2022-07-05T08:19:26.449867Z","iopub.status.idle":"2022-07-05T08:19:30.555691Z","shell.execute_reply.started":"2022-07-05T08:19:26.449833Z","shell.execute_reply":"2022-07-05T08:19:30.554988Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Metrics","metadata":{}},{"cell_type":"code","source":"class DiceMetric(Metric):\n    def __init__(self):\n        super().__init__()\n\n        self.post_processing = monai.transforms.Compose(\n            [\n                monai.transforms.Activations(sigmoid=True),\n                monai.transforms.AsDiscrete(threshold=0.5),\n            ]\n        )\n        self.add_state(\"dice\", default=[])\n\n    def update(self, y_pred, y_true):\n        y_pred = self.post_processing(y_pred)\n        self.dice.append(monai.metrics.compute_meandice(y_pred, y_true))\n\n    def compute(self):\n        if len(self.dice) == 1:\n            return self.dice[0]\n        \n        return torch.mean(torch.stack(self.dice))","metadata":{"execution":{"iopub.status.busy":"2022-07-05T08:19:30.557015Z","iopub.execute_input":"2022-07-05T08:19:30.557513Z","iopub.status.idle":"2022-07-05T08:19:30.566282Z","shell.execute_reply.started":"2022-07-05T08:19:30.557471Z","shell.execute_reply":"2022-07-05T08:19:30.565710Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# LitModule","metadata":{}},{"cell_type":"code","source":"class LitModule(pl.LightningModule):\n    LOSSES = {\n        \"BCE\": torch.nn.BCEWithLogitsLoss(),\n        \"Dice\": monai.losses.DiceLoss(\n            sigmoid=True, smooth_nr=0.01, smooth_dr=0.01, include_background=True, batch=True, squared_pred=True\n        ),\n    }\n    \n    def __init__(\n        self,\n        loss: str,\n        learning_rate: float,\n        weight_decay: float,\n        scheduler: Optional[str],\n        T_max: int,\n        T_0: int,\n        min_lr: int,\n    ):\n        super().__init__()\n\n        self.save_hyperparameters()\n\n        self.model = self._init_model()\n\n        self.loss_fn = self._init_loss_fn()\n\n        self.metrics = self._init_metrics()\n\n    def _init_model(self):\n        return monai.networks.nets.UNet(\n            spatial_dims=3,\n            in_channels=1,\n            out_channels=3,\n            channels=(32, 64, 128, 256, 512),\n            strides=(2, 2, 2, 2),\n            kernel_size=3,\n            up_kernel_size=3,\n            num_res_units=2,\n            act=\"PRELU\",\n            norm=\"BATCH\",\n            dropout=0.2,\n            bias=True,\n            dimensions=None,\n        )\n\n    def _init_loss_fn(self):\n        losses = self.hparams.loss.split(\"_\")\n        loss_fns = [self.LOSSES[loss] for loss in losses]\n\n        def loss_fn(y_pred, y_true):\n            return sum(loss_fn(y_pred, y_true) for loss_fn in loss_fns) / len(loss_fns)\n\n        return loss_fn\n\n    def _init_metrics(self):\n        val_metrics = MetricCollection({\"val_dice\": DiceMetric()})\n        test_metrics = MetricCollection({\"test_dice\": DiceMetric()})\n        \n        return torch.nn.ModuleDict(\n            {\n                \"val_metrics\": val_metrics,\n                \"test_metrics\": test_metrics,\n            }\n        )\n\n    def configure_optimizers(self):\n        optimizer = torch.optim.Adam(params=self.parameters(), lr=self.hparams.learning_rate, weight_decay=self.hparams.weight_decay)\n\n        if self.hparams.scheduler is not None:\n            if self.hparams.scheduler == \"CosineAnnealingLR\":\n                scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n                    optimizer, T_max=self.hparams.T_max, eta_min=self.hparams.min_lr\n                )\n            elif self.hparams.scheduler == \"CosineAnnealingWarmRestarts\":\n                scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(\n                    optimizer, T_0=self.hparams.T_0, T_mult=1, eta_min=self.hparams.min_lr\n                )\n            else:\n                raise ValueError(f\"Unknown scheduler: {self.hparams.scheduler}\")\n\n            return {\"optimizer\": optimizer, \"lr_scheduler\": {\"scheduler\": scheduler, \"interval\": \"step\"}}\n        else:\n            return {\"optimizer\": optimizer}\n\n    def forward(self, images):\n        return self.model(images)\n\n    def training_step(self, batch, batch_idx):\n        images, masks = batch[\"image_3d\"], batch[\"mask_3d\"]\n        y_pred = self(images)\n\n        loss = self.loss_fn(y_pred, masks)\n\n        self.log(\"train_loss\", loss, on_epoch=True, batch_size=images.shape[0])\n\n        return loss\n\n    def validation_step(self, batch, batch_idx):\n        self.shared_step(batch, \"val\")\n\n    def test_step(self, batch, batch_idx):\n        self.shared_step(batch, \"test\")\n\n    def shared_step(self, batch, stage, log=True):\n        images, masks = batch[\"image_3d\"], batch[\"mask_3d\"]\n        y_pred = monai.inferers.sliding_window_inference(\n            inputs=images,\n            roi_size=(224, 224, 80),\n            sw_batch_size=4,\n            predictor=self,\n        )\n\n        loss = self.loss_fn(y_pred, masks)\n\n        metrics = self.metrics[f\"{stage}_metrics\"](y_pred, masks)\n\n        if log:\n            batch_size = images.shape[0]\n            self._log(loss, batch_size, metrics, stage)\n\n        return loss\n\n    def _log(self, loss, batch_size, metrics, stage):\n        self.log(f\"{stage}_loss\", loss, on_step=False, on_epoch=True, batch_size=batch_size)\n        self.log_dict(metrics, on_step=False, on_epoch=True, batch_size=batch_size)\n\n    @classmethod\n    def load_eval_checkpoint(cls, checkpoint_path, device):\n        module = cls.load_from_checkpoint(checkpoint_path=checkpoint_path).to(device)\n        module.eval()\n\n        return module","metadata":{"execution":{"iopub.status.busy":"2022-07-05T08:19:30.567639Z","iopub.execute_input":"2022-07-05T08:19:30.568127Z","iopub.status.idle":"2022-07-05T08:19:30.596940Z","shell.execute_reply.started":"2022-07-05T08:19:30.568090Z","shell.execute_reply":"2022-07-05T08:19:30.596142Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train","metadata":{}},{"cell_type":"code","source":"def train(\n    random_seed: int = RANDOM_SEED,\n    train_csv_path: str = \"train_preprocessed_3d.csv\",\n    spatial_size: Tuple[int, int, int] = SPATIAL_SIZE,\n    val_fold: str = VAL_FOLD,\n    batch_size: int = BATCH_SIZE,\n    num_workers: int = NUM_WORKERS,\n    loss: str = LOSS,\n    learning_rate: float = LEARNING_RATE,\n    weight_decay: float = WEIGHT_DECAY,\n    scheduler: Optional[str] = SCHEDULER,\n    min_lr: float = MIN_LR,\n    gpus: int = GPUS,\n    fast_dev_run: bool = FAST_DEV_RUN,\n    max_epochs: int = MAX_EPOCHS,\n    precision: int = PRECISION,\n    debug: bool = DEBUG,\n):\n    pl.seed_everything(random_seed)\n\n    if debug:\n        max_epochs = 2\n\n    data_module = LitDataModule(\n        train_csv_path=train_csv_path,\n        test_csv_path=None,\n        spatial_size=spatial_size,\n        val_fold=val_fold,\n        batch_size=batch_size,\n        num_workers=num_workers,\n    )\n    data_module.setup()\n\n    module = LitModule(\n        loss=loss,\n        learning_rate=learning_rate,\n        weight_decay=weight_decay,\n        scheduler=scheduler,\n        T_max=int(30_000 / batch_size * max_epochs) + 50,\n        T_0=100 * (len(data_module.train_dataset) // batch_size),\n        min_lr=min_lr,\n    )\n\n    trainer = pl.Trainer(\n        fast_dev_run=fast_dev_run,\n        gpus=gpus,\n        log_every_n_steps=1,\n        logger=pl.loggers.CSVLogger(save_dir='logs/'),\n        max_epochs=max_epochs,\n        precision=precision,\n    )\n\n    trainer.fit(module, datamodule=data_module)\n\n    if not fast_dev_run:\n        trainer.test(module, datamodule=data_module)\n        \n    return trainer","metadata":{"execution":{"iopub.status.busy":"2022-07-05T08:19:30.598269Z","iopub.execute_input":"2022-07-05T08:19:30.598611Z","iopub.status.idle":"2022-07-05T08:19:30.612729Z","shell.execute_reply.started":"2022-07-05T08:19:30.598574Z","shell.execute_reply":"2022-07-05T08:19:30.611859Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"trainer = train()","metadata":{"execution":{"iopub.status.busy":"2022-07-05T08:19:30.614134Z","iopub.execute_input":"2022-07-05T08:19:30.614460Z","iopub.status.idle":"2022-07-05T08:20:01.702999Z","shell.execute_reply.started":"2022-07-05T08:19:30.614426Z","shell.execute_reply":"2022-07-05T08:20:01.702040Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# From https://www.kaggle.com/code/jirkaborovec?scriptVersionId=93358967&cellId=22\nmetrics = pd.read_csv(f\"{trainer.logger.log_dir}/metrics.csv\")[[\"epoch\", \"train_loss_epoch\", \"val_loss\"]]\nmetrics.set_index(\"epoch\", inplace=True)\n\nsns.relplot(data=metrics, kind=\"line\", height=5, aspect=1.5)\nplt.grid()","metadata":{"execution":{"iopub.status.busy":"2022-07-05T08:20:01.704657Z","iopub.execute_input":"2022-07-05T08:20:01.705230Z","iopub.status.idle":"2022-07-05T08:20:02.064244Z","shell.execute_reply.started":"2022-07-05T08:20:01.705189Z","shell.execute_reply":"2022-07-05T08:20:02.063524Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Infer","metadata":{}},{"cell_type":"markdown","source":"### Load Test Data","metadata":{}},{"cell_type":"code","source":"def extract_metadata_from_id(df):\n    df[[\"case\", \"day\", \"slice\"]] = df[\"id\"].str.split(\"_\", n=2, expand=True)\n\n    df[\"case\"] = df[\"case\"].str.replace(\"case\", \"\").astype(int)\n    df[\"day\"] = df[\"day\"].str.replace(\"day\", \"\").astype(int)\n    df[\"slice\"] = df[\"slice\"].str.replace(\"slice_\", \"\").astype(int)\n\n    return df\n\n\ndef extract_metadata_from_path(path_df):\n    path_df[[\"parent\", \"case_day\", \"scans\", \"file_name\"]] = path_df[\"image_path\"].str.rsplit(\"/\", n=3, expand=True)\n\n    path_df[[\"case\", \"day\"]] = path_df[\"case_day\"].str.split(\"_\", expand=True)\n    path_df[\"case\"] = path_df[\"case\"].str.replace(\"case\", \"\")\n    path_df[\"day\"] = path_df[\"day\"].str.replace(\"day\", \"\")\n\n    path_df[[\"slice\", \"width\", \"height\", \"spacing\", \"spacing_\"]] = (\n        path_df[\"file_name\"].str.replace(\"slice_\", \"\").str.replace(\".png\", \"\").str.split(\"_\", expand=True)\n    )\n    path_df = path_df.drop(columns=[\"parent\", \"case_day\", \"scans\", \"file_name\", \"spacing_\"])\n\n    numeric_cols = [\"case\", \"day\", \"slice\", \"width\", \"height\", \"spacing\"]\n    path_df[numeric_cols] = path_df[numeric_cols].apply(pd.to_numeric)\n\n    return path_df","metadata":{"execution":{"iopub.status.busy":"2022-07-05T08:20:02.065496Z","iopub.execute_input":"2022-07-05T08:20:02.066542Z","iopub.status.idle":"2022-07-05T08:20:02.077386Z","shell.execute_reply.started":"2022-07-05T08:20:02.066496Z","shell.execute_reply":"2022-07-05T08:20:02.076518Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sub_df = pd.read_csv(INPUT_DATA_DIR / \"sample_submission.csv\")\ntest_set_hidden = not bool(len(sub_df))\n\nif test_set_hidden:\n    test_df = pd.read_csv(INPUT_DATA_DIR / \"train.csv\")[: 1000 * 3]\n    test_df = test_df.drop(columns=[\"class\", \"segmentation\"]).drop_duplicates()\n    image_paths = [str(path) for path in (INPUT_DATA_DIR / \"train\").rglob(\"*.png\")]\nelse:\n    test_df = sub_df.drop(columns=[\"class\", \"predicted\"]).drop_duplicates()\n    image_paths = [str(path) for path in (INPUT_DATA_DIR / \"test\").rglob(\"*.png\")]\n\ntest_df = extract_metadata_from_id(test_df)\n\npath_df = pd.DataFrame(image_paths, columns=[\"image_path\"])\npath_df = extract_metadata_from_path(path_df)\n\ntest_df = test_df.merge(path_df, on=[\"case\", \"day\", \"slice\"], how=\"left\")\ntest_df = add_3d_paths(test_df, stage=\"test\")\n\nprint(len(test_df))\ntest_df.head()","metadata":{"execution":{"iopub.status.busy":"2022-07-05T08:20:02.078601Z","iopub.execute_input":"2022-07-05T08:20:02.078936Z","iopub.status.idle":"2022-07-05T08:20:07.623441Z","shell.execute_reply.started":"2022-07-05T08:20:02.078897Z","shell.execute_reply":"2022-07-05T08:20:07.622709Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_df = create_3d_npy_data(test_df, stage=\"test\")","metadata":{"execution":{"iopub.status.busy":"2022-07-05T08:20:07.624709Z","iopub.execute_input":"2022-07-05T08:20:07.625227Z","iopub.status.idle":"2022-07-05T08:20:09.879933Z","shell.execute_reply.started":"2022-07-05T08:20:07.625187Z","shell.execute_reply":"2022-07-05T08:20:09.879184Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Save Test DataFrame","metadata":{}},{"cell_type":"code","source":"test_df.to_csv(\"test_preprocessed_3d.csv\", index=False)","metadata":{"execution":{"iopub.status.busy":"2022-07-05T08:20:09.881126Z","iopub.execute_input":"2022-07-05T08:20:09.881381Z","iopub.status.idle":"2022-07-05T08:20:09.888848Z","shell.execute_reply.started":"2022-07-05T08:20:09.881334Z","shell.execute_reply":"2022-07-05T08:20:09.888158Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Run inference","metadata":{}},{"cell_type":"code","source":"def mask2rle(mask):\n    \"\"\"\n    img: numpy array, 1 - mask, 0 - background\n    Returns run length as string formated\n    \"\"\"\n    mask = cp.array(mask)\n    pixels = mask.flatten()\n    pad = cp.array([0])\n    pixels = cp.concatenate([pad, pixels, pad])\n    runs = cp.where(pixels[1:] != pixels[:-1])[0] + 1\n    runs[1::2] -= runs[::2]\n\n    return \" \".join(str(x) for x in runs)\n\n\ndef masks2rles(masks, ids):\n    pred_strings = []\n    pred_ids = []\n    pred_classes = []\n\n    for idx in tqdm(range(masks.shape[0])):\n        mask = masks[idx]\n        \n        rle = [None] * 3\n        for midx in [0, 1, 2]:\n            rle[midx] = mask2rle(mask[..., midx])\n\n        pred_strings.extend(rle)\n        pred_ids.extend([ids[idx]] * len(rle))\n        pred_classes.extend([\"large_bowel\", \"small_bowel\", \"stomach\"])\n\n    return pred_strings, pred_ids, pred_classes\n\n\n@torch.no_grad()\ndef infer(model_paths, spatial_size, device, thr):\n    data_module = LitDataModule(\n        train_csv_path=\"train_preprocessed_3d.csv\",\n        test_csv_path=\"test_preprocessed_3d.csv\",\n        spatial_size=spatial_size,\n        val_fold=0,\n        batch_size=1,\n        num_workers=0,\n    )\n    \n    data_module.setup(stage=\"test\")\n    test_dataloader = data_module.test_dataloader()\n    \n    pred_strings = []\n    pred_ids = []\n    pred_classes = []\n\n    for batch in tqdm(test_dataloader):\n        images_3d = batch[\"image_3d\"].to(device, dtype=torch.float)\n        ids, height, width = batch[\"id\"][0], batch[\"height\"][0], batch[\"width\"][0]\n        \n        ids = eval(ids)\n        height = int(height)\n        width = int(width)\n\n        size_3d = images_3d.size()\n        masks_3d = torch.zeros((size_3d[0], 3, size_3d[2], size_3d[3], size_3d[4]), device=device, dtype=torch.float32)\n\n        for path in model_paths:\n            model = LitModule.load_eval_checkpoint(path, device=device)\n            out_3d = monai.inferers.sliding_window_inference(\n                inputs=images_3d,\n                roi_size=(224, 224, 80),\n                sw_batch_size=4,\n                predictor=model,\n            )\n            #out_3d = model(images_3d)\n            out_3d = torch.nn.Sigmoid()(out_3d)\n            masks_3d += out_3d / len(model_paths)\n\n        # Remove batch dim\n        images_3d = images_3d[0]\n        masks_3d = masks_3d[0]\n        \n        # Resize to original shape\n        spatial_size = (height, width, len(ids))\n        resize_transform = monai.transforms.Resize(spatial_size=spatial_size, mode=\"nearest\")\n        images_3d = resize_transform(images_3d)\n        masks_3d = resize_transform(masks_3d)\n\n        # Use depth as batch dim\n        images_3d = images_3d.permute((3, 0, 1, 2))\n        masks_3d = masks_3d.permute((3, 0, 1, 2))\n            \n        masks = (masks_3d.permute((0, 2, 3, 1)) > thr).to(torch.uint8).cpu().detach().numpy()  # shape: (n, h, w, c)\n        result = masks2rles(masks, ids)\n        pred_strings.extend(result[0])\n        pred_ids.extend(result[1])\n        pred_classes.extend(result[2])\n        \n    pred_df = pd.DataFrame({\"id\": pred_ids, \"class\": pred_classes, \"predicted\": pred_strings})\n\n    return pred_df, images_3d, masks_3d, ids","metadata":{"execution":{"iopub.status.busy":"2022-07-05T08:22:25.244327Z","iopub.execute_input":"2022-07-05T08:22:25.244730Z","iopub.status.idle":"2022-07-05T08:22:25.266287Z","shell.execute_reply.started":"2022-07-05T08:22:25.244692Z","shell.execute_reply":"2022-07-05T08:22:25.265430Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model_paths = list((Path(trainer.logger.log_dir) / \"checkpoints\").glob(\"*.ckpt\"))\nmodel_paths","metadata":{"execution":{"iopub.status.busy":"2022-07-05T08:22:25.876307Z","iopub.execute_input":"2022-07-05T08:22:25.877017Z","iopub.status.idle":"2022-07-05T08:22:25.883258Z","shell.execute_reply.started":"2022-07-05T08:22:25.876980Z","shell.execute_reply":"2022-07-05T08:22:25.882333Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pred_df, images_3d, masks_3d, ids = infer(model_paths, SPATIAL_SIZE, DEVICE, THR)","metadata":{"execution":{"iopub.status.busy":"2022-07-05T08:22:26.382853Z","iopub.execute_input":"2022-07-05T08:22:26.383329Z","iopub.status.idle":"2022-07-05T08:22:42.012757Z","shell.execute_reply.started":"2022-07-05T08:22:26.383290Z","shell.execute_reply":"2022-07-05T08:22:42.011913Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ncols = 7\nnrows = 7\nfig, _ = plt.subplots(figsize=(5 * ncols, 5 * nrows))\nfor i, (image, mask, id) in enumerate(zip(images_3d[6 * 12 : 10 * 12 + 1], masks_3d[6 * 12 : 10 * 12 + 1], ids[6 * 12 : 10 * 12 + 1])):\n    image = image.permute((1, 2, 0)).detach().cpu().numpy()\n    mask = mask.permute((1, 2, 0)).detach().cpu().numpy()\n\n    plt.subplot(nrows, ncols, i + 1)\n    plt.tight_layout()\n    plt.title(id)\n    \n    plt.imshow(image, cmap=\"bone\")\n    plt.imshow(mask, alpha=0.5)\n    plt.axis(\"off\")","metadata":{"execution":{"iopub.status.busy":"2022-07-05T08:22:42.014518Z","iopub.execute_input":"2022-07-05T08:22:42.014859Z","iopub.status.idle":"2022-07-05T08:22:54.193684Z","shell.execute_reply.started":"2022-07-05T08:22:42.014820Z","shell.execute_reply":"2022-07-05T08:22:54.191455Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Submit","metadata":{}},{"cell_type":"code","source":"if not test_set_hidden:\n    sub_df = pd.read_csv(\"../input/uw-madison-gi-tract-image-segmentation/sample_submission.csv\")\n    del sub_df[\"predicted\"]\nelse:\n    sub_df = pd.read_csv(\"../input/uw-madison-gi-tract-image-segmentation/train.csv\")[: 1000 * 3]\n    del sub_df[\"segmentation\"]\n\nsub_df = sub_df.merge(pred_df, on=[\"id\", \"class\"])\nsub_df.to_csv(\"submission.csv\", index=False)\ndisplay(sub_df.head(5))","metadata":{"execution":{"iopub.status.busy":"2022-07-05T08:22:54.195133Z","iopub.execute_input":"2022-07-05T08:22:54.195625Z","iopub.status.idle":"2022-07-05T08:22:54.562681Z","shell.execute_reply.started":"2022-07-05T08:22:54.195581Z","shell.execute_reply":"2022-07-05T08:22:54.561826Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## ","metadata":{}}]}