{"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":"I've trained model for tumor 3D segmentation based on data from [Task1](https://www.kaggle.com/dschettler8845/brats-2021-task1). [DenseVNet](https://www.ncbi.nlm.nih.gov/pmc/articles/PMC6076994/) model was used for segmentation. You can somehow use model and code provided for\nMGMT classification.","metadata":{}},{"cell_type":"markdown","source":"## Model","metadata":{}},{"cell_type":"code","source":"\"\"\"Module with DenseVNet\"\"\"\n\nimport numpy as np\nimport torch\n\n\nclass DenseVNet(torch.nn.Module):\n    def __init__(self, in_channels: int = 1, out_channels: int = 1):\n        super().__init__()\n\n        kernel_size = [5, 3, 3]\n        num_downsample_channels = [24, 24, 24]\n        num_skip_channels = [12, 24, 24]\n        units = [5, 10, 10]\n        growth_rate = [4, 8, 16]\n\n        self.dfs_blocks = torch.nn.ModuleList()\n        for i in range(3):\n            self.dfs_blocks.append(\n                DownsampleWithDfs(\n                    in_channels=in_channels,\n                    downsample_channels=num_downsample_channels[i],\n                    skip_channels=num_skip_channels[i],\n                    kernel_size=kernel_size[i],\n                    units=units[i],\n                    growth_rate=growth_rate[i],\n                )\n            )\n            in_channels = num_downsample_channels[i] + units[i] * growth_rate[i]\n\n        self.upsample_1 = torch.nn.Upsample(scale_factor=2, mode='trilinear')\n        self.upsample_2 = torch.nn.Upsample(scale_factor=4, mode='trilinear')\n\n        self.out_conv = ConvBlock(\n            in_channels=sum(num_skip_channels),\n            out_channels=out_channels,\n            kernel_size=3,\n            batch_norm=True,\n            preactivation=True,\n        )\n        self.upsample_out = torch.nn.Upsample(scale_factor=2, mode='trilinear')\n\n    def forward(self, x):\n        x, skip_1 = self.dfs_blocks[0](x)\n        x, skip_2 = self.dfs_blocks[1](x)\n        _, skip_3 = self.dfs_blocks[2](x)\n\n        skip_2 = self.upsample_1(skip_2)\n        skip_3 = self.upsample_2(skip_3)\n\n        out = self.out_conv(torch.cat([skip_1, skip_2, skip_3], 1))\n        out = self.upsample_out(out)\n\n        return out\n\n\nclass ConvBlock(torch.nn.Module):\n    def __init__(\n        self,\n        in_channels,\n        out_channels,\n        kernel_size,\n        dilation=1,\n        stride=1,\n        batch_norm=True,\n        preactivation=False,\n    ):\n        super().__init__()\n\n        if dilation != 1:\n            raise NotImplementedError()\n\n        padding = kernel_size - stride\n        if padding % 2 != 0:\n            pad = torch.nn.ConstantPad3d(\n                tuple([padding % 2, padding - padding % 2] * 3), 0\n            )\n        else:\n            pad = torch.nn.ConstantPad3d(padding // 2, 0)\n\n        if preactivation:\n            layers = [\n                torch.nn.ReLU(),\n                pad,\n                torch.nn.Conv3d(\n                    in_channels=in_channels,\n                    out_channels=out_channels,\n                    kernel_size=kernel_size,\n                    stride=stride,\n                ),\n            ]\n            if batch_norm:\n                layers = [torch.nn.BatchNorm3d(in_channels)] + layers\n        else:\n            layers = [\n                pad,\n                torch.nn.Conv3d(\n                    in_channels=in_channels,\n                    out_channels=out_channels,\n                    kernel_size=kernel_size,\n                    stride=stride,\n                ),\n            ]\n            if batch_norm:\n                layers.append(torch.nn.BatchNorm3d(out_channels))\n            layers.append(torch.nn.ReLU())\n\n        self.conv = torch.nn.Sequential(*layers)\n\n    def forward(self, x):\n        return self.conv(x)\n\n\nclass DenseFeatureStack(torch.nn.Module):\n    def __init__(\n        self,\n        in_channels,\n        units,\n        growth_rate,\n        kernel_size,\n        dilation=1,\n        batch_norm=True,\n        batchwise_spatial_dropout=False,\n    ):\n        super().__init__()\n\n        self.units = torch.nn.ModuleList()\n        for _ in range(units):\n            if batchwise_spatial_dropout:\n                raise NotImplementedError\n\n            self.units.append(\n                ConvBlock(\n                    in_channels=in_channels,\n                    out_channels=growth_rate,\n                    kernel_size=kernel_size,\n                    dilation=dilation,\n                    stride=1,\n                    batch_norm=batch_norm,\n                    preactivation=True,\n                )\n            )\n            in_channels += growth_rate\n\n    def forward(self, x):\n        feature_stack = [x]\n\n        for unit in self.units:\n            inputs = torch.cat(feature_stack, 1)\n            out = unit(inputs)\n            feature_stack.append(out)\n\n        return torch.cat(feature_stack, 1)\n\n\nclass DownsampleWithDfs(torch.nn.Module):\n    def __init__(\n        self,\n        in_channels,\n        downsample_channels,\n        skip_channels,\n        kernel_size,\n        units,\n        growth_rate,\n    ):\n        super().__init__()\n\n        self.downsample = ConvBlock(\n            in_channels=in_channels,\n            out_channels=downsample_channels,\n            kernel_size=kernel_size,\n            stride=2,\n            batch_norm=True,\n            preactivation=True,\n        )\n        self.dfs = DenseFeatureStack(\n            downsample_channels, units, growth_rate, 3, batch_norm=True\n        )\n        self.skip = ConvBlock(\n            in_channels=downsample_channels + units * growth_rate,\n            out_channels=skip_channels,\n            kernel_size=3,\n            batch_norm=True,\n            preactivation=True,\n        )\n\n    def forward(self, x):\n        x = self.downsample(x)\n        x = self.dfs(x)\n        x_skip = self.skip(x)\n\n        return x, x_skip\n","metadata":{"execution":{"iopub.status.busy":"2021-09-21T15:24:06.724186Z","iopub.execute_input":"2021-09-21T15:24:06.724616Z","iopub.status.idle":"2021-09-21T15:24:06.755602Z","shell.execute_reply.started":"2021-09-21T15:24:06.724580Z","shell.execute_reply":"2021-09-21T15:24:06.754630Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model Evaluator","metadata":{}},{"cell_type":"code","source":"\"\"\"Module with model evaluator\"\"\"\n\nfrom typing import Tuple\n\nimport numpy as np\nimport scipy.ndimage\nimport torch\n\n\nclass ModelEvaluator:\n    def __init__(\n        self,\n        model: DenseVNet,\n        min_value: int = -200,\n        max_value: int = 200,\n        image_size: Tuple[int, int, int] = (184, 184, 128),\n        device: torch.device = torch.device('cpu'),\n    ):\n        self.model = model\n        self.model.eval()\n        self.model.to(device)\n\n        self.min_value = min_value\n        self.max_value = max_value\n        self.image_size = image_size\n\n        self.device = device\n\n    def evaluate(self, image: np.ndarray) -> np.ndarray:\n        image = self._resize_ct(\n            ct=image, \n            ct_size=self.image_size\n        )\n        image_tensor = self.preprocess_image(\n            image=image,\n            min_value=self.min_value,\n            max_value=self.max_value,\n        )\n        image_tensor = image_tensor.to(self.device)\n\n        image_tensor = image_tensor.unsqueeze(0).float()\n        \n        mask = self.model(x=image_tensor)\n\n        mask = torch.sigmoid(mask)\n        mask = mask.cpu().detach().numpy()\n        mask = mask[0]\n\n        return mask\n\n    @staticmethod\n    def preprocess_image(\n        image: np.ndarray,\n        min_value: int = -200,\n        max_value: int = 200,\n    ) -> torch.Tensor:\n        image = np.clip(image, min_value, max_value)\n        image = (image - min_value) / (max_value - min_value)\n\n        image_tensor = torch.tensor(image)\n        image_tensor = image_tensor.unsqueeze(0)\n\n        return image_tensor\n    \n    @staticmethod\n    def _resize_ct(\n        ct: np.ndarray, ct_size: Tuple[int, int, int] = (512, 512, 256)\n    ) -> np.ndarray:\n        zoom_factor = [\n            first_value / second_value\n            for first_value, second_value in zip(ct_size, ct.shape)\n        ]\n        res_ct = scipy.ndimage.zoom(ct, zoom=zoom_factor)\n\n        assert res_ct.shape == ct_size, f'Bad result size: {res_ct.shape}!={ct_size}'\n\n        return res_ct\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2021-09-21T15:24:07.509911Z","iopub.execute_input":"2021-09-21T15:24:07.510349Z","iopub.status.idle":"2021-09-21T15:24:07.524126Z","shell.execute_reply.started":"2021-09-21T15:24:07.510317Z","shell.execute_reply":"2021-09-21T15:24:07.523362Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Dicom loading functions","metadata":{}},{"cell_type":"code","source":"from pathlib import Path\nfrom typing import Dict, List, Tuple, Union\n\nimport numpy as np\nimport nibabel as nib\nimport SimpleITK as sitk\n\n\ndef _load_and_align_ct(dicom_folder_path: str) -> np.ndarray:\n    image_and_meta = _load_dicom(dicom_folder_path=dicom_folder_path)\n    \n    image = sitk.GetArrayFromImage(image=image_and_meta)\n    affine = _make_affine(image_and_meta=image_and_meta)\n    \n    ct_all_info = nib.Nifti1Image(image, affine)\n    image = _align_ct(ct_all_info=ct_all_info)\n\n    return image.get_fdata(dtype=np.float64)\n\ndef _load_dicom(dicom_folder_path: Union[str, Path]) -> sitk.Image:\n    sitk.ProcessObject_SetGlobalWarningDisplay(False)\n\n    series_ids = sitk.ImageSeriesReader.GetGDCMSeriesIDs(\n        directory=str(dicom_folder_path)\n    )\n    series_file_names = sitk.ImageSeriesReader.GetGDCMSeriesFileNames(\n        str(dicom_folder_path), series_ids[0]\n    )\n    series_reader = sitk.ImageSeriesReader()\n    series_reader.SetFileNames(series_file_names)\n    series_reader.LoadPrivateTagsOn()\n    image_and_meta = series_reader.Execute()\n    \n    return image_and_meta\n    \n\ndef _align_ct(ct_all_info: nib.Nifti1Image) -> nib.Nifti1Image:\n    orig_ornt = nib.io_orientation(ct_all_info.affine)\n    targ_ornt = nib.orientations.axcodes2ornt(axcodes='LPS')\n    transform = nib.orientations.ornt_transform(\n        start_ornt=orig_ornt, end_ornt=targ_ornt\n    )\n\n    img_ornt = ct_all_info.as_reoriented(ornt=transform)\n\n    return img_ornt\ndef _make_affine(image_and_meta: sitk.Image):\n    # get affine transform in LPS\n    c = [image_and_meta.TransformContinuousIndexToPhysicalPoint(p)\n         for p in ((1, 0, 0),\n                   (0, 1, 0),\n                   (0, 0, 1),\n                   (0, 0, 0))]\n    c = np.array(c)\n    affine = np.concatenate([\n        np.concatenate([c[0:3] - c[3:], c[3:]], axis=0),\n        [[0.], [0.], [0.], [1.]]\n    ], axis=1)\n    affine = np.transpose(affine)\n    # convert to RAS to match nibabel\n    affine = np.matmul(np.diag([-1., -1., 1., 1.]), affine)\n    return affine\n","metadata":{"execution":{"iopub.status.busy":"2021-09-21T15:24:08.165558Z","iopub.execute_input":"2021-09-21T15:24:08.165935Z","iopub.status.idle":"2021-09-21T15:24:08.178939Z","shell.execute_reply.started":"2021-09-21T15:24:08.165900Z","shell.execute_reply":"2021-09-21T15:24:08.178156Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model and image loading","metadata":{}},{"cell_type":"code","source":"def rename_keys(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:\n    new_state_dict = {}\n\n    for layer_name, layer_weights in state_dict.items():\n        layer_name = layer_name.replace('model.', '')\n\n        new_state_dict[layer_name] = layer_weights\n\n    return new_state_dict\n","metadata":{"execution":{"iopub.status.busy":"2021-09-21T15:24:08.892866Z","iopub.execute_input":"2021-09-21T15:24:08.893628Z","iopub.status.idle":"2021-09-21T15:24:08.898859Z","shell.execute_reply.started":"2021-09-21T15:24:08.893591Z","shell.execute_reply":"2021-09-21T15:24:08.898094Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"flair_dicom_path = '../input/rsna-miccai-brain-tumor-radiogenomic-classification/test/00013/FLAIR'\nmodel_path = '../input/densevnettumorsegmentaion/DenseVNet_27epoch_best.ckpt'\n\nin_channels = 1\nout_channels = 2\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')","metadata":{"execution":{"iopub.status.busy":"2021-09-21T15:25:05.912142Z","iopub.execute_input":"2021-09-21T15:25:05.912852Z","iopub.status.idle":"2021-09-21T15:25:05.917318Z","shell.execute_reply.started":"2021-09-21T15:25:05.912815Z","shell.execute_reply":"2021-09-21T15:25:05.916429Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"image = _load_and_align_ct(dicom_folder_path=flair_dicom_path)\nmodel_meta = torch.load(model_path, map_location='cpu')\n\nmodel_state_dict = model_meta['state_dict']\nmodel_state_dict = rename_keys(state_dict=model_state_dict)\n\nmodel = DenseVNet(in_channels=in_channels, out_channels=out_channels)\nmodel.load_state_dict(model_state_dict)\nmodel.to(device=device)\nmodel.eval()\n\nmodel_evaluator = ModelEvaluator(\n    model=model,\n    device=device,\n    min_value=-200,\n    max_value=2500,\n)","metadata":{"execution":{"iopub.status.busy":"2021-09-21T15:25:06.837538Z","iopub.execute_input":"2021-09-21T15:25:06.838161Z","iopub.status.idle":"2021-09-21T15:25:08.657023Z","shell.execute_reply.started":"2021-09-21T15:25:06.838121Z","shell.execute_reply":"2021-09-21T15:25:08.656217Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Mask calculation and visualization","metadata":{}},{"cell_type":"code","source":"mask = model_evaluator.evaluate(image=image)","metadata":{"execution":{"iopub.status.busy":"2021-09-21T15:25:08.658774Z","iopub.execute_input":"2021-09-21T15:25:08.659114Z","iopub.status.idle":"2021-09-21T15:25:12.630055Z","shell.execute_reply.started":"2021-09-21T15:25:08.659078Z","shell.execute_reply":"2021-09-21T15:25:12.629311Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import matplotlib.pyplot as plt\nfrom cv2 import cv2\nfrom matplotlib import animation, rc\nrc('animation', html='jshtml')\n\n\ndef create_animation(\n    image,\n    mask,\n    min_value: int = -200,\n    max_value: int = 200,\n    edge: float = 0.5,\n    image_size: Tuple[int, int] = (512, 512),\n):\n    curr_slice = combine_image_and_mask_slice(\n        image=image,\n        mask=mask, \n        slice_idx=0,\n        min_value=min_value,\n        max_value=max_value,\n        edge=edge,\n        image_size=image_size\n    )\n    \n    fig = plt.figure(figsize=(6, 6))\n    plt.axis('off')\n    im = plt.imshow(\n        curr_slice, \n        vmin=0, \n        vmax=255\n    )\n    \n    def animate_func(i):\n        curr_slice = combine_image_and_mask_slice(\n            image=image,\n            mask=mask, \n            slice_idx=i,\n            min_value=min_value,\n            max_value=max_value,\n            edge=edge\n        )\n        \n        im.set_array(curr_slice)\n        return [im]\n\n    return animation.FuncAnimation(\n        fig,\n        animate_func,\n        frames = image.shape[2],\n        interval = 1000//24\n    )\n\n\ndef combine_image_and_mask_slice(\n    image: np.ndarray,\n    mask: np.ndarray,\n    slice_idx: int,\n    min_value: int = -200,\n    max_value: int = 200,\n    edge: float = 0.5, \n    image_size: Tuple[int, int] = (512, 512),\n):\n    image_slice = image[:, :, slice_idx]\n    mask_slice = mask[:, :, slice_idx]\n    \n    image_slice = preprocess_image(\n        image=image_slice, \n        min_value=min_value,\n        max_value=max_value,\n        image_size=image_size,\n    )\n    mask_slice = preprocess_mask(\n        mask=mask_slice,\n        edge=edge,\n        image_size=image_size,\n    )\n    \n    result = cv2.addWeighted(image_slice, 0.7, mask_slice, 0.3, 0.0)\n\n    return result\n    \ndef preprocess_image(\n    image: np.ndarray,\n    min_value: int = -200,\n    max_value: int = 200,\n    image_size: Tuple[int, int] = (512, 512),\n) -> np.ndarray:\n    image = np.clip(image, min_value, max_value)\n    image = (image - min_value) / (max_value - min_value)\n    image = cv2.resize(src=image, dsize=image_size)\n    \n    image = np.uint8(image * 255)\n    image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)\n\n    return image\n\ndef preprocess_mask(\n    mask: np.ndarray,\n    edge: float = 0.5,\n    image_size: Tuple[int, int] = (512, 512),\n) -> np.ndarray:\n    mask = np.uint8(mask > edge)  \n    mask = np.uint8(mask * 255) \n    mask = cv2.resize(src=mask, dsize=image_size)\n    mask = mask[..., np.newaxis]\n\n    mask = np.concatenate(\n        [\n            mask,\n            np.zeros_like(mask),\n            np.zeros_like(mask),\n        ],\n        axis=-1,\n    )\n\n    return mask\n    ","metadata":{"execution":{"iopub.status.busy":"2021-09-21T15:26:34.551063Z","iopub.execute_input":"2021-09-21T15:26:34.551388Z","iopub.status.idle":"2021-09-21T15:26:34.581329Z","shell.execute_reply.started":"2021-09-21T15:26:34.551352Z","shell.execute_reply":"2021-09-21T15:26:34.580612Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def _resize_ct(\n    ct: np.ndarray, ct_size: Tuple[int, int, int] = (512, 512, 256)\n) -> np.ndarray:\n    zoom_factor = [\n        first_value / second_value\n        for first_value, second_value in zip(ct_size, ct.shape)\n    ]\n    res_ct = scipy.ndimage.zoom(ct, zoom=zoom_factor)\n\n    assert res_ct.shape == ct_size, f'Bad result size: {res_ct.shape}!={ct_size}'\n\n    return res_ct","metadata":{"execution":{"iopub.status.busy":"2021-09-21T15:26:34.803524Z","iopub.execute_input":"2021-09-21T15:26:34.803785Z","iopub.status.idle":"2021-09-21T15:26:34.809798Z","shell.execute_reply.started":"2021-09-21T15:26:34.803754Z","shell.execute_reply":"2021-09-21T15:26:34.808770Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"image_to_draw = _resize_ct(\n    ct=image,\n    ct_size=(184, 184, 128)\n)\n\ncreate_animation(\n    image=image_to_draw, \n    mask=mask[1],\n    min_value=image.min(),\n    max_value=image.max()\n)","metadata":{"execution":{"iopub.status.busy":"2021-09-21T15:26:35.060866Z","iopub.execute_input":"2021-09-21T15:26:35.061321Z","iopub.status.idle":"2021-09-21T15:26:46.018392Z","shell.execute_reply.started":"2021-09-21T15:26:35.061290Z","shell.execute_reply":"2021-09-21T15:26:46.017723Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Conclusion\n\nYou can use resulted masks as additional channel for your classification model. Or in some other way. Have fun!","metadata":{}}]}