rib_fracture_viz.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h1>Rib fracture, RoI and prediction visualisations</h1>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Package imports</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import all the necessary packages\n",
    "import numpy as np\n",
    "import nibabel as nib\n",
    "import itk\n",
    "import itkwidgets\n",
    "from ipywidgets import interact, interactive, IntSlider, ToggleButtons\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "from dataset.fracnet_dataset import FracNetTrainDataset\n",
    "from dataset import transforms as tsfm\n",
    "\n",
    "# Matplotlib setting\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Function definition</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Image loading function\n",
    "def load_img_data(image_path):\n",
    "    image_obj = nib.load(image_path)\n",
    "    return image_obj.get_fdata()\n",
    "\n",
    "# Convert a grayscale image (i, j, k) to a rgb image (i, j, k, c)\n",
    "def convert_gray_to_rgb(gray_img):\n",
    "    stacked_img = np.stack((gray_img,)*3, axis=-1)\n",
    "    stacked_img_sca = (stacked_img - stacked_img.min()) / (stacked_img.max() - stacked_img.min())\n",
    "    return (stacked_img_sca*255).astype(int)\n",
    "\n",
    "# Axial plane view function for single image\n",
    "def explore_axial_single(layer):\n",
    "    plt.figure(figsize=(20, 10))\n",
    "    plt.imshow(gray_img1[:, :, layer], cmap='gray');\n",
    "    plt.title('Explore Layers of Ribs', fontsize=20)\n",
    "    plt.axis('off')\n",
    "    return layer\n",
    "\n",
    "# Saggital plane view function for two images\n",
    "def explore_saggital_plane(img1, img2, title1, title2, fig_size=(18, 18), title_font_size=22):\n",
    "    def explore_plane_func(layer):\n",
    "        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=fig_size)\n",
    "        ax1.imshow(img1[layer, :, :, :])\n",
    "        ax2.imshow(img2[layer, :, :, :])\n",
    "        ax1.axis('off')\n",
    "        ax2.axis('off')\n",
    "        ax1.set_title(title1, fontsize=title_font_size)\n",
    "        ax2.set_title(title2, fontsize=title_font_size)\n",
    "        return layer\n",
    "    return explore_plane_func\n",
    "\n",
    "# Coronal plane view function for two images\n",
    "def explore_coronal_plane(img1, img2, title1, title2, fig_size=(18, 18), title_font_size=22):\n",
    "    def explore_plane_func(layer):\n",
    "        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=fig_size)\n",
    "        ax1.imshow(img1[:, layer, :, :])\n",
    "        ax2.imshow(img2[:, layer, :, :])\n",
    "        ax1.axis('off')\n",
    "        ax2.axis('off')\n",
    "        ax1.set_title(title1, fontsize=title_font_size)\n",
    "        ax2.set_title(title2, fontsize=title_font_size)\n",
    "        return layer    \n",
    "    return explore_plane_func\n",
    "\n",
    "# Axial plane view function for two images\n",
    "def explore_axial_plane(img1, img2, title1, title2, fig_size=(18, 18), title_font_size=22):\n",
    "    def explore_plane_func(layer):\n",
    "        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=fig_size)\n",
    "        ax1.imshow(img1[:, :, layer, :])\n",
    "        ax2.imshow(img2[:, :, layer, :])\n",
    "        ax1.axis('off')\n",
    "        ax2.axis('off')\n",
    "        ax1.set_title(title1, fontsize=title_font_size)\n",
    "        ax2.set_title(title2, fontsize=title_font_size)\n",
    "        return layer\n",
    "    return explore_plane_func\n",
    "\n",
    "# Function to get RoIs from centroid\n",
    "def get_roi_coordinates(img_arr, centroid, crop_size=64):\n",
    "    src_beg = [max(0, centroid[i] - crop_size // 2) for i in range(len(centroid))]\n",
    "    src_end = [min(img_arr.shape[i], centroid[i] + crop_size // 2) for i in range(len(centroid))]\n",
    "    dst_beg = [max(0, crop_size // 2 - centroid[i]) for i in range(len(centroid))]\n",
    "    dst_end = [min(img_arr.shape[i] - (centroid[i] - crop_size // 2), crop_size) for i in range(len(centroid))]\n",
    "    return src_beg, src_end, dst_beg, dst_end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Data Loading</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Load the image data\n",
    "gray_img1 = load_img_data(\"../FracNet/data/train/ribfrac-train-images/RibFrac301-image.nii.gz\")\n",
    "gray_img2 = load_img_data(\"../FracNet/data/train/ribfrac-train-images/RibFrac381-image.nii.gz\")\n",
    "gray_img3 = load_img_data(\"../FracNet/data/train/ribfrac-train-images/RibFrac382-image.nii.gz\")\n",
    "label_img1 = load_img_data(\"../FracNet/data/train/ribfrac-train-labels/RibFrac301-label.nii.gz\")\n",
    "label_img2 = load_img_data(\"../FracNet/data/train/ribfrac-train-labels/RibFrac381-label.nii.gz\")\n",
    "label_img3 = load_img_data(\"../FracNet/data/train/ribfrac-train-labels/RibFrac382-label.nii.gz\")\n",
    "pred_img2 = load_img_data(\"../FracNet/data/train/ribfrac-train-pred_w_const_gauss_noise/RibFrac381-label.nii.gz\")\n",
    "pred_img3 = load_img_data(\"../FracNet/data/train/ribfrac-train-pred_w_const_gauss_noise/RibFrac382-label.nii.gz\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Colours</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the colours to apply\n",
    "red = np.array([255, 0, 0])\n",
    "blue = np.array([0, 0, 140])\n",
    "green = np.array([0, 100, 0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert the gray scale images to rgb\n",
    "rgb_img1 = convert_gray_to_rgb(gray_img1)\n",
    "rgb_img2 = convert_gray_to_rgb(gray_img2)\n",
    "rgb_img3 = convert_gray_to_rgb(gray_img3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "gray_shape: ((512, 512, 407), (512, 512, 349), (512, 512, 313))\n",
      "rgb_shape: ((512, 512, 407, 3), (512, 512, 349, 3), (512, 512, 313, 3))\n"
     ]
    }
   ],
   "source": [
    "# Print the array shapes\n",
    "gray_shape = (gray_img1.shape, gray_img2.shape, gray_img3.shape)\n",
    "rgb_shape = (rgb_img1.shape, rgb_img2.shape, rgb_img3.shape)\n",
    "print(f'gray_shape: {gray_shape}')\n",
    "print(f'rgb_shape: {rgb_shape}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "<h3>Ribs visualisation (Axial plane)</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "81c522c9613c484eb6745dec19aa6313",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=203, description='layer', max=406), Output()), _dom_classes=('widget-int…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.explore_axial_single(layer)>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create a widget to explore the image\n",
    "interact(explore_axial_single, layer=(0, gray_img1.shape[2] - 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Fracture visualisation</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy the rgb image\n",
    "rgb_img1_w_labels = rgb_img1.copy()\n",
    "\n",
    "# Set the fracture red\n",
    "rgb_img1_w_labels[np.where(label_img1 != 0)] = red\n",
    "\n",
    "# Set the title of the images\n",
    "title1 = 'Original image'\n",
    "title2 = 'Image with fracture'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Fracture visualisation (Saggital place)</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9c1cc1aff6124202ad89f570aec544fa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=255, description='layer', max=511), Output()), _dom_classes=('widget-int…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.explore_saggital_plane.<locals>.explore_plane_func(layer)>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create the plane exploration function\n",
    "explore_func = explore_saggital_plane(img1=rgb_img1, img2=rgb_img1_w_labels, title1=title1, title2=title2)\n",
    "\n",
    "# Create an interactive ipywidgets visualisation\n",
    "interact(explore_func, layer=(0, rgb_img1.shape[0] - 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Fracture visualisation (Coronal place)</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1a56744cca4f4720bda580874932578b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=255, description='layer', max=511), Output()), _dom_classes=('widget-int…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.explore_coronal_plane.<locals>.explore_plane_func(layer)>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create the plane exploration function\n",
    "explore_func = explore_coronal_plane(img1=rgb_img1, img2=rgb_img1_w_labels, title1=title1, title2=title2)\n",
    "\n",
    "# Create an interactive ipywidgets visualisation\n",
    "interact(explore_func, layer=(0, rgb_img1.shape[1] - 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Fracture visualisation (Axial place)</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "44c39042b62f4b6eb7d3c8d58fcb1567",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=203, description='layer', max=406), Output()), _dom_classes=('widget-int…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.explore_axial_plane.<locals>.explore_plane_func(layer)>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create the plane exploration function\n",
    "explore_func = explore_axial_plane(img1=rgb_img1, img2=rgb_img1_w_labels, title1=title1, title2=title2)\n",
    "\n",
    "# Create an interactive ipywidgets visualisation\n",
    "interact(explore_func, layer=(0, rgb_img1.shape[2] - 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Prediction visualisation</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy the rgb image\n",
    "rgb_img2_w_preds = rgb_img2.copy()\n",
    "\n",
    "# Set the predictions red\n",
    "rgb_img2_w_preds[np.where(pred_img2 != 0)] = red\n",
    "\n",
    "# Set the title of the images\n",
    "title1 = 'Original image'\n",
    "title2 = 'Image with facture prediction'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Prediction visualisation (Saggital place)</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ecc1a5d49fd746ca9f7988465579e40e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=255, description='layer', max=511), Output()), _dom_classes=('widget-int…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.explore_saggital_plane.<locals>.explore_plane_func(layer)>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create the plane exploration function\n",
    "explore_func = explore_saggital_plane(img1=rgb_img2, img2=rgb_img2_w_preds, title1=title1, title2=title2)\n",
    "\n",
    "# Create an interactive ipywidgets visualisation\n",
    "interact(explore_func, layer=(0, rgb_img2.shape[0] - 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Prediction visualisation (Coronal place)</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0b3ca283a9944dc2b1c0bd1886ed971a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=255, description='layer', max=511), Output()), _dom_classes=('widget-int…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.explore_coronal_plane.<locals>.explore_plane_func(layer)>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create the plane exploration function\n",
    "explore_func = explore_coronal_plane(img1=rgb_img2, img2=rgb_img2_w_preds, title1=title1, title2=title2)\n",
    "\n",
    "# Create an interactive ipywidgets visualisation\n",
    "interact(explore_func, layer=(0, rgb_img2.shape[1] - 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Prediction visualisation (Axial place)</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fe6612bf3dc54e6082d67ce8e980805a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=174, description='layer', max=348), Output()), _dom_classes=('widget-int…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.explore_axial_plane.<locals>.explore_plane_func(layer)>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create the plane exploration function\n",
    "explore_func = explore_axial_plane(img1=rgb_img2, img2=rgb_img2_w_preds, title1=title1, title2=title2)\n",
    "\n",
    "# Create an interactive ipywidgets visualisation\n",
    "interact(explore_func, layer=(0, rgb_img2.shape[2] - 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Prediction vs label visualisation</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy the rgb image\n",
    "rgb_img2_w_preds_n_labels = rgb_img2.copy()\n",
    "\n",
    "# Confusion matrix\n",
    "true_pos_cond = (pred_img2 != 0) & (label_img2 != 0)\n",
    "false_pos_cond = (pred_img2 != 0) & (label_img2 == 0)\n",
    "false_neg_cond = (pred_img2 == 0) & (label_img2 != 0)\n",
    "\n",
    "# Set colours\n",
    "rgb_img2_w_preds_n_labels[np.where(true_pos_cond)] = red\n",
    "rgb_img2_w_preds_n_labels[np.where(false_pos_cond)] = blue\n",
    "rgb_img2_w_preds_n_labels[np.where(false_neg_cond)] = green\n",
    "\n",
    "# Set the title of the images\n",
    "title1 = 'Original image'\n",
    "title2 = 'Image w inference (TP: red, FP: blue, FN: green)'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Prediction vs label visualisation (Saggital place)</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c40997240a2b417086e1a8e12c1ab0e8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=255, description='layer', max=511), Output()), _dom_classes=('widget-int…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.explore_saggital_plane.<locals>.explore_plane_func(layer)>"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create the plane exploration function\n",
    "explore_func = explore_saggital_plane(img1=rgb_img2, img2=rgb_img2_w_preds_n_labels, title1=title1, title2=title2)\n",
    "\n",
    "# Create an interactive ipywidgets visualisation\n",
    "interact(explore_func, layer=(0, rgb_img2.shape[0] - 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Prediction vs label visualisation (Coronal place)</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f994b720ad5140808b39434c0de6cb84",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=255, description='layer', max=511), Output()), _dom_classes=('widget-int…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.explore_coronal_plane.<locals>.explore_plane_func(layer)>"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create the plane exploration function\n",
    "explore_func = explore_coronal_plane(img1=rgb_img2, img2=rgb_img2_w_preds_n_labels, title1=title1, title2=title2)\n",
    "\n",
    "# Create an interactive ipywidgets visualisation\n",
    "interact(explore_func, layer=(0, rgb_img2.shape[1] - 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Prediction vs label visualisation (Axial place)</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7881d0cde06f4df59701422ae0c8589b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=174, description='layer', max=348), Output()), _dom_classes=('widget-int…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.explore_axial_plane.<locals>.explore_plane_func(layer)>"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create the plane exploration function\n",
    "explore_func = explore_axial_plane(img1=rgb_img2, img2=rgb_img2_w_preds_n_labels, title1=title1, title2=title2)\n",
    "\n",
    "# Create an interactive ipywidgets visualisation\n",
    "interact(explore_func, layer=(0, rgb_img2.shape[2] - 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Original vs modified positive sampling visualisation</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy rgb image\n",
    "rgb_img2_w_rois = rgb_img2.copy()\n",
    "\n",
    "# Set transform for FracNetTrainDataset\n",
    "transforms = [\n",
    "    tsfm.Window(-200, 1000),\n",
    "    tsfm.MinMaxNorm(-200, 1000)\n",
    "]\n",
    "\n",
    "# Instantiate FracNetTrainDataset\n",
    "frac_dataset = FracNetTrainDataset(\n",
    "    image_dir='/home/ryan/PycharmProjects/FracNet/data/val/ribfrac-val-images/', \n",
    "    label_dir='/home/ryan/PycharmProjects/FracNet/data/val/ribfrac-val-labels/',\n",
    "    transforms=transforms\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we simulate 10 epochs of positive sampling regions of interest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retrieve ROIs\n",
    "epochs_sim = 10\n",
    "pos_centroids = []\n",
    "pos_centroids_new = []\n",
    "\n",
    "for e in range(epochs_sim):\n",
    "    pos_centroids += frac_dataset._get_pos_centroids(label_img2.astype(int))\n",
    "    pos_centroids_new += frac_dataset._get_pos_centroids_new(label_img2.astype(int))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the positive RoIs\n",
    "pos_roi_coordinates = [get_roi_coordinates(rgb_img2_w_rois, roi) for roi in pos_centroids]\n",
    "pos_roi_coordinates_new = [get_roi_coordinates(rgb_img2_w_rois, roi) for roi in pos_centroids_new]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy the rgb image\n",
    "rgb_img2_w_rois_new = rgb_img2_w_rois.copy()\n",
    "\n",
    "# Colour the RoIs of the original positive sampling\n",
    "for src_beg, src_end, dst_beg, dst_end in pos_roi_coordinates:\n",
    "    rgb_img2_w_rois[\n",
    "        src_beg[0]:src_end[0],\n",
    "        src_beg[1]:src_end[1],\n",
    "        src_beg[2]:src_end[2],\n",
    "        :\n",
    "    ] = blue\n",
    "\n",
    "# Colour the RoIs of the modified positive sampling\n",
    "for src_beg, src_end, dst_beg, dst_end in pos_roi_coordinates_new:\n",
    "    rgb_img2_w_rois_new[\n",
    "        src_beg[0]:src_end[0],\n",
    "        src_beg[1]:src_end[1],\n",
    "        src_beg[2]:src_end[2],\n",
    "        :\n",
    "    ] = blue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Colour the fractures\n",
    "rgb_img2_w_rois[np.where(label_img2 != 0)] = red\n",
    "rgb_img2_w_rois_new[np.where(label_img2 != 0)] = red"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the title of the images\n",
    "title1 = 'Original positive sampling'\n",
    "title2 = 'Modified positive sampling'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Original vs modified positive sampling visualisation (Saggital place)</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c37b530f05bb49469dd64c29bcfb4e8d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=255, description='layer', max=511), Output()), _dom_classes=('widget-int…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.explore_saggital_plane.<locals>.explore_plane_func(layer)>"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create the plane exploration function\n",
    "explore_func = explore_saggital_plane(img1=rgb_img2_w_rois, img2=rgb_img2_w_rois_new, title1=title1, title2=title2)\n",
    "\n",
    "# Create an interactive ipywidgets visualisation\n",
    "interact(explore_func, layer=(0, rgb_img2.shape[0] - 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Original vs modified positive sampling visualisation (Coronal place)</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9e1977e3b5fb4a06b30fddae3eb5d8cd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=255, description='layer', max=511), Output()), _dom_classes=('widget-int…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.explore_coronal_plane.<locals>.explore_plane_func(layer)>"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create the plane exploration function\n",
    "explore_func = explore_coronal_plane(img1=rgb_img2_w_rois, img2=rgb_img2_w_rois_new, title1=title1, title2=title2)\n",
    "\n",
    "# Create an interactive ipywidgets visualisation\n",
    "interact(explore_func, layer=(0, rgb_img2.shape[1] - 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Original vs modified positive sampling visualisation (Axial place)</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "490f9fc8015b4c2fb1ad4472a27ab8d2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=174, description='layer', max=348), Output()), _dom_classes=('widget-int…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.explore_axial_plane.<locals>.explore_plane_func(layer)>"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create the plane exploration function\n",
    "explore_func = explore_axial_plane(img1=rgb_img2_w_rois, img2=rgb_img2_w_rois_new, title1=title1, title2=title2)\n",
    "\n",
    "# Create an interactive ipywidgets visualisation\n",
    "interact(explore_func, layer=(0, rgb_img2.shape[2] - 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Original vs modified negative sampling visualisation</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy rgb image\n",
    "rgb_img2_w_rois = rgb_img2.copy()\n",
    "\n",
    "# Set transform for FracNetTrainDataset\n",
    "transforms = [\n",
    "    tsfm.Window(-200, 1000),\n",
    "    tsfm.MinMaxNorm(-200, 1000)\n",
    "]\n",
    "\n",
    "# Instantiate FracNetTrainDataset\n",
    "frac_dataset = FracNetTrainDataset(\n",
    "    image_dir='/home/ryan/PycharmProjects/FracNet/data/val/ribfrac-val-images/', \n",
    "    label_dir='/home/ryan/PycharmProjects/FracNet/data/val/ribfrac-val-labels/',\n",
    "    transforms=transforms\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we simulate 10 epochs of positive sampling regions of interest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retrieve ROIs\n",
    "epochs_sim = 10\n",
    "neg_centroids = []\n",
    "neg_centroids_new = []\n",
    "\n",
    "for e in range(epochs_sim):\n",
    "    neg_centroids += frac_dataset._get_neg_centroids(pos_centroids, label_img2.shape)\n",
    "    neg_centroids_temp = frac_dataset._get_neg_centroids(pos_centroids, label_img2.shape)\n",
    "    neg_centroids_new += neg_centroids_temp + frac_dataset._new_get_neg_centroids(gray_img2, n_centroids=len(neg_centroids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the positive RoIs\n",
    "neg_roi_coordinates = [get_roi_coordinates(rgb_img2_w_rois, roi) for roi in neg_centroids]\n",
    "neg_roi_coordinates_new = [get_roi_coordinates(rgb_img2_w_rois, roi) for roi in neg_centroids_new]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy the rgb image\n",
    "rgb_img2_w_rois_new = rgb_img2_w_rois.copy()\n",
    "\n",
    "# Colour the RoIs of the original negative sampling\n",
    "for src_beg, src_end, dst_beg, dst_end in neg_roi_coordinates:\n",
    "    rgb_img2_w_rois[\n",
    "        src_beg[0]:src_end[0],\n",
    "        src_beg[1]:src_end[1],\n",
    "        src_beg[2]:src_end[2],\n",
    "        :\n",
    "    ] = blue\n",
    "\n",
    "# Colour the RoIs of the modified negative sampling\n",
    "for src_beg, src_end, dst_beg, dst_end in neg_roi_coordinates_new:\n",
    "    rgb_img2_w_rois_new[\n",
    "        src_beg[0]:src_end[0],\n",
    "        src_beg[1]:src_end[1],\n",
    "        src_beg[2]:src_end[2],\n",
    "        :\n",
    "    ] = blue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Colour the fractures\n",
    "rgb_img2_w_rois[np.where(label_img2 != 0)] = red\n",
    "rgb_img2_w_rois_new[np.where(label_img2 != 0)] = red"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the title of the images\n",
    "title1 = 'Original negative sampling'\n",
    "title2 = 'Modified negative sampling'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Original vs modified negative sampling visualisation (Saggital place)</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d58299aacf0c4f78adabf14619eb6078",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=255, description='layer', max=511), Output()), _dom_classes=('widget-int…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.explore_saggital_plane.<locals>.explore_plane_func(layer)>"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create the plane exploration function\n",
    "explore_func = explore_saggital_plane(img1=rgb_img2_w_rois, img2=rgb_img2_w_rois_new, title1=title1, title2=title2)\n",
    "\n",
    "# Create an interactive ipywidgets visualisation\n",
    "interact(explore_func, layer=(0, rgb_img2.shape[0] - 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Original vs modified negative sampling visualisation (Coronal place)</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "29f34501a3ca48af8899583395aef0e7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=255, description='layer', max=511), Output()), _dom_classes=('widget-int…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.explore_coronal_plane.<locals>.explore_plane_func(layer)>"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create the plane exploration function\n",
    "explore_func = explore_coronal_plane(img1=rgb_img2_w_rois, img2=rgb_img2_w_rois_new, title1=title1, title2=title2)\n",
    "\n",
    "# Create an interactive ipywidgets visualisation\n",
    "interact(explore_func, layer=(0, rgb_img2.shape[1] - 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Original vs modified negative sampling visualisation (Axial place)</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "643dc23e3064453ba86971615521bb0f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=174, description='layer', max=348), Output()), _dom_classes=('widget-int…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.explore_axial_plane.<locals>.explore_plane_func(layer)>"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create the plane exploration function\n",
    "explore_func = explore_axial_plane(img1=rgb_img2_w_rois, img2=rgb_img2_w_rois_new, title1=title1, title2=title2)\n",
    "\n",
    "# Create an interactive ipywidgets visualisation\n",
    "interact(explore_func, layer=(0, rgb_img2.shape[2] - 1))"
   ]
  }
 ],
 "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}