interpretation/ENDOs_T2_128.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a44204f5",
   "metadata": {},
   "source": [
    "Sample code to generate 128 dimensional ENDOs from T2-FLAIR. Similar code can be used to generate ENDOs from T1 (change checkpoint and use T1 MRI)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0735d31",
   "metadata": {},
   "outputs": [],
   "source": [
    "#imports \n",
    "\n",
    "#PyTorch\n",
    "import torch\n",
    "from torch.nn import functional as F\n",
    "from torch import nn\n",
    "\n",
    "#general\n",
    "from glob import glob\n",
    "import pandas as pd\n",
    "import nibabel as nib\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "\n",
    "#imaging\n",
    "from monai import transforms\n",
    "import nibabel\n",
    "\n",
    "#custom\n",
    "from model128 import *\n",
    "from dataset_visualize import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "676408de",
   "metadata": {},
   "outputs": [],
   "source": [
    "transforms_monai = transforms.Compose(\n",
    "    [transforms.AddChannel(), transforms.ToTensor(),]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2f2131b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "class aedataset_T2(torch.utils.data.Dataset):\n",
    "    def __init__(self, datafile, transforms):\n",
    "        \"\"\"\n",
    "        Provide list of files for unbiased brain\n",
    "        \"\"\"\n",
    "        self.unbiased_brain = datafile\n",
    "        self.transforms = transforms\n",
    "\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.unbiased_brain)\n",
    "\n",
    "    def __getitem__(self, idxx=int):\n",
    "        img = nib.load(self.unbiased_brain[idxx])\n",
    "        img = img.get_fdata()\n",
    "        img = (img - img[img != 0].mean()) / img[img != 0].std()\n",
    "        img = self.transforms(img)\n",
    "\n",
    "        img = img.type(torch.float)\n",
    "        MRI_name = self.unbiased_brain[idxx]\n",
    "\n",
    "        return img, MRI_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "52d70066",
   "metadata": {},
   "outputs": [],
   "source": [
    "## T2-FLAIR MRI (linearly registered to MNI152, brain extracted)\n",
    "T2 = glob(\"/T2/*\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c1fd2d76",
   "metadata": {},
   "outputs": [],
   "source": [
    "#T2 ckpt\n",
    "T2_ckpt = \"T2.ckpt\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8f67fc4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:6\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "4206dcd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "T2_ds = aedataset_T2(\n",
    "    datafile=T2,\n",
    "    transforms=transforms_monai,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0bf8a3eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "T2_dataloader = torch.utils.data.DataLoader(\n",
    "    T2_ds, batch_size=16, pin_memory=True, num_workers=16, shuffle=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "b5592054",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = torch.device(\"cuda:1\")\n",
    "model_T2 = engine_AE(0.001)\n",
    "\n",
    "model_T2 = model_T2.to(device)\n",
    "\n",
    "checkpoint = torch.load(T2_ckpt, map_location=device)\n",
    "model_T2.load_state_dict(checkpoint[\"state_dict\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "31c2ad20",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2703/2703 [1:03:11<00:00,  1.40s/it]\n"
     ]
    }
   ],
   "source": [
    "bottle_neck_T2 = []\n",
    "img_names_T2 = []\n",
    "model_T2 = model_T2.eval()\n",
    "with torch.no_grad():\n",
    "    for data in tqdm(T2_dataloader , total=len(T2_dataloader)):\n",
    "        img, name = data\n",
    "        img = img.to(device)\n",
    "        recon, lin1 = model_T2(img)\n",
    "        bottle_neck_T2.extend(lin1.cpu().numpy())\n",
    "        img_names_T2.extend(name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "a8b182a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "b8f4d1fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('T2_128.pkl', 'wb') as f:\n",
    "    pickle.dump(bottle_neck_T2, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "ec36ef21",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"T2_img_names\", \"wb\") as f:\n",
    "    pickle.dump(img_names_T2, f)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}