fmri_vnet_interface.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyMyGM4XmJq2jKz/brFZ08Pi",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "gpuClass": "standard",
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Achillesy/Fetal_Functional_MRI_Segmentation/blob/master/fmri_vnet_interface.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Automated Brain Masking of Fetal Functional MRI with Open Data"
      ],
      "metadata": {
        "id": "j09lPZQqdzIL"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Preprocessing, downloading resources from my github repository\n",
        "\n",
        "https://github.com/Achillesy/Fetal_Functional_MRI_Segmentation\n",
        "![xuchu](https://avatars.githubusercontent.com/u/5572441?s=16) "
      ],
      "metadata": {
        "id": "fFGYAoRIo5eg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!wget https://github.com/Achillesy/Fetal_Functional_MRI_Segmentation/releases/download/v1.0.0/fold4_train_metric_vnet.pth"
      ],
      "metadata": {
        "id": "QajrY-mkfbx2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Install ![monai](https://monai.io/assets/img/MONAI-logo_color.png) and check environment variables"
      ],
      "metadata": {
        "id": "vdwr3buiz3tz"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install monai\n",
        "\n",
        "from monai.config import print_config\n",
        "print_config()\n"
      ],
      "metadata": {
        "id": "DD8b1scpLgIA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import numpy as np\n",
        "import nibabel as nib\n",
        "from glob import glob\n",
        "\n",
        "import torch\n",
        "from types import SimpleNamespace\n",
        "from google.colab import files\n",
        "\n",
        "cfg = SimpleNamespace(**{})\n",
        "cfg.pixdim = (3.5, 3.5, 3.5)\n",
        "cfg.roi_size = [64, 64, 64]\n",
        "cfg.sw_batch_size = 4\n",
        "cfg.file_pth = \"fold4_train_metric_vnet.pth\"\n",
        "\n",
        "cfg.mri_dir = \"mri\"\n",
        "cfg.mask_dir = \"mask\"\n",
        "os.makedirs(cfg.mri_dir, exist_ok=True)\n",
        "os.makedirs(cfg.mask_dir, exist_ok=True)\n",
        "\n",
        "cfg.device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n"
      ],
      "metadata": {
        "id": "hvYoA-BrGwp_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Load VNet model and parameters"
      ],
      "metadata": {
        "id": "EP4q1W8A0aBa"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from monai.networks.nets import VNet\n",
        "\n",
        "model = VNet(\n",
        "  spatial_dims=3,\n",
        "  in_channels=1,\n",
        "  out_channels=2,\n",
        "  act=(\"elu\", {\"inplace\": True}),\n",
        "  dropout_dim=3,\n",
        "  bias=False,\n",
        ").to(cfg.device)\n",
        "model.load_state_dict(torch.load(cfg.file_pth, cfg.device))\n"
      ],
      "metadata": {
        "id": "BTEnbatnLSXr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from monai.transforms import (\n",
        "  AsDiscreted,\n",
        "  Compose,\n",
        "  EnsureChannelFirstd,\n",
        "  Invertd,\n",
        "  LoadImaged,\n",
        "  NormalizeIntensityd,\n",
        "  Orientationd,\n",
        "  SaveImaged,\n",
        "  Spacingd,\n",
        ")\n",
        "\n",
        "test_transforms = Compose(\n",
        "  [\n",
        "    LoadImaged(keys=[\"image\"]),\n",
        "    EnsureChannelFirstd(keys=[\"image\"]),\n",
        "    Orientationd(keys=[\"image\"], axcodes=\"RAS\"),\n",
        "    Spacingd(keys=[\"image\"], pixdim=cfg.pixdim, mode=\"bilinear\"),\n",
        "    NormalizeIntensityd(keys=\"image\", nonzero=True),\n",
        "  ]\n",
        ")\n",
        "\n",
        "post_transforms = Compose(\n",
        "  [\n",
        "    Invertd(\n",
        "      keys=\"pred\",\n",
        "      transform=test_transforms,\n",
        "      orig_keys=\"image\",\n",
        "      meta_keys=\"pred_meta_dict\",\n",
        "      orig_meta_keys=\"image_meta_dict\",\n",
        "      meta_key_postfix=\"meta_dict\",\n",
        "      nearest_interp=False,\n",
        "      to_tensor=True,\n",
        "    ),\n",
        "    AsDiscreted(keys=\"pred\", argmax=True),\n",
        "    SaveImaged(\n",
        "      keys=\"pred\",\n",
        "      meta_keys=\"pred_meta_dict\",\n",
        "      output_dir=cfg.mask_dir,\n",
        "      output_postfix=\"vnet\",\n",
        "      resample=False,\n",
        "    ),\n",
        "  ]\n",
        ")"
      ],
      "metadata": {
        "id": "I_pwLzSNLOQ2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Please upload your Fetal Functional MRI files"
      ],
      "metadata": {
        "id": "XKlCTgZQgaRV"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "![upload_guide](https://github.com/Achillesy/Fetal_Functional_MRI_Segmentation/blob/master/figures/upload_guide.png?raw=1)\n",
        "1. Click the **Files** icon on the left\n",
        "2. Click the **Upload to session storage** icon above\n",
        "3. Your uploaded files will be displayed here\n",
        "----\n",
        "After double-checking your uploaded files, by <font color=\"green\">pressing the **Enter** key in the input box below</font>, the fMRI mask will be automatically generated in a short time."
      ],
      "metadata": {
        "id": "gfVakLRprliM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "input()\n"
      ],
      "metadata": {
        "id": "j7YfZn2DdMhl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Split into multiple 3D MRI files"
      ],
      "metadata": {
        "id": "txPF0mKF1aZi"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "frmi_files = glob(\"*.nii.gz\")\n",
        "for fmri_data in frmi_files:\n",
        "  fmri_data_name = os.path.basename(fmri_data).replace(\".nii.gz\", \"\")\n",
        "  image = nib.load(fmri_data)\n",
        "  data = image.get_fdata()\n",
        "  if len(data.shape) != 4:\n",
        "    raise ValueError(\"Invalid shape of fMRI file format. Expected 4D shape: [x, y, z, t]\")\n",
        "  channel_list = np.split(data, data.shape[-1], axis=-1)\n",
        "  for i, channel in enumerate(channel_list):\n",
        "    channel_image = nib.Nifti1Image(channel, image.affine)\n",
        "    channel_file_name = os.path.join(cfg.mri_dir, f\"{fmri_data_name}_{i+1}.nii.gz\")\n",
        "    nib.save(channel_image, channel_file_name)\n"
      ],
      "metadata": {
        "id": "2z_9HOboPQky"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "test_files = []\n",
        "rmi_files = glob(os.path.join(cfg.mri_dir, \"*.nii.gz\"))\n",
        "for f_file in rmi_files:\n",
        "  test_files.append({\"image\": f_file})\n",
        "print(test_files)\n"
      ],
      "metadata": {
        "id": "UwGKtCioPYwx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Generate MRI mask"
      ],
      "metadata": {
        "id": "214XWdsQ1-WD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from monai.inferers import sliding_window_inference\n",
        "from monai.data import DataLoader, Dataset, decollate_batch\n",
        "\n",
        "test_ds = Dataset(data=test_files, transform=test_transforms)\n",
        "test_loader = DataLoader(test_ds, batch_size=1)\n"
      ],
      "metadata": {
        "id": "rxSmofwbPaoi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "with torch.no_grad():\n",
        "  for test_data in test_loader:\n",
        "    test_inputs = test_data[\"image\"].to(cfg.device)\n",
        "    test_data[\"pred\"] = sliding_window_inference(\n",
        "      test_inputs, cfg.roi_size, cfg.sw_batch_size, model\n",
        "    )\n",
        "    test_data = [post_transforms(i) for i in decollate_batch(test_data)]"
      ],
      "metadata": {
        "id": "mMwtJ79tPc8K"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Combined into fMRI mask"
      ],
      "metadata": {
        "id": "3PuORVRi2fK6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "for fmri_data in frmi_files:\n",
        "  image = nib.load(fmri_data)\n",
        "  data = image.get_fdata()\n",
        "\n",
        "  fmri_data_name = os.path.basename(fmri_data).replace(\".nii.gz\", \"\")\n",
        "  mask_data = np.zeros_like(data)\n",
        "  for i in range(data.shape[-1]):\n",
        "    i_mask_file = os.path.join(cfg.mask_dir,  f\"{fmri_data_name}_{i+1}\", f\"{fmri_data_name}_{i+1}_vnet.nii.gz\")\n",
        "    i_mask_data = nib.load(i_mask_file).get_fdata()\n",
        "    mask_data[:,:,:,i] = i_mask_data\n",
        "  fmri_mask = nib.Nifti1Image(mask_data, affine=image.affine, header=image.header)\n",
        "  mask_data_name = f\"{fmri_data_name}_vnet.nii.gz\"\n",
        "  fmri_mask.to_filename(mask_data_name)\n",
        "  # files.download(mask_data_name)"
      ],
      "metadata": {
        "id": "0TCtG03GPhVh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Please <font color=\"red\">right-click</font> the generated mask files (**your_frmi_name_vnet.nii.gz**) to open the download link."
      ],
      "metadata": {
        "id": "nVptp4hP5NDL"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Clean up temporary files"
      ],
      "metadata": {
        "id": "YUewdRt_3ICj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!rm -rf {cfg.mri_dir}\n",
        "!rm -rf {cfg.mask_dir}\n"
      ],
      "metadata": {
        "id": "HfzH1p9gPiZY"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}