predict128.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9d9ed1e2",
   "metadata": {},
   "source": [
    "# Prediction for 3D U-Net model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f64ba82b",
   "metadata": {},
   "source": [
    "### Installation of patchify library"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3c439f19",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: patchify in c:\\users\\ajkuna seipi\\anaconda3\\lib\\site-packages (0.2.3)\n",
      "Requirement already satisfied: numpy<2,>=1 in c:\\users\\ajkuna seipi\\anaconda3\\lib\\site-packages (from patchify) (1.21.5)\n"
     ]
    }
   ],
   "source": [
    "#Use patchify to break large volumes into smaller for training \n",
    "#and also to put patches back together after prediction.\n",
    "!pip install patchify"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f37e0a16",
   "metadata": {},
   "source": [
    "### GPU availability "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c9de5a9d",
   "metadata": {},
   "outputs": [
    {
     "ename": "SystemError",
     "evalue": "GPU device not found",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mSystemError\u001b[0m                               Traceback (most recent call last)",
      "\u001b[1;32m~\\AppData\\Local\\Temp\\ipykernel_12616\\1449781351.py\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      3\u001b[0m \u001b[0mdevice_name\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mtf\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mtest\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mgpu_device_name\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      4\u001b[0m \u001b[1;32mif\u001b[0m \u001b[0mdevice_name\u001b[0m \u001b[1;33m!=\u001b[0m \u001b[1;34m'/device:GPU:0'\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 5\u001b[1;33m     \u001b[1;32mraise\u001b[0m \u001b[0mSystemError\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;34m'GPU device not found'\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      6\u001b[0m \u001b[0mprint\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;34m'Found GPU at: {}'\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mformat\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mdevice_name\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mSystemError\u001b[0m: GPU device not found"
     ]
    }
   ],
   "source": [
    "#Make sure the GPU is available. \n",
    "import tensorflow as tf\n",
    "device_name = tf.test.gpu_device_name()\n",
    "if device_name != '/device:GPU:0':\n",
    "    raise SystemError('GPU device not found')\n",
    "print('Found GPU at: {}'.format(device_name))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27ac8e59",
   "metadata": {},
   "source": [
    "### Loading the four pairs of images and masks tif files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6509fb1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from skimage import io\n",
    "from patchify import patchify, unpatchify\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from keras import backend as K\n",
    "from tensorflow.keras.utils import to_categorical\n",
    "from sklearn.model_selection import train_test_split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "eaaf8b2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Load input images and masks. \n",
    "#Here we load 180x1024x1024 pixel volume. We will break it into patches of 16x128x128 for training, with a separation step of 5x128x128.\n",
    "#This means that there is an overlap between the first patch and the next one of 5 in the depth dimension. \n",
    "\n",
    "#IMPORTANT: you need to change the file paths in case you do not upload the tif files in the same folder.. \n",
    "image = []\n",
    "\n",
    "image.append(io.imread('wt_pom1D_01_07_R3D_REF_image.tif'))\n",
    "image.append(io.imread('wt_pom1D_01_15_R3D_REF_image.tif'))\n",
    "image.append(io.imread('wt_pom1D_01_20_R3D_REF_image.tif'))\n",
    "image.append(io.imread('wt_pom1D_01_30_R3D_REF_image.tif'))\n",
    "\n",
    "img_patches = []\n",
    "img_patches.append(patchify(image[0], (16, 128, 128), step=(5, 128, 128)))\n",
    "img_patches.append(patchify(image[1], (16, 128, 128), step=(5, 128, 128))) \n",
    "img_patches.append(patchify(image[2], (16, 128, 128), step=(5, 128, 128)))  \n",
    "img_patches.append(patchify(image[3], (16, 128, 128), step=(5, 128, 128)))  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "096603f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "mask = []\n",
    "\n",
    "mask.append(io.imread('wt_pom1D_01_07_R3D_REF_mask.tif'))\n",
    "mask.append(io.imread('wt_pom1D_01_15_R3D_REF_mask.tif'))\n",
    "mask.append(io.imread('wt_pom1D_01_20_R3D_REF_mask.tif'))\n",
    "mask.append(io.imread('wt_pom1D_01_30_R3D_REF_mask.tif'))\n",
    "\n",
    "mask_patches = []\n",
    "mask_patches.append(patchify(mask[0], (16, 128, 128), step=(5, 128, 128)))\n",
    "mask_patches.append(patchify(mask[1], (16, 128, 128), step=(5, 64, 64)))  \n",
    "mask_patches.append(patchify(mask[2], (16, 128, 128), step=(5, 128, 128)))  \n",
    "mask_patches.append(patchify(mask[3], (16, 128, 128), step=(5, 128, 128)))  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "733c18bc",
   "metadata": {},
   "source": [
    "## Data pre processing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4516fa29",
   "metadata": {},
   "source": [
    "### Reshape the inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4720be8",
   "metadata": {},
   "outputs": [],
   "source": [
    "#We reshape the patches to get an input image and input mask of shape (N, D, H, W), where N is the total number of patches, \n",
    "# D is the depth size of the patches, H is the height size of the patches, and W is the width size of the patches. \n",
    "input_img = np.reshape(img_patches[0], (-1, img_patches[0].shape[3], img_patches[0].shape[4], img_patches[0].shape[5]))\n",
    "input_mask = np.reshape(mask_patches[0], (-1, mask_patches[0].shape[3], mask_patches[0].shape[4], mask_patches[0].shape[5]))\n",
    "\n",
    "for i in range(1, 4):\n",
    "    input_img += np.reshape(img_patches[i], (-1, img_patches[i].shape[3], img_patches[i].shape[4], img_patches[i].shape[5]))\n",
    "    input_mask += np.reshape(mask_patches[i], (-1, mask_patches[i].shape[3], mask_patches[i].shape[4], mask_patches[i].shape[5]))\n",
    "\n",
    "input_img = np.array(input_img)\n",
    "input_mask = np.array(input_mask)\n",
    "\n",
    "print(input_img.shape)\n",
    "print(input_mask.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d1328c5",
   "metadata": {},
   "source": [
    "### Removing empy patches "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5159abac",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Keep in a variable all the indices where the patches are empty in the whole sequences \n",
    "idx_img = np.where(input_mask.mean(axis=(1,2,3)) != 0)[0]\n",
    "\n",
    "input_img = input_img[idx_img]\n",
    "input_mask = input_mask[idx_img]\n",
    "\n",
    "#Print the number of patches we have \n",
    "print(input_img.shape[0])\n",
    "print(input_mask.shape[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "fc98620c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Standardize the input array of pixels by the maximum value the pixels\n",
    "train_img = input_img / input_img.max() \n",
    "\n",
    "#Expand the dimension of the training sets by 1 to match with the input of the model (i.e. chanel number). \n",
    "train_img = np.expand_dims(train_img, axis=4)\n",
    "train_mask = np.expand_dims(input_mask, axis=4)\n",
    "\n",
    "#Since we are performing a binary segmentation, we need to binarize the masks. \n",
    "train_mask[train_mask>1] = 1\n",
    "\n",
    "#Finally, we perform one hot encoding with the function to_categorical with a chosen number of classes of 2. \n",
    "n_classes=2\n",
    "train_mask_cat = to_categorical(train_mask, num_classes=n_classes)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00d0e018",
   "metadata": {},
   "source": [
    "### Split randomly into training and validation set "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9da08034",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'train_test_split' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[1;32m~\\AppData\\Local\\Temp\\ipykernel_1000\\3353077275.py\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[1;32m----> 1\u001b[1;33m \u001b[0mX_train\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mX_test\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0my_train\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0my_test\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mtrain_test_split\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mtrain_img\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtrain_mask_cat\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtest_size\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;36m0.20\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mrandom_state\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;36m0\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m: name 'train_test_split' is not defined"
     ]
    }
   ],
   "source": [
    "X_train, X_test, y_train, y_test = train_test_split(train_img, train_mask_cat, test_size = 0.20, random_state = 0)\n",
    "print(X_train.shape)\n",
    "print(X_test.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f010d148",
   "metadata": {},
   "source": [
    "### Prediction with the model we trained"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df33e8a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from keras.models import load_model\n",
    "\n",
    "# Load the pre-trained model and predict from it \n",
    "# You need to specify another path file if the pre-trained model is not in the one one we specified, or simply add it to the /saved_models folder\n",
    "model = load_model('/saved_models/3dunetmodel_leaky_bs4_16x128x128_epoch120.h5', compile=False)\n",
    "\n",
    "# Predict with the model trained\n",
    "y_pred=model.predict(X_test)\n",
    "\n",
    "#Predict on the test data\n",
    "y_pred_argmax=np.argmax(y_pred, axis=4)\n",
    "y_test_argmax = np.argmax(y_test, axis=4)\n",
    "\n",
    "print(y_pred_argmax.shape)\n",
    "print(y_test_argmax.shape)\n",
    "print(np.unique(y_pred_argmax))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3adcb28d",
   "metadata": {},
   "source": [
    "### Mean IoU "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ec10065",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Using built in keras function for IoU\n",
    "#Only works on TF > 2.0\n",
    "from keras.metrics import MeanIoU\n",
    "n_classes = 2\n",
    "IOU_keras = MeanIoU(num_classes=n_classes)  \n",
    "IOU_keras.update_state(y_test_argmax, y_pred_argmax)\n",
    "print(\"Mean IoU =\", IOU_keras.result().numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a84e8728",
   "metadata": {},
   "source": [
    "### Testing random images "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6372634c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Test some random images\n",
    "import random\n",
    "\n",
    "test_img_number = random.randint(0, len(X_test)-1)\n",
    "test_img = X_test[test_img_number]\n",
    "ground_truth=y_test[test_img_number]\n",
    "\n",
    "test_img_input=np.expand_dims(test_img, 0)\n",
    "\n",
    "\n",
    "test_pred = model.predict(test_img_input)\n",
    "test_prediction = np.argmax(test_pred, axis=4)[0,:,:,:]\n",
    "\n",
    "ground_truth_argmax = np.argmax(ground_truth, axis=3)\n",
    "print(ground_truth_argmax.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74178bcb",
   "metadata": {},
   "source": [
    "#### Plotting the testing image, ground truth mask and the prediction "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0ca3013",
   "metadata": {},
   "outputs": [],
   "source": [
    "slice = random.randint(0, ground_truth_argmax.shape[0]-1)\n",
    "plt.figure(figsize=(12, 8))\n",
    "plt.subplot(231)\n",
    "plt.title('Testing Image')\n",
    "plt.imshow(test_img[slice,:,:,0], cmap='gray')\n",
    "plt.subplot(232)\n",
    "plt.title('Testing Label')\n",
    "plt.imshow(ground_truth_argmax[slice,:,:])\n",
    "plt.subplot(233)\n",
    "plt.title('Prediction on test image')\n",
    "plt.imshow(test_prediction[slice,:,:])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "598d7c69",
   "metadata": {},
   "source": [
    "### Now segment the full image using our trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb699aaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from skimage import io\n",
    "from patchify import patchify, unpatchify\n",
    "\n",
    "#Break the large image (volume) into patches of same size as the training images (patches)\n",
    "path_image = 'wt_pom1D_01_07_R3D_REF_image.tif'\n",
    "large_image = io.imread(path_image)\n",
    "patches = patchify(large_image, (16, 128, 128), step=128)  \n",
    "print(large_image.shape)\n",
    "print(patches.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2d4cb0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Predict each 3D patch   \n",
    "predicted_patches = []\n",
    "for i in range(patches.shape[0]):\n",
    "    for j in range(patches.shape[1]):\n",
    "        for k in range(patches.shape[2]):\n",
    "            print(i,j,k)\n",
    "            single_patch = patches[i,j,k, :,:,:]\n",
    "            single_patch = single_patch/patches.max()\n",
    "            single_patch_3ch_input = np.expand_dims(single_patch, axis=0)\n",
    "            single_patch_prediction = model.predict(single_patch_3ch_input)\n",
    "            single_patch_prediction_argmax = np.argmax(single_patch_prediction, axis=4)[0,:,:,:]\n",
    "            predicted_patches.append(single_patch_prediction_argmax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3970da7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Convert list to numpy array\n",
    "predicted_patches = np.array(predicted_patches)\n",
    "print(predicted_patches.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10f9cf52",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Reshape to the shape we had after patchifying\n",
    "predicted_patches_reshaped = np.reshape(predicted_patches, \n",
    "                                        (patches.shape[0], patches.shape[1], patches.shape[2],\n",
    "                                         patches.shape[3], patches.shape[4], patches.shape[5]) )\n",
    "print(predicted_patches_reshaped.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67eb20b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Repach individual patches into the orginal volume shape\n",
    "reconstructed_image = unpatchify(predicted_patches_reshaped, large_image.shape)\n",
    "print(reconstructed_image.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4ed3d69",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(reconstructed_image.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0014e31",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Convert to uint8 so we can open image in most image viewing software packages\n",
    "reconstructed_image=reconstructed_image.astype(np.uint8)\n",
    "print(reconstructed_image.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "216173f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Now save it as segmented volume.\n",
    "from tifffile import imsave\n",
    "imsave('images_segmented.tif', reconstructed_image)"
   ]
  }
 ],
 "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}