test_example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "induced-lounge",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tempfile\n",
    "import warnings\n",
    "import glob\n",
    "import torch\n",
    "import utils\n",
    "import SimpleITK as sitk\n",
    "import numpy as np\n",
    "import pystrum.pynd.ndutils as nd\n",
    "import matplotlib.pyplot as plt\n",
    "import torch.nn as nn\n",
    "from model.teacher_student_model import teacher_student_model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "banned-insert",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "burning-honor",
   "metadata": {},
   "outputs": [],
   "source": [
    "class AverageMeter(object):\n",
    "    \"\"\"Computes and stores the average and current value\"\"\"\n",
    "    def __init__(self):\n",
    "        self.reset()\n",
    "\n",
    "    def reset(self):\n",
    "        self.val = 0\n",
    "        self.avg = 0\n",
    "        self.sum = 0\n",
    "        self.count = 0\n",
    "        self.vals = []\n",
    "        self.std = 0\n",
    "\n",
    "    def update(self, val, n=1):\n",
    "        self.val = val\n",
    "        self.sum += val * n\n",
    "        self.count += n\n",
    "        self.avg = self.sum / self.count\n",
    "        self.vals.append(val)\n",
    "        self.std = np.std(self.vals)\n",
    "\n",
    "def multi_class_dice_coefficient(label1, label2,num_classes, include_background=True):\n",
    "    dice_coefficients = []\n",
    "\n",
    "    for class_idx in range(0 if include_background else 1, num_classes):\n",
    "        \n",
    "        label1_binary = (label1 == class_idx).astype(int)\n",
    "        label2_binary = (label2 == class_idx).astype(int)\n",
    "        intersection = np.sum(label1_binary * label2_binary)\n",
    "        union = np.sum(label1_binary) + np.sum(label2_binary)\n",
    "\n",
    "        dice_coefficient = (2.0 * intersection) / (union + 1e-8) \n",
    "        dice_coefficients.append(dice_coefficient)\n",
    "    \n",

    "    filtered_dice_coefficients = [x for x in dice_coefficients if x != 0]\n",

    "    avg_dice = round(np.mean(filtered_dice_coefficients), 3)   \n",
    "    #avg_dice = round(np.mean(dice_coefficients), 3)           \n",
    "    return dice_coefficients,avg_dice"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "906a4d2e-ac1a-41ea-bc08-0e7e28804273",
   "metadata": {},
   "source": [
    "# Param"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3494080f-20fc-419e-a44b-2a8592cb7517",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 1, 224, 192, 160])\n"
     ]
    }
   ],
   "source": [
    "dir_save = \"your model path\"\n",
    "batch_size = 1 \n",
    "num_classes = 36 #36 for OASIS,6 for CT, 31 for IXI\n",
    "folder_names = os.path.join(dir_save,\"test_output\")\n",
    "load_pretrained_model_weights = True\n",
    "pth_filename =  \"OS-MedSeg_pre-trained_model_on_OASIS*\" # pre-trained model filename\n",
    "\n",
    "source_image_path = \"D:/datasets/OASIS/OASIS_OAS1_0001_MR1/aligned_norm.nii.gz\" # your image path\n",
    "source_label_path = \"D:/datasets/OASIS/OASIS_OAS1_0001_MR1/aligned_seg35.nii.gz\"\n",
    "source_image_sitk = sitk.ReadImage(source_image_path)\n",
    "source_label_sitk = sitk.ReadImage(source_label_path)\n",
    "\n",
    "source_image_np = sitk.GetArrayFromImage(source_image_sitk) #Converting sitk_metadata to image Array\n",
    "source_label_np = sitk.GetArrayFromImage(source_label_sitk)\n",
    "\n",
    "source_image = torch.Tensor(source_image_np).unsqueeze(dim = 0).unsqueeze(dim = 0)\n",
    "source_label = torch.Tensor(source_label_np).unsqueeze(dim = 0).unsqueeze(dim = 0)\n",
    "print(source_image.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "floppy-history",
   "metadata": {},
   "source": [
    "# Visualize registration performance of trained network\n",
    "## Load pretrained model and perform forward pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "swiss-nickel",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Best model: .\\experiments\\OASIS\\syn_distill_model_hint_1_seg_1_recon_1_lr=1e-3\\distill_model_kpt_best_dice_69_0.8459.pth\n",
      "Dice | 0.851\n",
      "Dice Coefficient for class 1: 0.939\n",
      "Dice Coefficient for class 2: 0.896\n",
      "Dice Coefficient for class 3: 0.944\n",
      "Dice Coefficient for class 4: 0.536\n",
      "Dice Coefficient for class 5: 0.924\n",
      "Dice Coefficient for class 6: 0.938\n",
      "Dice Coefficient for class 7: 0.917\n",
      "Dice Coefficient for class 8: 0.893\n",
      "Dice Coefficient for class 9: 0.926\n",
      "Dice Coefficient for class 10: 0.911\n",
      "Dice Coefficient for class 11: 0.895\n",
      "Dice Coefficient for class 12: 0.900\n",
      "Dice Coefficient for class 13: 0.954\n",
      "Dice Coefficient for class 14: 0.863\n",
      "Dice Coefficient for class 15: 0.904\n",
      "Dice Coefficient for class 16: 0.749\n",
      "Dice Coefficient for class 17: 0.878\n",
      "Dice Coefficient for class 18: 0.826\n",
      "Dice Coefficient for class 19: 0.336\n",
      "Dice Coefficient for class 20: 0.951\n",
      "Dice Coefficient for class 21: 0.905\n",
      "Dice Coefficient for class 22: 0.950\n",
      "Dice Coefficient for class 23: 0.755\n",
      "Dice Coefficient for class 24: 0.923\n",
      "Dice Coefficient for class 25: 0.948\n",
      "Dice Coefficient for class 26: 0.922\n",
      "Dice Coefficient for class 27: 0.916\n",
      "Dice Coefficient for class 28: 0.904\n",
      "Dice Coefficient for class 29: 0.872\n",
      "Dice Coefficient for class 30: 0.902\n",
      "Dice Coefficient for class 31: 0.903\n",
      "Dice Coefficient for class 32: 0.679\n",
      "Dice Coefficient for class 33: 0.890\n",
      "Dice Coefficient for class 34: 0.749\n",
      "Dice Coefficient for class 35: 0.492\n"
     ]
    }
   ],
   "source": [
    "# ==============================================\n",
    "# Test\n",
    "# ==============================================\n",
    "\n",
    "if load_pretrained_model_weights:\n",
    "    dir_load = dir_save  # folder where network weights are stored\n",
    "    # instantiate model\n",
    "    model = teacher_student_model(\n",
    "    in_channel=1,\n",
    "    num_class=36,\n",
    "    channel_list=(16, 32, 64, 128, 256),\n",
    "    residual=True,\n",
    "    vae=False,\n",
    "    device='cuda'\n",
    "    )\n",
    "    \n",
    "    # load model weights\n",
    "    #filename_best_model = glob.glob(os.path.join(dir_load, \"distill_model_kpt_best_dice*\"))[0]\n",
    "    filename_best_model = glob.glob(os.path.join(dir_load, pth_filename))[0] # your pth filename\n",
    "    model.load_state_dict(torch.load(filename_best_model))\n",
    "    print('Best model: {}'.format(filename_best_model))\n",
    "    #print(model)\n",
    "    \n",
    "model.eval()\n",
    "eval_dsc = AverageMeter()\n",
    "with torch.no_grad():\n",
    "\n",
    "    x = source_image.cuda()\n",
    "    x_seg = source_label.cuda()\n",
    "    \n",
    "    pred_label,_ = model(x,x,is_Training=False)\n",
    "    pred_label,student_features = model(x,x,is_Training=False)\n",
    "    pred_label = torch.argmax(pred_label, dim=1)\n",
    "    #print(\"x_seg shape: \",x_seg.shape)\n",
    "    #print(\"pred label shape: \",pred_label.squeeze().cpu().numpy().shape)\n",
    "\n",
    "    dice_coefficients,avg_dice = multi_class_dice_coefficient(pred_label.squeeze().cpu().numpy(), \\\n",
    "                                                              x_seg.squeeze().cpu().numpy(),num_classes, include_background=False)#num_classes=36 for OASIS,27 for CT, 31 for IXI \n",
    "    eval_dsc.update(avg_dice)\n",
    "\n",
    "    # Save Images\n",
    "    if not os.path.exists(str('./{}'.format(folder_names))):\n",
    "        os.makedirs(str('./{}'.format(folder_names)))\n",
    "\n",
    "    fixed_image_sitk = source_image_sitk\n",
    "\n",
    "    pred_label = sitk.GetImageFromArray(pred_label.data.cpu().numpy().squeeze().astype(np.int16))\n",
    "    pred_label.SetSpacing(fixed_image_sitk.GetSpacing())\n",
    "    pred_label.SetDirection(fixed_image_sitk.GetDirection())\n",
    "    pred_label.SetOrigin(fixed_image_sitk.GetOrigin())       \n",
    "\n",
    "    sitk.WriteImage(pred_label, str('./{}/pred_label.nii.gz'.format(folder_names)))\n",
    "\n",
    "# metrics\n",
    "\n",
    "print(f'Dice | {eval_dsc.avg:.3f}')\n",
    "id_list = [i for i in range(1,num_classes)]\n",
    "\n",
    "for idx, avg_dice in enumerate(dice_coefficients, 1):\n",
    "    print(f\"Dice Coefficient for class {id_list[idx-1]}: {avg_dice:.3f}\")\n",
    "\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c2cbfc7-cc6e-419f-960f-107933e253b7",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}